add macro_context_generator on adapter

.changes/unreleased/Under the Hood-20231208-004854.yaml

@@ -0,0 +1,6 @@
+kind: Under the Hood
+body: add macro_context_generator on adapter
+time: 2023-12-08T00:48:54.506911+09:00
+custom:
+  Author: michelleark
+  Issue: "9247"

core/dbt/adapters/base/impl.py

@@ -54,7 +54,10 @@
     MacroResultError,
 )
 
-from dbt.adapters.protocol import AdapterConfig
+from dbt.adapters.protocol import (
+    AdapterConfig,
+    MacroContextGeneratorCallable,
+)
 from dbt.common.clients.agate_helper import (
     empty_table,
     get_column_value_uncased,
@@ -76,7 +79,11 @@
 from dbt.common.utils import filter_null_values, executor, cast_to_str, AttrDict
 
 from dbt.adapters.contracts.relation import RelationConfig
-from dbt.adapters.base.connections import Connection, AdapterResponse, BaseConnectionManager
+from dbt.adapters.base.connections import (
+    Connection,
+    AdapterResponse,
+    BaseConnectionManager,
+)
 from dbt.adapters.base.meta import AdapterMeta, available
 from dbt.adapters.base.relation import (
     ComponentName,
@@ -258,6 +265,7 @@
         self.cache = RelationsCache(log_cache_events=config.log_cache_events)
         self.connections = self.ConnectionManager(config, mp_context)
         self._macro_resolver: Optional[MacroResolverProtocol] = None
+        self._macro_context_generator: Optional[MacroContextGeneratorCallable] = None
 
     ###
     # Methods to set / access a macro resolver
@@ -269,9 +277,15 @@
         return self._macro_resolver
 
     def clear_macro_resolver(self) -> None:
         if self._macro_resolver is not None:
             self._macro_resolver = None
 
+    def set_macro_context_generator(
+        self,
+        macro_context_generator: MacroContextGeneratorCallable,
+    ) -> None:
+        self._macro_context_generator = macro_context_generator
+
     ###
     # Methods that pass through to the connection manager
     ###
@@ -1057,7 +1071,10 @@
 
         resolver = macro_resolver or self._macro_resolver
         if resolver is None:
-            raise DbtInternalError("macro resolver was None when calling execute_macro!")
+            raise DbtInternalError("Macro resolver was None when calling execute_macro!")
+
+        if self._macro_context_generator is None:
+            raise DbtInternalError("Macro context generator was None when calling execute_macro!")
 
         macro = resolver.find_macro_by_name(macro_name, self.config.project_name, project)
         if macro is None:
@@ -1071,17 +1088,8 @@
                     macro_name, package_name
                 )
             )
-        # This causes a reference cycle, as generate_runtime_macro_context()
-        # ends up calling get_adapter, so the import has to be here.
-        from dbt.context.providers import generate_runtime_macro_context
-
-        macro_context = generate_runtime_macro_context(
-            # TODO CT-211
-            macro=macro,
-            config=self.config,
-            manifest=resolver,  # type: ignore[arg-type]
-            package_name=project,
-        )
+
+        macro_context = self._macro_context_generator(macro, self.config, resolver, project)
         macro_context.update(context_override)
 
         macro_function = CallableMacroGenerator(macro, macro_context)

core/dbt/adapters/protocol.py

@@ -19,6 +19,7 @@
 from dbt.adapters.contracts.macros import MacroResolverProtocol
 from dbt.adapters.contracts.relation import Policy, HasQuoting, RelationConfig
 from dbt.common.contracts.config.base import BaseConfig
+from dbt.common.clients.jinja import MacroProtocol
 
 
 @dataclass
@@ -55,6 +56,17 @@
 Column_T = TypeVar("Column_T", bound=ColumnProtocol)
 
 
+class MacroContextGeneratorCallable(Protocol):
+    def __call__(
+        self,
+        macro_protocol: MacroProtocol,
+        config: AdapterRequiredConfig,
+        macro_resolver: MacroResolverProtocol,
+        package_name: Optional[str],
+    ) -> Dict[str, Any]:
+        ...
+
+
 # TODO CT-211
 class AdapterProtocol(  # type: ignore[misc]
     Protocol,
@@ -78,14 +90,20 @@
         ...
 
     def set_macro_resolver(self, macro_resolver: MacroResolverProtocol) -> None:
         ...
 
     def get_macro_resolver(self) -> Optional[MacroResolverProtocol]:
         ...
 
     def clear_macro_resolver(self) -> None:
         ...
 
+    def set_macro_context_generator(
+        self,
+        macro_context_generator: MacroContextGeneratorCallable,
+    ) -> None:
+        ...
+
     @classmethod
     def type(cls) -> str:
         pass

core/dbt/cli/requires.py

@@ -1,7 +1,7 @@
 import dbt.tracking
 from dbt.common.invocation import reset_invocation_id
 from dbt.version import installed as installed_version
-from dbt.adapters.factory import adapter_management, register_adapter
+from dbt.adapters.factory import adapter_management, register_adapter, get_adapter
 from dbt.flags import set_flags, get_flag_dict
 from dbt.cli.exceptions import (
     ExceptionExit,
@@ -10,6 +10,7 @@
 from dbt.cli.flags import Flags
 from dbt.config import RuntimeConfig
 from dbt.config.runtime import load_project, load_profile, UnsetProfile
+from dbt.context.providers import generate_runtime_macro_context
 
 from dbt.common.events.base_types import EventLevel
 from dbt.common.events.functions import (
@@ -274,6 +275,8 @@ def wrapper(*args, **kwargs):
 
             runtime_config = ctx.obj["runtime_config"]
             register_adapter(runtime_config)
+            adapter = get_adapter(runtime_config)
+            adapter.set_macro_context_generator(generate_runtime_macro_context)
 
             # a manifest has already been set on the context, so don't overwrite it
             if ctx.obj.get("manifest") is None:

core/dbt/parser/manifest.py

@@ -19,6 +19,7 @@
 from itertools import chain
 import time
 
+from dbt.context.manifest import generate_query_header_context
 from dbt.contracts.graph.semantic_manifest import SemanticManifest
 from dbt.common.events.base_types import EventLevel
 import json
@@ -72,7 +73,6 @@
 )
 from dbt.config import Project, RuntimeConfig
 from dbt.context.docs import generate_runtime_docs_context
-from dbt.context.manifest import generate_query_header_context
 from dbt.context.macro_resolver import MacroResolver, TestMacroNamespace
 from dbt.context.configured import generate_macro_context
 from dbt.context.providers import ParseProvider
@@ -237,7 +237,7 @@
         self,
         root_project: RuntimeConfig,
         all_projects: Mapping[str, Project],
-        macro_hook: Optional[Callable[[Dict[str, Any]], Any]] = None,
+        macro_hook: Optional[Callable[[Manifest], Any]] = None,
         file_diff: Optional[FileDiff] = None,
     ) -> None:
         self.root_project: RuntimeConfig = root_project
@@ -251,9 +251,9 @@
         # This is a MacroQueryStringSetter callable, which is called
         # later after we set the MacroManifest in the adapter. It sets
         # up the query headers.
-        self.macro_hook: Callable[[Dict[str, Any]], Any]
+        self.macro_hook: Callable[[Manifest], Any]
         if macro_hook is None:
-            self.macro_hook = lambda c: None
+            self.macro_hook = lambda m: None
         else:
             self.macro_hook = macro_hook
 
@@ -287,7 +287,7 @@
         # the config and adapter may be persistent.
         if reset:
             config.clear_dependencies()
             adapter.clear_macro_resolver()
         macro_hook = adapter.connections.set_query_header
 
         flags = get_flags()
@@ -1002,6 +1002,8 @@
     def save_macros_to_adapter(self, adapter):
         macro_manifest = MacroManifest(self.manifest.macros)
         adapter.set_macro_resolver(macro_manifest)
+        # This executes the callable macro_hook and sets the
+        # query headers
         # This executes the callable macro_hook and sets the query headers
         query_header_context = generate_query_header_context(adapter.config, macro_manifest)
         self.macro_hook(query_header_context)
@@ -1033,7 +1035,7 @@
     def load_macros(
         cls,
         root_config: RuntimeConfig,
-        macro_hook: Callable[[Dict[str, Any]], Any],
+        macro_hook: Callable[[Manifest], Any],
         base_macros_only=False,
     ) -> Manifest:
         with PARSING_STATE:

core/dbt/tests/fixtures/project.py

@@ -9,6 +9,7 @@
 
 from dbt.parser.manifest import ManifestLoader
 from dbt.common.exceptions import CompilationError, DbtDatabaseError
+from dbt.context.providers import generate_runtime_macro_context
 import dbt.flags as flags
 from dbt.config.runtime import RuntimeConfig
 from dbt.adapters.factory import get_adapter, register_adapter, reset_adapters, get_adapter_by_type
@@ -297,6 +298,7 @@ def adapter(
     )
 
     adapter.set_macro_resolver(manifest)
+    adapter.set_macro_context_generator(generate_runtime_macro_context)
     yield adapter
     adapter.cleanup_connections()
     reset_adapters()

tests/functional/artifacts/data/results/v6/run_results.json

@@ -1 +1 @@
-{"metadata": {"dbt_schema_version": "https://schemas.getdbt.com/dbt/run-results/v6.json", "dbt_version": "1.8.0a1", "generated_at": "2023-12-06T23:58:24.592514Z", "invocation_id": "e1267ed0-3444-45b0-89d6-789d548c1b1a", "env": {}}, "results": [{"status": "success", "timing": [{"name": "compile", "started_at": "2023-12-06T23:58:24.505025Z", "completed_at": "2023-12-06T23:58:24.513104Z"}, {"name": "execute", "started_at": "2023-12-06T23:58:24.515046Z", "completed_at": "2023-12-06T23:58:24.570681Z"}], "thread_id": "Thread-99 (worker)", "execution_time": 0.07069063186645508, "adapter_response": {"_message": "CREATE VIEW", "code": "CREATE VIEW", "rows_affected": -1}, "message": "CREATE VIEW", "failures": null, "unique_id": "model.test.metricflow_time_spine", "compiled": true, "compiled_code": "SELECT to_date('02/20/2023', 'mm/dd/yyyy') as date_day", "relation_name": "\"dbt\".\"test17019071022207203621_test_previous_version_state\".\"metricflow_time_spine\""}, {"status": "success", "timing": [{"name": "compile", "started_at": "2023-12-06T23:58:24.509989Z", "completed_at": "2023-12-06T23:58:24.513530Z"}, {"name": "execute", "started_at": "2023-12-06T23:58:24.517941Z", "completed_at": "2023-12-06T23:58:24.571727Z"}], "thread_id": "Thread-100 (worker)", "execution_time": 0.07077383995056152, "adapter_response": {"_message": "CREATE VIEW", "code": "CREATE VIEW", "rows_affected": -1}, "message": "CREATE VIEW", "failures": null, "unique_id": "model.test.my_model", "compiled": true, "compiled_code": "select 1 as id", "relation_name": "\"dbt\".\"test17019071022207203621_test_previous_version_state\".\"my_model\""}], "elapsed_time": 0.6561338901519775, "args": {"log_file_max_bytes": 10485760, "use_colors_file": true, "strict_mode": false, "cache_selected_only": false, "enable_legacy_logger": false, "log_level_file": "debug", "vars": {}, "project_dir": "/private/var/folders/k6/gtt07v8j2vn51m_z05xk_fjc0000gp/T/pytest-of-michelleark/pytest-955/project7", "which": "run", "defer": false, "write_json": true, "log_path": "/Users/michelleark/src/dbt-core/logs/test17019071022207203621", "log_format": "default", "log_level": "info", "invocation_command": "dbt tests/functional/artifacts", "printer_width": 80, "favor_state": false, "exclude": [], "partial_parse_file_diff": true, "profiles_dir": "/private/var/folders/k6/gtt07v8j2vn51m_z05xk_fjc0000gp/T/pytest-of-michelleark/pytest-955/profile7", "use_colors": true, "macro_debugging": false, "introspect": true, "quiet": false, "show_resource_report": false, "select": [], "version_check": true, "populate_cache": true, "empty": false, "print": true, "send_anonymous_usage_stats": false, "log_format_file": "debug", "indirect_selection": "eager", "partial_parse": true, "warn_error_options": {"include": [], "exclude": []}, "static_parser": true}}
+{"metadata": {"dbt_schema_version": "https://schemas.getdbt.com/dbt/run-results/v6.json", "dbt_version": "1.8.0a1", "generated_at": "2023-12-08T18:33:26.385305Z", "invocation_id": "74d0d61f-8496-4661-9d15-1a48f395508b", "env": {}}, "results": [{"status": "success", "timing": [{"name": "compile", "started_at": "2023-12-08T18:33:26.238276Z", "completed_at": "2023-12-08T18:33:26.245635Z"}, {"name": "execute", "started_at": "2023-12-08T18:33:26.246342Z", "completed_at": "2023-12-08T18:33:26.361436Z"}], "thread_id": "Thread-107 (worker)", "execution_time": 0.12888598442077637, "adapter_response": {"_message": "CREATE VIEW", "code": "CREATE VIEW", "rows_affected": -1}, "message": "CREATE VIEW", "failures": null, "unique_id": "model.test.my_model", "compiled": true, "compiled_code": "select 1 as id", "relation_name": "\"dbt\".\"test17020604041734807591_test_previous_version_state\".\"my_model\""}, {"status": "success", "timing": [{"name": "compile", "started_at": "2023-12-08T18:33:26.235770Z", "completed_at": "2023-12-08T18:33:26.240830Z"}, {"name": "execute", "started_at": "2023-12-08T18:33:26.241467Z", "completed_at": "2023-12-08T18:33:26.371786Z"}], "thread_id": "Thread-106 (worker)", "execution_time": 0.13865900039672852, "adapter_response": {"_message": "CREATE VIEW", "code": "CREATE VIEW", "rows_affected": -1}, "message": "CREATE VIEW", "failures": null, "unique_id": "model.test.metricflow_time_spine", "compiled": true, "compiled_code": "SELECT to_date('02/20/2023', 'mm/dd/yyyy') as date_day", "relation_name": "\"dbt\".\"test17020604041734807591_test_previous_version_state\".\"metricflow_time_spine\""}], "elapsed_time": 0.27949094772338867, "args": {"partial_parse_file_diff": true, "populate_cache": true, "log_path": "/Users/colinrogers/projects/dbt-core/logs/test17020604041734807591", "project_dir": "/private/var/folders/5_/1cykfnvj641gf4rsdwj6v8p00000gp/T/pytest-of-colinrogers/pytest-12/project12", "profiles_dir": "/private/var/folders/5_/1cykfnvj641gf4rsdwj6v8p00000gp/T/pytest-of-colinrogers/pytest-12/profile12", "write_json": true, "log_file_max_bytes": 10485760, "defer": false, "use_colors": true, "version_check": true, "printer_width": 80, "log_level": "info", "cache_selected_only": false, "enable_legacy_logger": false, "exclude": [], "send_anonymous_usage_stats": false, "log_format_file": "debug", "log_level_file": "debug", "warn_error_options": {"include": [], "exclude": []}, "favor_state": false, "use_colors_file": true, "static_parser": true, "log_format": "default", "macro_debugging": false, "introspect": true, "print": true, "select": [], "empty": false, "invocation_command": "dbt --cov=core --cov-append --cov-report=xml tests/functional -k not tests/functional/graph_selection", "show_resource_report": false, "vars": {}, "quiet": false, "partial_parse": true, "indirect_selection": "eager", "which": "run", "strict_mode": false}}

tests/unit/test_postgres_adapter.py

@@ -8,17 +8,17 @@
 
 from dbt.adapters.base import BaseRelation
 from dbt.adapters.contracts.relation import Path
+from dbt.context.manifest import generate_query_header_context
 from dbt.task.debug import DebugTask
 
-from dbt.adapters.base.query_headers import MacroQueryStringSetter
 from dbt.adapters.postgres import PostgresAdapter
 from dbt.adapters.postgres import Plugin as PostgresPlugin
 from dbt.contracts.files import FileHash
 from dbt.contracts.graph.manifest import ManifestStateCheck
-from dbt.context.manifest import generate_query_header_context
 from dbt.common.clients import agate_helper
 from dbt.exceptions import DbtConfigError
 from dbt.common.exceptions import DbtValidationError
+from dbt.context.providers import generate_runtime_macro_context
 from psycopg2 import extensions as psycopg2_extensions
 from psycopg2 import DatabaseError
 
@@ -430,14 +430,10 @@ def _mock_state_check(self):
         self.psycopg2.connect.return_value = self.handle
         self.adapter = PostgresAdapter(self.config, self.mp_context)
         self.adapter.set_macro_resolver(load_internal_manifest_macros(self.config))
-
-        query_header_context = generate_query_header_context(
-            self.adapter.config, self.adapter.get_macro_resolver()
+        self.adapter.set_macro_context_generator(generate_runtime_macro_context)
+        self.adapter.connections.set_query_header(
+            generate_query_header_context(self.config, self.adapter.get_macro_resolver())
         )
-        self.adapter.connections.query_header = MacroQueryStringSetter(
-            self.config, query_header_context
-        )
-
         self.qh_patch = mock.patch.object(self.adapter.connections.query_header, "add")
         self.mock_query_header_add = self.qh_patch.start()
         self.mock_query_header_add.side_effect = lambda q: "/* dbt */\n{}".format(q)