include project macros in the manifest the adapter stores locally

dbt-labs · Aug 5, 2020 · 285479c · 285479c
1 parent e641ec1
commit 285479c
Show file tree

Hide file tree

Showing 11 changed files with 103 additions and 78 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -3,6 +3,7 @@
 
 ### Features
 - Add support for impersonating a service account using `impersonate_service_account` in the BigQuery profile configuration ([#2677](https://github.com/fishtown-analytics/dbt/issues/2677)) ([docs](https://docs.getdbt.com/reference/warehouse-profiles/bigquery-profile#service-account-impersonation))
+- Macros in the current project can override internal dbt macros that are called through `execute_macros`. ([#2301](https://github.com/fishtown-analytics/dbt/issues/2301), [#2686](https://github.com/fishtown-analytics/dbt/pull/2686))
 
 
 ### Breaking changes

diff --git a/core/dbt/adapters/base/impl.py b/core/dbt/adapters/base/impl.py
@@ -158,7 +158,7 @@ def __init__(self, config):
         self.config = config
         self.cache = RelationsCache()
         self.connections = self.ConnectionManager(config)
-        self._internal_manifest_lazy: Optional[Manifest] = None
+        self._macro_manifest_lazy: Optional[Manifest] = None
 
     ###
     # Methods that pass through to the connection manager
@@ -239,24 +239,26 @@ def type(cls) -> str:
         return cls.ConnectionManager.TYPE
 
     @property
-    def _internal_manifest(self) -> Manifest:
-        if self._internal_manifest_lazy is None:
-            return self.load_internal_manifest()
-        return self._internal_manifest_lazy
+    def _macro_manifest(self) -> Manifest:
+        if self._macro_manifest_lazy is None:
+            return self.load_macro_manifest()
+        return self._macro_manifest_lazy
 
-    def check_internal_manifest(self) -> Optional[Manifest]:
+    def check_macro_manifest(self) -> Optional[Manifest]:
         """Return the internal manifest (used for executing macros) if it's
         been initialized, otherwise return None.
         """
-        return self._internal_manifest_lazy
+        return self._macro_manifest_lazy
 
-    def load_internal_manifest(self) -> Manifest:
-        if self._internal_manifest_lazy is None:
+    def load_macro_manifest(self) -> Manifest:
+        if self._macro_manifest_lazy is None:
             # avoid a circular import
-            from dbt.parser.manifest import load_internal_manifest
-            manifest = load_internal_manifest(self.config)
-            self._internal_manifest_lazy = manifest
-        return self._internal_manifest_lazy
+            from dbt.parser.manifest import load_macro_manifest
+            manifest = load_macro_manifest(
+                self.config, self.connections.set_query_header
+            )
+            self._macro_manifest_lazy = manifest
+        return self._macro_manifest_lazy
 
     ###
     # Caching methods
@@ -941,7 +943,7 @@ def execute_macro(
             context_override = {}
 
         if manifest is None:
-            manifest = self._internal_manifest
+            manifest = self._macro_manifest
 
         macro = manifest.find_macro_by_name(
             macro_name, self.config.project_name, project

diff --git a/core/dbt/parser/manifest.py b/core/dbt/parser/manifest.py
@@ -11,8 +11,6 @@
 from dbt import deprecations
 from dbt.adapters.factory import (
     get_relation_class_by_name,
-    get_adapter_package_names,
-    get_include_paths,
 )
 from dbt.helper_types import PathSet
 from dbt.logger import GLOBAL_LOGGER as logger, DbtProcessState
@@ -123,28 +121,6 @@ def __init__(
         )
         self._loaded_file_cache: Dict[str, FileBlock] = {}
 
-    def _load_macros(
-        self,
-        old_results: Optional[ParseResult],
-        internal_manifest: Optional[Manifest] = None,
-    ) -> None:
-        projects = self.all_projects
-        if internal_manifest is not None:
-            # skip internal packages
-            packages = get_adapter_package_names(
-                self.root_project.credentials.type
-            )
-            projects = {
-                k: v for k, v in self.all_projects.items() if k not in packages
-            }
-            self.results.macros.update(internal_manifest.macros)
-            self.results.files.update(internal_manifest.files)
-
-        for project in projects.values():
-            parser = MacroParser(self.results, project)
-            for path in parser.search():
-                self.parse_with_cache(path, parser, old_results)
-
     def parse_with_cache(
         self,
         path: FilePath,
@@ -201,25 +177,26 @@ def parse_project(
 
     def load_only_macros(self) -> Manifest:
         old_results = self.read_parse_results()
-        self._load_macros(old_results, internal_manifest=None)
+
+        for project in self.all_projects.values():
+            parser = MacroParser(self.results, project)
+            for path in parser.search():
+                self.parse_with_cache(path, parser, old_results)
+
         # make a manifest with just the macros to get the context
         macro_manifest = Manifest.from_macros(
             macros=self.results.macros,
             files=self.results.files
         )
+        self.macro_hook(macro_manifest)
         return macro_manifest
 
-    def load(self, internal_manifest: Optional[Manifest] = None):
+    def load(self, macro_manifest: Manifest):
         old_results = self.read_parse_results()
         if old_results is not None:
             logger.debug('Got an acceptable cached parse result')
-        self._load_macros(old_results, internal_manifest=internal_manifest)
-        # make a manifest with just the macros to get the context
-        macro_manifest = Manifest.from_macros(
-            macros=self.results.macros,
-            files=self.results.files
-        )
-        self.macro_hook(macro_manifest)
+        self.results.macros.update(macro_manifest.macros)
+        self.results.files.update(macro_manifest.files)
 
         for project in self.all_projects.values():
             # parse a single project
@@ -352,7 +329,7 @@ def create_manifest(self) -> Manifest:
     def load_all(
         cls,
         root_config: RuntimeConfig,
-        internal_manifest: Optional[Manifest],
+        macro_manifest: Manifest,
         macro_hook: Callable[[Manifest], Any],
     ) -> Manifest:
         with PARSING_STATE:
@@ -367,18 +344,22 @@ def load_all(
                     project_names=''.join(v1_configs)
                 )
             loader = cls(root_config, projects, macro_hook)
-            loader.load(internal_manifest=internal_manifest)
+            loader.load(macro_manifest=macro_manifest)
             loader.write_parse_results()
             manifest = loader.create_manifest()
             _check_manifest(manifest, root_config)
             manifest.build_flat_graph()
             return manifest
 
     @classmethod
-    def load_internal(cls, root_config: RuntimeConfig) -> Manifest:
+    def load_macros(
+        cls,
+        root_config: RuntimeConfig,
+        macro_hook: Callable[[Manifest], Any],
+    ) -> Manifest:
         with PARSING_STATE:
-            projects = load_internal_projects(root_config)
-            loader = cls(root_config, projects)
+            projects = root_config.load_dependencies()
+            loader = cls(root_config, projects, macro_hook)
             return loader.load_only_macros()
 
 
@@ -681,18 +662,16 @@ def process_node(
     _process_docs_for_node(ctx, node)
 
 
-def load_internal_projects(config):
-    project_paths = get_include_paths(config.credentials.type)
-    return dict(_load_projects(config, project_paths))
-
-
-def load_internal_manifest(config: RuntimeConfig) -> Manifest:
-    return ManifestLoader.load_internal(config)
+def load_macro_manifest(
+    config: RuntimeConfig,
+    macro_hook: Callable[[Manifest], Any],
+) -> Manifest:
+    return ManifestLoader.load_macros(config, macro_hook)
 
 
 def load_manifest(
     config: RuntimeConfig,
-    internal_manifest: Optional[Manifest],
+    macro_manifest: Manifest,
     macro_hook: Callable[[Manifest], Any],
 ) -> Manifest:
-    return ManifestLoader.load_all(config, internal_manifest, macro_hook)
+    return ManifestLoader.load_all(config, macro_manifest, macro_hook)
diff --git a/core/dbt/perf_utils.py b/core/dbt/perf_utils.py
@@ -15,9 +15,10 @@ def get_full_manifest(config: RuntimeConfig) -> Manifest:
     attached to the adapter for any methods that need it.
     """
     adapter = get_adapter(config)  # type: ignore
-    internal: Manifest = adapter.load_internal_manifest()
+    internal: Manifest = adapter.load_macro_manifest()
 
-    def set_header(manifest: Manifest) -> None:
-        adapter.connections.set_query_header(manifest)
-
-    return load_manifest(config, internal, set_header)
+    return load_manifest(
+        config,
+        internal,
+        adapter.connections.set_query_header,
+    )
diff --git a/test/integration/016_macro_tests/override-get-columns-macros/macros.sql b/test/integration/016_macro_tests/override-get-columns-macros/macros.sql
@@ -0,0 +1,3 @@
+{% macro get_columns_in_relation(relation) %}
+	{{ return('a string') }}
+{% endmacro %}
diff --git a/test/integration/016_macro_tests/override-get-columns-models/model.sql b/test/integration/016_macro_tests/override-get-columns-models/model.sql
@@ -0,0 +1,5 @@
+{% set result = adapter.get_columns_in_relation(this) %}
+{% if execute and result != 'a string' %}
+  {% do exceptions.raise_compiler_error('overriding get_columns_in_relation failed') %}
+{% endif %}
+select 1 as id
diff --git a/test/integration/016_macro_tests/test_macros.py b/test/integration/016_macro_tests/test_macros.py
@@ -93,3 +93,27 @@ def test_postgres_invalid_macro(self):
             self.run_dbt(['run'])
 
         assert "In dispatch: No macro named 'dispatch_to_nowhere' found" in str(exc.value)
+
+
+class TestMacroOverrideBuiltin(DBTIntegrationTest):
+    @property
+    def schema(self):
+        return "test_macros_016"
+
+    @property
+    def models(self):
+        return 'override-get-columns-models'
+
+    @property
+    def project_config(self):
+        return {
+            'config-version': 2,
+            'macro-paths': ['override-get-columns-macros'],
+        }
+
+
+    @use_profile('postgres')
+    def test_postgres_overrides(self):
+        # the first time, the model doesn't exist
+        self.run_dbt()
+        self.run_dbt()
diff --git a/test/unit/test_graph.py b/test/unit/test_graph.py
@@ -104,7 +104,7 @@ def _mock_parse_result(config, all_projects):
         self.mock_source_file = self.load_source_file_patcher.start()
         self.mock_source_file.side_effect = lambda path: [n for n in self.mock_models if n.path == path][0]
 
-        self.internal_manifest = Manifest.from_macros(macros={
+        self.macro_manifest = Manifest.from_macros(macros={
             n.unique_id: n for n in generate_name_macros('test_models_compile')
         })
 
@@ -161,7 +161,7 @@ def use_models(self, models):
 
     def load_manifest(self, config):
         loader = dbt.parser.manifest.ManifestLoader(config, {config.project_name: config})
-        loader.load(internal_manifest=self.internal_manifest)
+        loader.load(macro_manifest=self.macro_manifest)
         return loader.create_manifest()
 
     def test__single_model(self):
@@ -319,7 +319,7 @@ def test__partial_parse(self):
         config = self.get_config()
 
         loader = dbt.parser.manifest.ManifestLoader(config, {config.project_name: config})
-        loader.load(internal_manifest=self.internal_manifest)
+        loader.load(macro_manifest=self.macro_manifest)
         loader.create_manifest()
         results = loader.results
 

diff --git a/test/unit/test_postgres_adapter.py b/test/unit/test_postgres_adapter.py
@@ -16,7 +16,7 @@
 from psycopg2 import extensions as psycopg2_extensions
 from psycopg2 import DatabaseError
 
-from .utils import config_from_parts_or_dicts, inject_adapter, mock_connection, TestAdapterConversions
+from .utils import config_from_parts_or_dicts, inject_adapter, mock_connection, TestAdapterConversions, load_internal_manifest_macros
 
 
 class TestPostgresAdapter(unittest.TestCase):
@@ -276,20 +276,21 @@ def setUp(self):
         self.patcher = mock.patch('dbt.adapters.postgres.connections.psycopg2')
         self.psycopg2 = self.patcher.start()
 
+        self.load_patch = mock.patch('dbt.parser.manifest.make_parse_result')
+        self.mock_parse_result = self.load_patch.start()
+        self.mock_parse_result.return_value = ParseResult.rpc()
+
         self.psycopg2.connect.return_value = self.handle
         self.adapter = PostgresAdapter(self.config)
-        self.adapter.connections.query_header = MacroQueryStringSetter(self.config, mock.MagicMock(macros={}))
+        self.adapter._macro_manifest_lazy = load_internal_manifest_macros(self.config)
+        self.adapter.connections.query_header = MacroQueryStringSetter(self.config, self.adapter._macro_manifest_lazy)
 
         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)
         self.adapter.acquire_connection()
         inject_adapter(self.adapter, PostgresPlugin)
 
-        self.load_patch = mock.patch('dbt.parser.manifest.make_parse_result')
-        self.mock_parse_result = self.load_patch.start()
-        self.mock_parse_result.return_value = ParseResult.rpc()
-
     def tearDown(self):
         # we want a unique self.handle every time.
         self.adapter.cleanup_connections()

diff --git a/test/unit/test_snowflake_adapter.py b/test/unit/test_snowflake_adapter.py
@@ -14,7 +14,7 @@
 from dbt.parser.results import ParseResult
 from snowflake import connector as snowflake_connector
 
-from .utils import config_from_parts_or_dicts, inject_adapter, mock_connection, TestAdapterConversions
+from .utils import config_from_parts_or_dicts, inject_adapter, mock_connection, TestAdapterConversions, load_internal_manifest_macros
 
 
 class TestSnowflakeAdapter(unittest.TestCase):
@@ -65,8 +65,8 @@ def setUp(self):
 
         self.snowflake.return_value = self.handle
         self.adapter = SnowflakeAdapter(self.config)
-
-        self.adapter.connections.query_header = MacroQueryStringSetter(self.config, mock.MagicMock(macros={}))
+        self.adapter._macro_manifest_lazy = load_internal_manifest_macros(self.config)
+        self.adapter.connections.query_header = MacroQueryStringSetter(self.config, self.adapter._macro_manifest_lazy)
 
         self.qh_patch = mock.patch.object(self.adapter.connections.query_header, 'add')
         self.mock_query_header_add = self.qh_patch.start()

diff --git a/test/unit/utils.py b/test/unit/utils.py
@@ -289,3 +289,12 @@ def MockDocumentation(package, name, **kwargs):
     )
     doc.name = name
     return doc
+
+
+def load_internal_manifest_macros(config, macro_hook = lambda m: None):
+    from dbt.adapters.factory import get_include_paths
+    from dbt.parser.manifest import ManifestLoader
+    paths = get_include_paths(config.credentials.type)
+    projects = {k: v for k, v in config.load_dependencies().items() if k.startswith('dbt')}
+    loader = ManifestLoader(config, projects, macro_hook)
+    return loader.load_only_macros()