from __future__ import annotations

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

from dvt import tracking, utils
from dvt.adapters.base import BaseAdapter, BaseRelation
from dvt.adapters.capability import Capability
from dvt.adapters.factory import get_adapter
from dvt.adapters.events.types import FinishedRunningStats
from dvt.adapters.exceptions import MissingMaterializationError
from dvt.artifacts.resources import Hook
from dvt.artifacts.schemas.batch_results import BatchResults, BatchType
from dvt.artifacts.schemas.results import (
    NodeStatus,
    RunningStatus,
    RunStatus,
    TimingInfo,
    collect_timing_info,
)
from dvt.artifacts.schemas.run import RunResult
from dvt.cli.flags import Flags
from dvt.clients.jinja import MacroGenerator
from dvt.config import RuntimeConfig
from dvt.context.providers import generate_runtime_model_context
from dvt.contracts.graph.manifest import Manifest
from dvt.contracts.graph.nodes import BatchContext, HookNode, ModelNode, ResultNode
from dvt.events.types import (
    GenericExceptionOnRun,
    LogBatchResult,
    LogHookEndLine,
    LogHookStartLine,
    LogModelResult,
    LogStartBatch,
    LogStartLine,
    MicrobatchExecutionDebug,
)
from dvt.exceptions import CompilationError, DbtInternalError, DbtRuntimeError
from dvt.graph import ResourceTypeSelector
from dvt.graph.thread_pool import DvtThreadPool
from dvt.hooks import get_hook_dict
from dvt.materializations.incremental.microbatch import MicrobatchBuilder
from dvt.node_types import NodeType, RunHookType
from dvt.task import group_lookup
from dvt.task.base import BaseRunner
from dvt.task.compile import CompileRunner, CompileTask
from dvt.task.printer import get_counts, print_run_end_messages
from dvt.utils.artifact_upload import add_artifact_produced
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
from dbt_common.invocation import get_invocation_started_at


@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 error_result(self, node, message, start_time, timing_info):
        """Necessary to return a result with a batch result

        Called by `BaseRunner.safe_run` when an error occurs
        """
        return self._build_run_result(
            node=node,
            start_time=start_time,
            status=RunStatus.Error,
            timing_info=timing_info,
            message=message,
            batch_results=BatchResults(failed=[self.batches[self.batch_idx]]),
        )

    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[DvtThreadPool] = None

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

    def set_pool(self, pool: DvtThreadPool) -> 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) -> DvtThreadPool:
        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:
        # Intially set the start/end to values from args
        event_time_start = getattr(self.config.args, "EVENT_TIME_START", None)
        event_time_end = getattr(self.config.args, "EVENT_TIME_END", None)

        # If we're in sample mode, alter start/end to sample values
        if getattr(self.config.args, "SAMPLE", None) is not None:
            event_time_start = self.config.args.sample.start
            event_time_end = self.config.args.sample.end

        return MicrobatchBuilder(
            model=model,
            is_incremental=self._is_incremental(model),
            event_time_start=event_time_start,
            event_time_end=event_time_end,
            default_end_time=get_invocation_started_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: DvtThreadPool, callback: Callable) -> None:
        assert self.job_queue is not None
        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]

        # Ensure we don't submit more MicrobatchModelRunners than the pool allows
        self._maybe_wait_for_microbatch(pool, runner)
        self._submit(pool, args, callback)

    def _maybe_wait_for_microbatch(
        self, pool: DvtThreadPool, runner: BaseRunner
    ) -> None:
        """
        Checks if the runner is a MicrobatchModelRunner, and waits until the
        number of microbatch models in progress is less than the max number
        of microbatch models allowed by the pool.
        """

        if isinstance(runner, MicrobatchModelRunner) and self.job_queue is not None:
            fire_event(
                MicrobatchExecutionDebug(
                    msg=f"Waiting for microbatch model to be run: {runner.node.name}.\n\tpool.max_microbatch_models: {pool.max_microbatch_models}\n\tlen(self.job_queue.in_progress_microbatch): {len(self.job_queue.in_progress_microbatch)}"
                )
            )
            while pool.max_microbatch_models < len(
                self.job_queue.in_progress_microbatch
            ):
                time.sleep(0.1)

        return

    def _submit_batch(
        self,
        node: ModelNode,
        adapter: BaseAdapter,
        relation_exists: bool,
        batches: Dict[int, BatchType],
        batch_idx: int,
        batch_results: List[RunResult],
        pool: DvtThreadPool,
        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():
            # Only run the batch in parallel IFF:
            # 1. The batch runner is not forced to run sequentially
            # 2. The batch runner should be run in parallel
            # 3. There are available threads in the pool
            #   a. This prevents deadlocks from occurring --threads=1
            if (
                not force_sequential_run
                and batch_runner.should_run_in_parallel()
                and self.job_queue is not None
                and len(self.job_queue.in_progress) < pool.max_threads
            ):
                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.now(
                        timezone.utc
                    ).replace(tzinfo=None)
                    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.now(
                        timezone.utc
                    ).replace(tzinfo=None)
                    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.now(timezone.utc).replace(tzinfo=None),
            )

            if self.args.write_json and hasattr(run_result, "write"):
                run_result.write(self.result_path())
                add_artifact_produced(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)
