From 47da4734387392bd011276cc3ff9214e6e36ca34 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Edgar=20Ram=C3=ADrez-Mondrag=C3=B3n?= Date: Tue, 1 Oct 2024 23:33:15 -0600 Subject: [PATCH] Fix `post_model_hook` in data and unit tests --- core/dbt/task/base.py | 2 +- core/dbt/task/run.py | 2 +- core/dbt/task/test.py | 8 ++++++-- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/core/dbt/task/base.py b/core/dbt/task/base.py index 4b0b29db4b1..11273fe8a12 100644 --- a/core/dbt/task/base.py +++ b/core/dbt/task/base.py @@ -6,7 +6,7 @@ from contextlib import nullcontext from datetime import datetime from pathlib import Path -from typing import Any, Dict, List, Optional, Set +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set import dbt.exceptions import dbt_common.exceptions.base diff --git a/core/dbt/task/run.py b/core/dbt/task/run.py index 93a8e261b50..25077aea1d7 100644 --- a/core/dbt/task/run.py +++ b/core/dbt/task/run.py @@ -427,7 +427,7 @@ def _execute_microbatch_model( else: return self._build_run_model_result(model, context) - def execute(self, model, manifest): + def execute(self, model: ModelNode, manifest: Manifest): context = generate_runtime_model_context(model, self.config, manifest) materialization_macro = manifest.find_materialization_macro_by_name( diff --git a/core/dbt/task/test.py b/core/dbt/task/test.py index 6580b85879b..4c614399a49 100644 --- a/core/dbt/task/test.py +++ b/core/dbt/task/test.py @@ -159,13 +159,15 @@ def execute_data_test(self, data_test: TestNode, manifest: Manifest) -> TestResu "Invalid materialization context generated, missing config: {}".format(context) ) + context_config = context["config"] + # generate materialization macro macro_func = MacroGenerator(materialization_macro, context) try: # execute materialization macro macro_func() finally: - self.adapter.post_model_hook(context, hook_ctx) + self.adapter.post_model_hook(context_config, hook_ctx) # load results from context # could eventually be returned directly by materialization @@ -238,6 +240,8 @@ def execute_unit_test( "Invalid materialization context generated, missing config: {}".format(context) ) + context_config = context["config"] + # generate materialization macro macro_func = MacroGenerator(materialization_macro, context) try: @@ -249,7 +253,7 @@ def execute_unit_test( f"There may be an error in the unit test definition: check the data types.\n {e}" ) finally: - self.adapter.post_model_hook(context, hook_ctx) + self.adapter.post_model_hook(context_config, hook_ctx) # load results from context # could eventually be returned directly by materialization