from __future__ import annotations

import functools
import threading
import time
from copy import deepcopy
from dataclasses import asdict
from datetime import datetime
from typing import AbstractSet, Any, Dict, Iterable, List, Optional, Set, Tuple, Type

from dbt import tracking, utils
from dbt.adapters.base import BaseAdapter, BaseRelation
from dbt.adapters.capability import Capability
from dbt.adapters.events.types import FinishedRunningStats
from dbt.adapters.exceptions import MissingMaterializationError
from dbt.artifacts.resources import Hook
from dbt.artifacts.schemas.batch_results import BatchResults, BatchType
from dbt.artifacts.schemas.results import (
    NodeStatus,
    RunningStatus,
    RunStatus,
    TimingInfo,
    collect_timing_info,
)
from dbt.artifacts.schemas.run import RunResult
from dbt.cli.flags import Flags
from dbt.clients.jinja import MacroGenerator
from dbt.config import RuntimeConfig
from dbt.context.providers import generate_runtime_model_context
from dbt.contracts.graph.manifest import Manifest
from dbt.contracts.graph.nodes import BatchContext, HookNode, ModelNode, ResultNode
from dbt.events.types import (
    GenericExceptionOnRun,
    LogBatchResult,
    LogHookEndLine,
    LogHookStartLine,
    LogModelResult,
    LogStartBatch,
    LogStartLine,
    MicrobatchExecutionDebug,
)
from dbt.exceptions import CompilationError, DbtInternalError, DbtRuntimeError
from dbt.graph import ResourceTypeSelector
from dbt.graph.thread_pool import DbtThreadPool
from dbt.hooks import get_hook_dict
from dbt.materializations.incremental.microbatch import MicrobatchBuilder
from dbt.node_types import NodeType, RunHookType
from dbt.task import group_lookup
from dbt.task.base import BaseRunner
from dbt.task.compile import CompileRunner, CompileTask
from dbt.task.printer import get_counts, print_run_end_messages
from dbt_common.clients.jinja import MacroProtocol
from dbt_common.dataclass_schema import dbtClassMixin
from dbt_common.events.base_types import EventLevel
from dbt_common.events.contextvars import log_contextvars
from dbt_common.events.functions import fire_event, get_invocation_id
from dbt_common.events.types import Formatting
from dbt_common.exceptions import DbtValidationError


@functools.total_ordering
class BiggestName(str):
    def __lt__(self, other):
        return True

    def __eq__(self, other):
        return isinstance(other, self.__class__)


def _hook_list() -> List[HookNode]:
    return []


def get_hooks_by_tags(
    nodes: Iterable[ResultNode],
    match_tags: Set[str],
) -> List[HookNode]:
    matched_nodes = []
    for node in nodes:
        if not isinstance(node, HookNode):
            continue
        node_tags = node.tags
        if len(set(node_tags) & match_tags):
            matched_nodes.append(node)
    return matched_nodes


def get_hook(source, index):
    hook_dict = get_hook_dict(source)
    hook_dict.setdefault("index", index)
    Hook.validate(hook_dict)
    return Hook.from_dict(hook_dict)


def get_execution_status(sql: str, adapter: BaseAdapter) -> Tuple[RunStatus, str]:
    if not sql.strip():
        return RunStatus.Success, "OK"

    try:
        response, _ = adapter.execute(sql, auto_begin=False, fetch=False)
        status = RunStatus.Success
        message = response._message
    except (KeyboardInterrupt, SystemExit):
        raise
    except DbtRuntimeError as exc:
        status = RunStatus.Error
        message = exc.msg
    except Exception as exc:
        status = RunStatus.Error
        message = str(exc)

    return (status, message)


def _get_adapter_info(adapter, run_model_result) -> Dict[str, Any]:
    """Each adapter returns a dataclass with a flexible dictionary for
    adapter-specific fields. Only the non-'model_adapter_details' fields
    are guaranteed cross adapter."""
    return asdict(adapter.get_adapter_run_info(run_model_result.node.config)) if adapter else {}


def track_model_run(index, num_nodes, run_model_result, adapter=None):
    if tracking.active_user is None:
        raise DbtInternalError("cannot track model run with no active user")
    invocation_id = get_invocation_id()
    node = run_model_result.node
    has_group = True if hasattr(node, "group") and node.group else False
    if node.resource_type == NodeType.Model:
        access = node.access.value if node.access is not None else None
        contract_enforced = node.contract.enforced
        versioned = True if node.version else False
        incremental_strategy = node.config.incremental_strategy
    else:
        access = None
        contract_enforced = False
        versioned = False
        incremental_strategy = None

    tracking.track_model_run(
        {
            "invocation_id": invocation_id,
            "index": index,
            "total": num_nodes,
            "execution_time": run_model_result.execution_time,
            "run_status": str(run_model_result.status).upper(),
            "run_skipped": run_model_result.status == NodeStatus.Skipped,
            "run_error": run_model_result.status == NodeStatus.Error,
            "model_materialization": node.get_materialization(),
            "model_incremental_strategy": incremental_strategy,
            "model_id": utils.get_hash(node),
            "hashed_contents": utils.get_hashed_contents(node),
            "timing": [t.to_dict(omit_none=True) for t in run_model_result.timing],
            "language": str(node.language),
            "has_group": has_group,
            "contract_enforced": contract_enforced,
            "access": access,
            "versioned": versioned,
            "adapter_info": _get_adapter_info(adapter, run_model_result),
        }
    )


# make sure that we got an ok result back from a materialization
def _validate_materialization_relations_dict(inp: Dict[Any, Any], model) -> List[BaseRelation]:
    try:
        relations_value = inp["relations"]
    except KeyError:
        msg = (
            'Invalid return value from materialization, "relations" '
            "not found, got keys: {}".format(list(inp))
        )
        raise CompilationError(msg, node=model) from None

    if not isinstance(relations_value, list):
        msg = (
            'Invalid return value from materialization, "relations" '
            "not a list, got: {}".format(relations_value)
        )
        raise CompilationError(msg, node=model) from None

    relations: List[BaseRelation] = []
    for relation in relations_value:
        if not isinstance(relation, BaseRelation):
            msg = (
                "Invalid return value from materialization, "
                '"relations" contains non-Relation: {}'.format(relation)
            )
            raise CompilationError(msg, node=model)

        assert isinstance(relation, BaseRelation)
        relations.append(relation)
    return relations


class ModelRunner(CompileRunner):
    def get_node_representation(self):
        display_quote_policy = {"database": False, "schema": False, "identifier": False}
        relation = self.adapter.Relation.create_from(
            self.config, self.node, quote_policy=display_quote_policy
        )
        # exclude the database from output if it's the default
        if self.node.database == self.config.credentials.database:
            relation = relation.include(database=False)
        return str(relation)

    def describe_node(self) -> str:
        # TODO CL 'language' will be moved to node level when we change representation
        return f"{self.node.language} {self.node.get_materialization()} model {self.get_node_representation()}"

    def print_start_line(self):
        fire_event(
            LogStartLine(
                description=self.describe_node(),
                index=self.node_index,
                total=self.num_nodes,
                node_info=self.node.node_info,
            )
        )

    def print_result_line(self, result):
        description = self.describe_node()
        group = group_lookup.get(self.node.unique_id)
        if result.status == NodeStatus.Error:
            status = result.status
            level = EventLevel.ERROR
        else:
            status = result.message
            level = EventLevel.INFO
        fire_event(
            LogModelResult(
                description=description,
                status=status,
                index=self.node_index,
                total=self.num_nodes,
                execution_time=result.execution_time,
                node_info=self.node.node_info,
                group=group,
            ),
            level=level,
        )

    def before_execute(self) -> None:
        self.print_start_line()

    def after_execute(self, result) -> None:
        track_model_run(self.node_index, self.num_nodes, result, adapter=self.adapter)
        self.print_result_line(result)

    def _build_run_model_result(self, model, context, elapsed_time: float = 0.0):
        result = context["load_result"]("main")
        if not result:
            raise DbtRuntimeError("main is not being called during running model")
        adapter_response = {}
        if isinstance(result.response, dbtClassMixin):
            adapter_response = result.response.to_dict(omit_none=True)
        return RunResult(
            node=model,
            status=RunStatus.Success,
            timing=[],
            thread_id=threading.current_thread().name,
            execution_time=elapsed_time,
            message=str(result.response),
            adapter_response=adapter_response,
            failures=result.get("failures"),
            batch_results=None,
        )

    def _materialization_relations(self, result: Any, model) -> List[BaseRelation]:
        if isinstance(result, str):
            msg = (
                'The materialization ("{}") did not explicitly return a '
                "list of relations to add to the cache.".format(str(model.get_materialization()))
            )
            raise CompilationError(msg, node=model)

        if isinstance(result, dict):
            return _validate_materialization_relations_dict(result, model)

        msg = (
            "Invalid return value from materialization, expected a dict "
            'with key "relations", got: {}'.format(str(result))
        )
        raise CompilationError(msg, node=model)

    def _execute_model(
        self,
        hook_ctx: Any,
        context_config: Any,
        model: ModelNode,
        context: Dict[str, Any],
        materialization_macro: MacroProtocol,
    ) -> RunResult:
        try:
            result = MacroGenerator(
                materialization_macro, context, stack=context["context_macro_stack"]
            )()
        finally:
            self.adapter.post_model_hook(context_config, hook_ctx)

        for relation in self._materialization_relations(result, model):
            self.adapter.cache_added(relation.incorporate(dbt_created=True))

        return self._build_run_model_result(model, context)

    def execute(self, model, manifest):
        context = generate_runtime_model_context(model, self.config, manifest)

        materialization_macro = manifest.find_materialization_macro_by_name(
            self.config.project_name, model.get_materialization(), self.adapter.type()
        )

        if materialization_macro is None:
            raise MissingMaterializationError(
                materialization=model.get_materialization(), adapter_type=self.adapter.type()
            )

        if "config" not in context:
            raise DbtInternalError(
                "Invalid materialization context generated, missing config: {}".format(context)
            )
        context_config = context["config"]

        mat_has_supported_langs = hasattr(materialization_macro, "supported_languages")
        model_lang_supported = model.language in materialization_macro.supported_languages
        if mat_has_supported_langs and not model_lang_supported:
            str_langs = [str(lang) for lang in materialization_macro.supported_languages]
            raise DbtValidationError(
                f'Materialization "{materialization_macro.name}" only supports languages {str_langs}; '
                f'got "{model.language}"'
            )

        hook_ctx = self.adapter.pre_model_hook(context_config)

        return self._execute_model(hook_ctx, context_config, model, context, materialization_macro)


class MicrobatchBatchRunner(ModelRunner):
    """Handles the running of individual batches"""

    def __init__(
        self,
        config,
        adapter,
        node,
        node_index: int,
        num_nodes: int,
        batch_idx: int,
        batches: Dict[int, BatchType],
        relation_exists: bool,
        incremental_batch: bool,
    ):
        super().__init__(config, adapter, node, node_index, num_nodes)

        self.batch_idx = batch_idx
        self.batches = batches
        self.relation_exists = relation_exists
        self.incremental_batch = incremental_batch

    def describe_batch(self) -> str:
        batch_start = self.batches[self.batch_idx][0]
        formatted_batch_start = MicrobatchBuilder.format_batch_start(
            batch_start, self.node.config.batch_size
        )
        return f"batch {formatted_batch_start} of {self.get_node_representation()}"

    def print_result_line(self, result: RunResult):
        if result.status == NodeStatus.Error:
            status = result.status
            level = EventLevel.ERROR
        elif result.status == NodeStatus.Skipped:
            status = result.status
            level = EventLevel.INFO
        else:
            status = result.message
            level = EventLevel.INFO

        fire_event(
            LogBatchResult(
                description=self.describe_batch(),
                status=status,
                batch_index=self.batch_idx + 1,
                total_batches=len(self.batches),
                execution_time=result.execution_time,
                node_info=self.node.node_info,
                group=group_lookup.get(self.node.unique_id),
            ),
            level=level,
        )

    def print_start_line(self) -> None:
        fire_event(
            LogStartBatch(
                description=self.describe_batch(),
                batch_index=self.batch_idx + 1,
                total_batches=len(self.batches),
                node_info=self.node.node_info,
            )
        )

    def should_run_in_parallel(self) -> bool:
        if not self.adapter.supports(Capability.MicrobatchConcurrency):
            run_in_parallel = False
        elif not self.relation_exists:
            # If the relation doesn't exist, we can't run in parallel
            run_in_parallel = False
        elif self.node.config.concurrent_batches is not None:
            # If the relation exists and the `concurrent_batches` config isn't None, use the config value
            run_in_parallel = self.node.config.concurrent_batches
        else:
            # If the relation exists, the `concurrent_batches` config is None, check if the model self references `this`.
            # If the model self references `this` then we assume the model batches _can't_ be run in parallel
            run_in_parallel = not self.node.has_this

        return run_in_parallel

    def on_skip(self):
        result = RunResult(
            node=self.node,
            status=RunStatus.Skipped,
            timing=[],
            thread_id=threading.current_thread().name,
            execution_time=0.0,
            message="SKIPPED",
            adapter_response={},
            failures=1,
            batch_results=BatchResults(failed=[self.batches[self.batch_idx]]),
        )
        self.print_result_line(result=result)
        return result

    def compile(self, manifest: Manifest):
        batch = self.batches[self.batch_idx]

        # LEGACY: Set start/end in context prior to re-compiling (Will be removed for 1.10+)
        # TODO: REMOVE before 1.10 GA
        self.node.config["__dbt_internal_microbatch_event_time_start"] = batch[0]
        self.node.config["__dbt_internal_microbatch_event_time_end"] = batch[1]
        # Create batch context on model node prior to re-compiling
        self.node.batch = BatchContext(
            id=MicrobatchBuilder.batch_id(batch[0], self.node.config.batch_size),
            event_time_start=batch[0],
            event_time_end=batch[1],
        )
        # Recompile node to re-resolve refs with event time filters rendered, update context
        self.compiler.compile_node(
            self.node,
            manifest,
            {},
            split_suffix=MicrobatchBuilder.format_batch_start(
                batch[0], self.node.config.batch_size
            ),
        )

        return self.node

    def _build_succesful_run_batch_result(
        self,
        model: ModelNode,
        context: Dict[str, Any],
        batch: BatchType,
        elapsed_time: float = 0.0,
    ) -> RunResult:
        run_result = self._build_run_model_result(model, context, elapsed_time)
        run_result.batch_results = BatchResults(successful=[batch])
        return run_result

    def _build_failed_run_batch_result(
        self,
        model: ModelNode,
        batch: BatchType,
        elapsed_time: float = 0.0,
    ) -> RunResult:
        return RunResult(
            node=model,
            status=RunStatus.Error,
            timing=[],
            thread_id=threading.current_thread().name,
            execution_time=elapsed_time,
            message="ERROR",
            adapter_response={},
            failures=1,
            batch_results=BatchResults(failed=[batch]),
        )

    def _execute_microbatch_materialization(
        self,
        model: ModelNode,
        context: Dict[str, Any],
        materialization_macro: MacroProtocol,
    ) -> RunResult:

        batch = self.batches[self.batch_idx]
        # call materialization_macro to get a batch-level run result
        start_time = time.perf_counter()
        try:
            # Update jinja context with batch context members
            jinja_context = MicrobatchBuilder.build_jinja_context_for_batch(
                model=model,
                incremental_batch=self.incremental_batch,
            )
            context.update(jinja_context)

            # Materialize batch and cache any materialized relations
            result = MacroGenerator(
                materialization_macro, context, stack=context["context_macro_stack"]
            )()
            for relation in self._materialization_relations(result, model):
                self.adapter.cache_added(relation.incorporate(dbt_created=True))

            # Build result of executed batch
            batch_run_result = self._build_succesful_run_batch_result(
                model, context, batch, time.perf_counter() - start_time
            )
            batch_result = batch_run_result

            # At least one batch has been inserted successfully!
            # Can proceed incrementally + in parallel
            self.relation_exists = True

        except (KeyboardInterrupt, SystemExit):
            # reraise it for GraphRunnableTask.execute_nodes to handle
            raise
        except Exception as e:
            fire_event(
                GenericExceptionOnRun(
                    unique_id=self.node.unique_id,
                    exc=f"Exception on worker thread. {str(e)}",
                    node_info=self.node.node_info,
                )
            )
            batch_run_result = self._build_failed_run_batch_result(
                model, batch, time.perf_counter() - start_time
            )

        batch_result = batch_run_result

        return batch_result

    def _execute_model(
        self,
        hook_ctx: Any,
        context_config: Any,
        model: ModelNode,
        context: Dict[str, Any],
        materialization_macro: MacroProtocol,
    ) -> RunResult:
        try:
            batch_result = self._execute_microbatch_materialization(
                model, context, materialization_macro
            )
        finally:
            self.adapter.post_model_hook(context_config, hook_ctx)

        return batch_result


class MicrobatchModelRunner(ModelRunner):
    """Handles the orchestration of batches to run for a given microbatch model"""

    def __init__(self, config, adapter, node, node_index: int, num_nodes: int):
        super().__init__(config, adapter, node, node_index, num_nodes)

        # The parent task is necessary because we need access to the `_submit_batch` and `submit` methods
        self._parent_task: Optional[RunTask] = None
        # The pool is necessary because we need to batches to be executed within the same thread pool
        self._pool: Optional[DbtThreadPool] = None

    def set_parent_task(self, parent_task: RunTask) -> None:
        self._parent_task = parent_task

    def set_pool(self, pool: DbtThreadPool) -> None:
        self._pool = pool

    @property
    def parent_task(self) -> RunTask:
        if self._parent_task is None:
            raise DbtInternalError(
                msg="Tried to access `parent_task` of `MicrobatchModelRunner` before it was set"
            )

        return self._parent_task

    @property
    def pool(self) -> DbtThreadPool:
        if self._pool is None:
            raise DbtInternalError(
                msg="Tried to access `pool` of `MicrobatchModelRunner` before it was set"
            )

        return self._pool

    def _has_relation(self, model: ModelNode) -> bool:
        """Check whether the relation for the model exists in the data warehouse"""
        relation_info = self.adapter.Relation.create_from(self.config, model)
        relation = self.adapter.get_relation(
            relation_info.database, relation_info.schema, relation_info.name
        )
        return relation is not None

    def _is_incremental(self, model) -> bool:
        """Check whether the model should be run `incrementally` or as `full refresh`"""
        # TODO: Remove this whole function. This should be a temporary method. We're working with adapters on
        # a strategy to ensure we can access the `is_incremental` logic without drift
        relation_info = self.adapter.Relation.create_from(self.config, model)
        relation = self.adapter.get_relation(
            relation_info.database, relation_info.schema, relation_info.name
        )
        if (
            relation is not None
            and relation.type == "table"
            and model.config.materialized == "incremental"
        ):
            if model.config.full_refresh is not None:
                return not model.config.full_refresh
            else:
                return not getattr(self.config.args, "FULL_REFRESH", False)
        else:
            return False

    def _initial_run_microbatch_model_result(self, model: ModelNode) -> RunResult:
        return RunResult(
            node=model,
            status=RunStatus.Success,
            timing=[],
            thread_id=threading.current_thread().name,
            # The execution_time here doesn't get propagated to logs because
            # `safe_run_hooks` handles the elapsed time at the node level
            execution_time=0,
            message="",
            adapter_response={},
            failures=0,
            batch_results=BatchResults(),
        )

    def describe_node(self) -> str:
        return f"{self.node.language} microbatch model {self.get_node_representation()}"

    def merge_batch_results(self, result: RunResult, batch_results: List[RunResult]):
        """merge batch_results into result"""
        if result.batch_results is None:
            result.batch_results = BatchResults()

        for batch_result in batch_results:
            if batch_result.batch_results is not None:
                result.batch_results += batch_result.batch_results
            result.execution_time += batch_result.execution_time

        num_successes = len(result.batch_results.successful)
        num_failures = len(result.batch_results.failed)
        if num_failures == 0:
            status = RunStatus.Success
            msg = "SUCCESS"
        elif num_successes == 0:
            status = RunStatus.Error
            msg = "ERROR"
        else:
            status = RunStatus.PartialSuccess
            msg = f"PARTIAL SUCCESS ({num_successes}/{num_successes + num_failures})"
        result.status = status
        result.message = msg

        result.batch_results.successful = sorted(result.batch_results.successful)
        result.batch_results.failed = sorted(result.batch_results.failed)

        # # If retrying, propagate previously successful batches into final result, even thoguh they were not run in this invocation
        if self.node.previous_batch_results is not None:
            result.batch_results.successful += self.node.previous_batch_results.successful

    def _update_result_with_unfinished_batches(
        self, result: RunResult, batches: Dict[int, BatchType]
    ) -> None:
        """This method is really only to be used when the execution of a microbatch model is halted before all batches have had a chance to run"""
        batches_finished: Set[BatchType] = set()

        if result.batch_results:
            # build list of finished batches
            batches_finished = batches_finished.union(set(result.batch_results.successful))
            batches_finished = batches_finished.union(set(result.batch_results.failed))
        else:
            # instantiate `batch_results` if it was `None`
            result.batch_results = BatchResults()

        # skipped batches are any batch that was expected but didn't finish
        batches_expected = {batch for _, batch in batches.items()}
        skipped_batches = batches_expected.difference(batches_finished)

        result.batch_results.failed.extend(list(skipped_batches))

        # We call this method, even though we are merging no new results, as it updates
        # the result witht he appropriate status (Success/Partial/Failed)
        self.merge_batch_results(result, [])

    def get_microbatch_builder(self, model: ModelNode) -> MicrobatchBuilder:
        return MicrobatchBuilder(
            model=model,
            is_incremental=self._is_incremental(model),
            event_time_start=getattr(self.config.args, "EVENT_TIME_START", None),
            event_time_end=getattr(self.config.args, "EVENT_TIME_END", None),
            default_end_time=self.config.invoked_at,
        )

    def get_batches(self, model: ModelNode) -> Dict[int, BatchType]:
        """Get the batches that should be run for the model"""

        # Note currently (02/23/2025) model.previous_batch_results is only ever _not_ `None`
        # IFF `dbt retry` is being run and the microbatch model had batches which
        # failed on the run of the model (which is being retried)
        if model.previous_batch_results is None:
            microbatch_builder = self.get_microbatch_builder(model)
            end = microbatch_builder.build_end_time()
            start = microbatch_builder.build_start_time(end)
            batches = microbatch_builder.build_batches(start, end)
        else:
            batches = model.previous_batch_results.failed

        return {batch_idx: batches[batch_idx] for batch_idx in range(len(batches))}

    def compile(self, manifest: Manifest):
        """Don't do anything here because this runner doesn't need to compile anything"""
        return self.node

    def execute(self, model: ModelNode, manifest: Manifest) -> RunResult:
        # Execution really means orchestration in this case

        batches = self.get_batches(model=model)
        relation_exists = self._has_relation(model=model)
        result = self._initial_run_microbatch_model_result(model=model)

        # No batches to run, so return initial result
        if len(batches) == 0:
            return result

        batch_results: List[RunResult] = []
        batch_idx = 0

        # Run first batch not in parallel
        relation_exists = self.parent_task._submit_batch(
            node=model,
            adapter=self.adapter,
            relation_exists=relation_exists,
            batches=batches,
            batch_idx=batch_idx,
            batch_results=batch_results,
            pool=self.pool,
            force_sequential_run=True,
            incremental_batch=self._is_incremental(model=model),
        )
        batch_idx += 1
        skip_batches = batch_results[0].status != RunStatus.Success

        # Run all batches except first and last batch, in parallel if possible
        while batch_idx < len(batches) - 1:
            relation_exists = self.parent_task._submit_batch(
                node=model,
                adapter=self.adapter,
                relation_exists=relation_exists,
                batches=batches,
                batch_idx=batch_idx,
                batch_results=batch_results,
                pool=self.pool,
                skip=skip_batches,
            )
            batch_idx += 1

        # Wait until all submitted batches have completed
        while len(batch_results) != batch_idx:
            # Check if the pool was closed, because if it was, then the main thread is trying to exit.
            # If the main thread is trying to exit, we need to shutdown. If we _don't_ shutdown, then
            # batches will continue to execute and we'll delay the run from stopping
            if self.pool.is_closed():
                # It's technically possible for more results to come in while we clean up
                # instead we're going to say the didn't finish, regardless of if they finished
                # or not. Thus, lets get a copy of the results as they exist right "now".
                frozen_batch_results = deepcopy(batch_results)
                self.merge_batch_results(result, frozen_batch_results)
                self._update_result_with_unfinished_batches(result, batches)
                return result

            # breifly sleep so that this thread doesn't go brrrrr while waiting
            time.sleep(0.1)

        # Only run "last" batch if there is more than one batch
        if len(batches) != 1:
            # Final batch runs once all others complete to ensure post_hook runs at the end
            self.parent_task._submit_batch(
                node=model,
                adapter=self.adapter,
                relation_exists=relation_exists,
                batches=batches,
                batch_idx=batch_idx,
                batch_results=batch_results,
                pool=self.pool,
                force_sequential_run=True,
                skip=skip_batches,
            )

        # Finalize run: merge results, track model run, and print final result line
        self.merge_batch_results(result, batch_results)

        return result


class RunTask(CompileTask):
    def __init__(
        self,
        args: Flags,
        config: RuntimeConfig,
        manifest: Manifest,
        batch_map: Optional[Dict[str, BatchResults]] = None,
    ) -> None:
        super().__init__(args, config, manifest)
        self.batch_map = batch_map

    def raise_on_first_error(self) -> bool:
        return False

    def get_hook_sql(self, adapter, hook, idx, num_hooks, extra_context) -> str:
        if self.manifest is None:
            raise DbtInternalError("compile_node called before manifest was loaded")

        compiled = self.compiler.compile_node(hook, self.manifest, extra_context)
        statement = compiled.compiled_code
        hook_index = hook.index or num_hooks
        hook_obj = get_hook(statement, index=hook_index)
        return hook_obj.sql or ""

    def handle_job_queue(self, pool, callback):
        node = self.job_queue.get()
        self._raise_set_error()
        runner = self.get_runner(node)
        # we finally know what we're running! Make sure we haven't decided
        # to skip it due to upstream failures
        if runner.node.unique_id in self._skipped_children:
            cause = self._skipped_children.pop(runner.node.unique_id)
            runner.do_skip(cause=cause)

        if isinstance(runner, MicrobatchModelRunner):
            runner.set_parent_task(self)
            runner.set_pool(pool)

        args = [runner]
        self._submit(pool, args, callback)

    def _submit_batch(
        self,
        node: ModelNode,
        adapter: BaseAdapter,
        relation_exists: bool,
        batches: Dict[int, BatchType],
        batch_idx: int,
        batch_results: List[RunResult],
        pool: DbtThreadPool,
        force_sequential_run: bool = False,
        skip: bool = False,
        incremental_batch: bool = True,
    ):
        node_copy = deepcopy(node)
        # Only run pre_hook(s) for first batch
        if batch_idx != 0:
            node_copy.config.pre_hook = []

        # Only run post_hook(s) for last batch
        if batch_idx != len(batches) - 1:
            node_copy.config.post_hook = []

        # TODO: We should be doing self.get_runner, however doing so
        # currently causes the tracking of how many nodes there are to
        # increment when we don't want it to
        batch_runner = MicrobatchBatchRunner(
            self.config,
            adapter,
            node_copy,
            self.run_count,
            self.num_nodes,
            batch_idx,
            batches,
            relation_exists,
            incremental_batch,
        )

        if skip:
            batch_runner.do_skip()

        if not pool.is_closed():
            if not force_sequential_run and batch_runner.should_run_in_parallel():
                fire_event(
                    MicrobatchExecutionDebug(
                        msg=f"{batch_runner.describe_batch()} is being run concurrently"
                    )
                )
                self._submit(pool, [batch_runner], batch_results.append)
            else:
                fire_event(
                    MicrobatchExecutionDebug(
                        msg=f"{batch_runner.describe_batch()} is being run sequentially"
                    )
                )
                batch_results.append(self.call_runner(batch_runner))
                relation_exists = batch_runner.relation_exists
        else:
            batch_results.append(
                batch_runner._build_failed_run_batch_result(node_copy, batches[batch_idx])
            )

        return relation_exists

    def _hook_keyfunc(self, hook: HookNode) -> Tuple[str, Optional[int]]:
        package_name = hook.package_name
        if package_name == self.config.project_name:
            package_name = BiggestName("")
        return package_name, hook.index

    def get_hooks_by_type(self, hook_type: RunHookType) -> List[HookNode]:

        if self.manifest is None:
            raise DbtInternalError("self.manifest was None in get_hooks_by_type")

        nodes = self.manifest.nodes.values()
        # find all hooks defined in the manifest (could be multiple projects)
        hooks: List[HookNode] = get_hooks_by_tags(nodes, {hook_type})
        hooks.sort(key=self._hook_keyfunc)
        return hooks

    def safe_run_hooks(
        self, adapter: BaseAdapter, hook_type: RunHookType, extra_context: Dict[str, Any]
    ) -> RunStatus:
        ordered_hooks = self.get_hooks_by_type(hook_type)

        if hook_type == RunHookType.End and ordered_hooks:
            fire_event(Formatting(""))

        # on-run-* hooks should run outside a transaction. This happens because psycopg2 automatically begins a transaction when a connection is created.
        adapter.clear_transaction()
        if not ordered_hooks:
            return RunStatus.Success

        status = RunStatus.Success
        failed = False
        num_hooks = len(ordered_hooks)

        for idx, hook in enumerate(ordered_hooks, 1):
            with log_contextvars(node_info=hook.node_info):
                hook.index = idx
                hook_name = f"{hook.package_name}.{hook_type}.{hook.index - 1}"
                execution_time = 0.0
                timing: List[TimingInfo] = []
                failures = 1

                if not failed:
                    with collect_timing_info("compile", timing.append):
                        sql = self.get_hook_sql(
                            adapter, hook, hook.index, num_hooks, extra_context
                        )

                    started_at = timing[0].started_at or datetime.utcnow()
                    hook.update_event_status(
                        started_at=started_at.isoformat(), node_status=RunningStatus.Started
                    )

                    fire_event(
                        LogHookStartLine(
                            statement=hook_name,
                            index=hook.index,
                            total=num_hooks,
                            node_info=hook.node_info,
                        )
                    )

                    with collect_timing_info("execute", timing.append):
                        status, message = get_execution_status(sql, adapter)

                    finished_at = timing[1].completed_at or datetime.utcnow()
                    hook.update_event_status(finished_at=finished_at.isoformat())
                    execution_time = (finished_at - started_at).total_seconds()
                    failures = 0 if status == RunStatus.Success else 1

                    if status == RunStatus.Success:
                        message = f"{hook_name} passed"
                    else:
                        message = f"{hook_name} failed, error:\n {message}"
                        failed = True
                else:
                    status = RunStatus.Skipped
                    message = f"{hook_name} skipped"

                hook.update_event_status(node_status=status)

                self.node_results.append(
                    RunResult(
                        status=status,
                        thread_id="main",
                        timing=timing,
                        message=message,
                        adapter_response={},
                        execution_time=execution_time,
                        failures=failures,
                        node=hook,
                    )
                )

                fire_event(
                    LogHookEndLine(
                        statement=hook_name,
                        status=status,
                        index=hook.index,
                        total=num_hooks,
                        execution_time=execution_time,
                        node_info=hook.node_info,
                    )
                )

        if hook_type == RunHookType.Start and ordered_hooks:
            fire_event(Formatting(""))

        return status

    def print_results_line(self, results, execution_time) -> None:
        nodes = [r.node for r in results if hasattr(r, "node")]
        stat_line = get_counts(nodes)

        execution = ""

        if execution_time is not None:
            execution = utils.humanize_execution_time(execution_time=execution_time)

        fire_event(Formatting(""))
        fire_event(
            FinishedRunningStats(
                stat_line=stat_line, execution=execution, execution_time=execution_time
            )
        )

    def populate_microbatch_batches(self, selected_uids: AbstractSet[str]):
        if self.batch_map is not None and self.manifest is not None:
            for uid in selected_uids:
                if uid in self.batch_map:
                    node = self.manifest.ref_lookup.perform_lookup(uid, self.manifest)
                    if isinstance(node, ModelNode):
                        node.previous_batch_results = self.batch_map[uid]

    def before_run(self, adapter: BaseAdapter, selected_uids: AbstractSet[str]) -> RunStatus:
        with adapter.connection_named("master"):
            self.defer_to_manifest()
            required_schemas = self.get_model_schemas(adapter, selected_uids)
            self.create_schemas(adapter, required_schemas)
            self.populate_adapter_cache(adapter, required_schemas)
            self.populate_microbatch_batches(selected_uids)
            group_lookup.init(self.manifest, selected_uids)
            run_hooks_status = self.safe_run_hooks(adapter, RunHookType.Start, {})
            return run_hooks_status

    def after_run(self, adapter, results) -> None:
        # in on-run-end hooks, provide the value 'database_schemas', which is a
        # list of unique (database, schema) pairs that successfully executed
        # models were in. For backwards compatibility, include the old
        # 'schemas', which did not include database information.

        database_schema_set: Set[Tuple[Optional[str], str]] = {
            (r.node.database, r.node.schema)
            for r in results
            if (hasattr(r, "node") and r.node.is_relational)
            and r.status not in (NodeStatus.Error, NodeStatus.Fail, NodeStatus.Skipped)
        }

        extras = {
            "schemas": list({s for _, s in database_schema_set}),
            "results": [
                r for r in results if r.thread_id != "main" or r.status == RunStatus.Error
            ],  # exclude that didn't fail to preserve backwards compatibility
            "database_schemas": list(database_schema_set),
        }

        try:
            with adapter.connection_named("master"):
                self.safe_run_hooks(adapter, RunHookType.End, extras)
        except (KeyboardInterrupt, SystemExit, DbtRuntimeError):
            run_result = self.get_result(
                results=self.node_results,
                elapsed_time=time.time() - self.started_at,
                generated_at=datetime.utcnow(),
            )

            if self.args.write_json and hasattr(run_result, "write"):
                run_result.write(self.result_path())

            print_run_end_messages(self.node_results, keyboard_interrupt=True)

            raise

    def get_node_selector(self) -> ResourceTypeSelector:
        if self.manifest is None or self.graph is None:
            raise DbtInternalError("manifest and graph must be set to get perform node selection")
        return ResourceTypeSelector(
            graph=self.graph,
            manifest=self.manifest,
            previous_state=self.previous_state,
            resource_types=[NodeType.Model],
        )

    def get_runner_type(self, node) -> Optional[Type[BaseRunner]]:
        if self.manifest is None:
            raise DbtInternalError("manifest must be set prior to calling get_runner_type")

        if (
            node.config.materialized == "incremental"
            and node.config.incremental_strategy == "microbatch"
            and self.manifest.use_microbatch_batches(project_name=self.config.project_name)
        ):
            return MicrobatchModelRunner
        else:
            return ModelRunner

    def task_end_messages(self, results) -> None:
        if results:
            print_run_end_messages(results)
