Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add macro_context_generator on adapter #9251

Merged
Show file tree
Hide file tree
Changes from 22 commits
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
2474722
moving types_pb2.py to common/events
colin-rogers-dbt Oct 25, 2023
07743b7
merge
colin-rogers-dbt Oct 25, 2023
252e3e3
Merge branch 'feature/decouple-adapters-from-core' of https://github.…
colin-rogers-dbt Oct 26, 2023
ff9d519
Merge branch 'feature/decouple-adapters-from-core' of https://github.…
colin-rogers-dbt Nov 2, 2023
bf079b1
Merge branch 'feature/decouple-adapters-from-core' of https://github.…
colin-rogers-dbt Nov 28, 2023
3187ded
Merge branch 'feature/decouple-adapters-from-core' of https://github.…
colin-rogers-dbt Nov 28, 2023
4da67bf
Merge branch 'feature/decouple-adapters-from-core' of https://github.…
colin-rogers-dbt Dec 5, 2023
c03823f
Merge branch 'feature/decouple-adapters-from-core' of https://github.…
colin-rogers-dbt Dec 6, 2023
2811ccd
Merge branch 'feature/decouple-adapters-from-core' of https://github.…
colin-rogers-dbt Dec 7, 2023
4dec100
remove manifest from adapter.execute_macro, replace with MacroResolve…
MichelleArk Dec 6, 2023
3e5515b
rename to MacroResolverProtocol
MichelleArk Dec 7, 2023
a70edd6
Merge branch 'feature/decouple-adapters-from-core' into macro-resolve…
MichelleArk Dec 7, 2023
7b1e332
pass MacroResolverProtcol in adapter.calculate_freshness_from_metadata
MichelleArk Dec 7, 2023
3620ce6
changelog entry
MichelleArk Dec 7, 2023
538481a
fix adapter.calculate_freshness call
MichelleArk Dec 7, 2023
f61da7b
add macro_context_generator on adapter
MichelleArk Dec 7, 2023
325717f
fix adapter test setup
MichelleArk Dec 7, 2023
bb12375
changelog entry
MichelleArk Dec 7, 2023
e88138a
Merge branch 'feature/decouple-adapters-from-core' of https://github.…
colin-rogers-dbt Dec 7, 2023
1740df5
Update parser to support conversion metrics (#9173)
WilliamDee Dec 7, 2023
977a842
Merge branch 'feature/decouple-adapters-from-core' of https://github.…
colin-rogers-dbt Dec 7, 2023
001bd51
merge feature/decouple-adapters-from-core
colin-rogers-dbt Dec 7, 2023
2bb2922
Add typing for macro_context_generator, fix query_header_context
colin-rogers-dbt Dec 8, 2023
003606a
merge main
colin-rogers-dbt Dec 8, 2023
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions .changes/unreleased/Under the Hood-20231208-004854.yaml
Original file line number Diff line number Diff line change
@@ -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"
36 changes: 22 additions & 14 deletions core/dbt/adapters/base/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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

Check warning on line 281 in core/dbt/adapters/base/impl.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/adapters/base/impl.py#L280-L281

Added lines #L280 - L281 were not covered by tests

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
###
Expand Down Expand Up @@ -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!")

Check warning on line 1074 in core/dbt/adapters/base/impl.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/adapters/base/impl.py#L1074

Added line #L1074 was not covered by tests

if self._macro_context_generator is None:
raise DbtInternalError("Macro context generator was None when calling execute_macro!")

Check warning on line 1077 in core/dbt/adapters/base/impl.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/adapters/base/impl.py#L1077

Added line #L1077 was not covered by tests

macro = resolver.find_macro_by_name(macro_name, self.config.project_name, project)
if macro is None:
Expand All @@ -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)
Expand Down
18 changes: 18 additions & 0 deletions core/dbt/adapters/protocol.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]:
...

Check warning on line 67 in core/dbt/adapters/protocol.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/adapters/protocol.py#L67

Added line #L67 was not covered by tests


# TODO CT-211
class AdapterProtocol( # type: ignore[misc]
Protocol,
Expand All @@ -78,14 +90,20 @@
...

def set_macro_resolver(self, macro_resolver: MacroResolverProtocol) -> None:
...

Check warning on line 93 in core/dbt/adapters/protocol.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/adapters/protocol.py#L93

Added line #L93 was not covered by tests

def get_macro_resolver(self) -> Optional[MacroResolverProtocol]:
...

Check warning on line 96 in core/dbt/adapters/protocol.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/adapters/protocol.py#L96

Added line #L96 was not covered by tests

def clear_macro_resolver(self) -> None:
...

Check warning on line 99 in core/dbt/adapters/protocol.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/adapters/protocol.py#L99

Added line #L99 was not covered by tests

def set_macro_context_generator(
self,
macro_context_generator: MacroContextGeneratorCallable,
) -> None:
...

Check warning on line 105 in core/dbt/adapters/protocol.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/adapters/protocol.py#L105

Added line #L105 was not covered by tests

@classmethod
def type(cls) -> str:
pass
Expand Down
5 changes: 4 additions & 1 deletion core/dbt/cli/requires.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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 (
Expand Down Expand Up @@ -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:
Expand Down
12 changes: 7 additions & 5 deletions core/dbt/parser/manifest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Copy link
Contributor

@peterallenwebb peterallenwebb Dec 8, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Love this change to the unused lambda parameter name.

else:
self.macro_hook = macro_hook

Expand Down Expand Up @@ -287,7 +287,7 @@
# the config and adapter may be persistent.
if reset:
config.clear_dependencies()
adapter.clear_macro_resolver()

Check warning on line 290 in core/dbt/parser/manifest.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/parser/manifest.py#L290

Added line #L290 was not covered by tests
macro_hook = adapter.connections.set_query_header

flags = get_flags()
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down
2 changes: 2 additions & 0 deletions core/dbt/tests/fixtures/project.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down
Original file line number Diff line number Diff line change
@@ -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}}
2 changes: 1 addition & 1 deletion tests/functional/artifacts/data/state/v12/manifest.json

Large diffs are not rendered by default.

14 changes: 5 additions & 9 deletions tests/unit/test_postgres_adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
Expand Down
Loading