Repository: dbt-labs/dbt-core

Issue:
[CT-3050] [Applied State] Implement Core Changes for Filtered Catalog Queries
### Short description

Building on some of the research done for #4997, we want to improve catalog creation performance by selecting only the needed relation information, instead of information for every relation in a schema.

### Acceptance criteria

- In dbt's base adapter implementation (base/impl.py), the `_get_one_catalog` method has been modified to accept an additional new parameter, `relations_by_schema: Optional[Dict[str, Optional[List[BaseRelation]]]] = None`, mapping schema names to the list of relations in that schema for which catalog information should be retrieved.
- When the new parameter defaults to None, behavior should be the same as before.
- When `_get_one_catalog` invokes the `get_catalog` macro, it should include the new parameter.
- Adapters which do not yet implement the new get_catalog macro parameter should continue to work as before. Therefore we need to evaluate the macro without generating an error and we also need to continue post-filtering the results.
- When the value of the `relations` parameter is `None`, the `_get_one_catalog` method, and `get_catalog` macros should continue to retrieve all relation information for every schema listed in the `schema` parameter, matching historical behavior.

### Impact to Other Teams

The adapter teams will need to implement changes to the get_catalog macro in each adapter implementation, in order to take advantage of the new parameter:
- https://github.com/dbt-labs/dbt-bigquery/issues/950
- https://github.com/dbt-labs/dbt-redshift/issues/625
- https://github.com/dbt-labs/dbt-snowflake/issues/801
- https://github.com/dbt-labs/dbt-spark/issues/900


Discussion:

Resolving patch:
diff --git a/core/dbt/adapters/base/impl.py b/core/dbt/adapters/base/impl.py
--- a/core/dbt/adapters/base/impl.py
+++ b/core/dbt/adapters/base/impl.py
@@ -74,6 +74,7 @@
 from dbt import deprecations
 
 GET_CATALOG_MACRO_NAME = "get_catalog"
+GET_CATALOG_RELATIONS_MACRO_NAME = "get_catalog_relations"
 FRESHNESS_MACRO_NAME = "collect_freshness"
 
 
@@ -161,6 +162,14 @@ def submit(self, compiled_code: str) -> Any:
         raise NotImplementedError("PythonJobHelper submit function is not implemented yet")
 
 
+class AdapterFeature(str, Enum):
+    """Enumeration of optional adapter features which can be probed using BaseAdapter.has_feature()"""
+
+    CatalogByRelations = "CatalogByRelations"
+    """Flags support for retrieving catalog information using a list of relations, rather than always retrieving all
+    the relations in a schema """
+
+
 class BaseAdapter(metaclass=AdapterMeta):
     """The BaseAdapter provides an abstract base class for adapters.
 
@@ -415,6 +424,29 @@ def _get_catalog_schemas(self, manifest: Manifest) -> SchemaSearchMap:
         lowercase strings.
         """
         info_schema_name_map = SchemaSearchMap()
+        relations = self._get_catalog_relations(manifest)
+        for relation in relations:
+            info_schema_name_map.add(relation)
+        # result is a map whose keys are information_schema Relations without
+        # identifiers that have appropriate database prefixes, and whose values
+        # are sets of lowercase schema names that are valid members of those
+        # databases
+        return info_schema_name_map
+
+    def _get_catalog_relations_by_info_schema(
+        self, manifest: Manifest
+    ) -> Dict[InformationSchema, List[BaseRelation]]:
+        relations = self._get_catalog_relations(manifest)
+        relations_by_info_schema: Dict[InformationSchema, List[BaseRelation]] = dict()
+        for relation in relations:
+            info_schema = relation.information_schema_only()
+            if info_schema not in relations_by_info_schema:
+                relations_by_info_schema[info_schema] = []
+            relations_by_info_schema[info_schema].append(relation)
+
+        return relations_by_info_schema
+
+    def _get_catalog_relations(self, manifest: Manifest) -> List[BaseRelation]:
         nodes: Iterator[ResultNode] = chain(
             [
                 node
@@ -423,14 +455,9 @@ def _get_catalog_schemas(self, manifest: Manifest) -> SchemaSearchMap:
             ],
             manifest.sources.values(),
         )
-        for node in nodes:
-            relation = self.Relation.create_from(self.config, node)
-            info_schema_name_map.add(relation)
-        # result is a map whose keys are information_schema Relations without
-        # identifiers that have appropriate database prefixes, and whose values
-        # are sets of lowercase schema names that are valid members of those
-        # databases
-        return info_schema_name_map
+
+        relations = [self.Relation.create_from(self.config, n) for n in nodes]
+        return relations
 
     def _relations_cache_for_schemas(
         self, manifest: Manifest, cache_schemas: Optional[Set[BaseRelation]] = None
@@ -1093,20 +1120,57 @@ def _get_one_catalog(
         results = self._catalog_filter_table(table, manifest)  # type: ignore[arg-type]
         return results
 
+    def _get_one_catalog_by_relations(
+        self,
+        information_schema: InformationSchema,
+        relations: List[BaseRelation],
+        manifest: Manifest,
+    ) -> agate.Table:
+
+        kwargs = {
+            "information_schema": information_schema,
+            "relations": relations,
+        }
+        table = self.execute_macro(
+            GET_CATALOG_RELATIONS_MACRO_NAME,
+            kwargs=kwargs,
+            # pass in the full manifest, so we get any local project
+            # overrides
+            manifest=manifest,
+        )
+
+        results = self._catalog_filter_table(table, manifest)  # type: ignore[arg-type]
+        return results
+
     def get_catalog(self, manifest: Manifest) -> Tuple[agate.Table, List[Exception]]:
-        schema_map = self._get_catalog_schemas(manifest)
 
         with executor(self.config) as tpe:
             futures: List[Future[agate.Table]] = []
-            for info, schemas in schema_map.items():
-                if len(schemas) == 0:
-                    continue
-                name = ".".join([str(info.database), "information_schema"])
-
-                fut = tpe.submit_connected(
-                    self, name, self._get_one_catalog, info, schemas, manifest
-                )
-                futures.append(fut)
+            relation_count = len(self._get_catalog_relations(manifest))
+            if relation_count <= 100 and self.has_feature(AdapterFeature.CatalogByRelations):
+                relations_by_schema = self._get_catalog_relations_by_info_schema(manifest)
+                for info_schema in relations_by_schema:
+                    name = ".".join([str(info_schema.database), "information_schema"])
+                    relations = relations_by_schema[info_schema]
+                    fut = tpe.submit_connected(
+                        self,
+                        name,
+                        self._get_one_catalog_by_relations,
+                        info_schema,
+                        relations,
+                        manifest,
+                    )
+                    futures.append(fut)
+            else:
+                schema_map: SchemaSearchMap = self._get_catalog_schemas(manifest)
+                for info, schemas in schema_map.items():
+                    if len(schemas) == 0:
+                        continue
+                    name = ".".join([str(info.database), "information_schema"])
+                    fut = tpe.submit_connected(
+                        self, name, self._get_one_catalog, info, schemas, manifest
+                    )
+                    futures.append(fut)
 
             catalogs, exceptions = catch_as_completed(futures)
 
@@ -1437,6 +1501,13 @@ def render_model_constraint(cls, constraint: ModelLevelConstraint) -> Optional[s
         else:
             return None
 
+    @classmethod
+    def has_feature(cls, feature: AdapterFeature) -> bool:
+        # The base adapter implementation does not implement any optional
+        # features, so always return false. Adapters which wish to provide
+        # optional features will have to override this function.
+        return False
+
 
 COLUMNS_EQUAL_SQL = """
 with diff_count as (
diff --git a/core/dbt/adapters/base/relation.py b/core/dbt/adapters/base/relation.py
--- a/core/dbt/adapters/base/relation.py
+++ b/core/dbt/adapters/base/relation.py
@@ -459,11 +459,11 @@ def add(self, relation: BaseRelation):
         self[key].add(schema)
 
     def search(self) -> Iterator[Tuple[InformationSchema, Optional[str]]]:
-        for information_schema_name, schemas in self.items():
+        for information_schema, schemas in self.items():
             for schema in schemas:
-                yield information_schema_name, schema
+                yield information_schema, schema
 
-    def flatten(self, allow_multiple_databases: bool = False):
+    def flatten(self, allow_multiple_databases: bool = False) -> "SchemaSearchMap":
         new = self.__class__()
 
         # make sure we don't have multiple databases if allow_multiple_databases is set to False
diff --git a/plugins/postgres/dbt/adapters/postgres/impl.py b/plugins/postgres/dbt/adapters/postgres/impl.py
--- a/plugins/postgres/dbt/adapters/postgres/impl.py
+++ b/plugins/postgres/dbt/adapters/postgres/impl.py
@@ -3,7 +3,7 @@
 from typing import Optional, Set, List, Any
 
 from dbt.adapters.base.meta import available
-from dbt.adapters.base.impl import AdapterConfig, ConstraintSupport
+from dbt.adapters.base.impl import AdapterConfig, AdapterFeature, ConstraintSupport
 from dbt.adapters.sql import SQLAdapter
 from dbt.adapters.postgres import PostgresConnectionManager
 from dbt.adapters.postgres.column import PostgresColumn
@@ -73,6 +73,10 @@ class PostgresAdapter(SQLAdapter):
         ConstraintType.foreign_key: ConstraintSupport.ENFORCED,
     }
 
+    CATALOG_BY_RELATION_SUPPORT = True
+
+    SUPPORTED_FEATURES: Set[AdapterFeature] = frozenset([AdapterFeature.CatalogByRelations])
+
     @classmethod
     def date_function(cls):
         return "now()"
@@ -113,9 +117,9 @@ def _link_cached_database_relations(self, schemas: Set[str]):
 
     def _get_catalog_schemas(self, manifest):
         # postgres only allow one database (the main one)
-        schemas = super()._get_catalog_schemas(manifest)
+        schema_search_map = super()._get_catalog_schemas(manifest)
         try:
-            return schemas.flatten()
+            return schema_search_map.flatten()
         except DbtRuntimeError as exc:
             raise CrossDbReferenceProhibitedError(self.type(), exc.msg)
 
@@ -143,3 +147,7 @@ def valid_incremental_strategies(self):
 
     def debug_query(self):
         self.execute("select 1 as id")
+
+    @classmethod
+    def has_feature(cls, feature: AdapterFeature) -> bool:
+        return feature in cls.SUPPORTED_FEATURES


Test patch:
diff --git a/tests/functional/artifacts/test_override.py b/tests/functional/artifacts/test_override.py
--- a/tests/functional/artifacts/test_override.py
+++ b/tests/functional/artifacts/test_override.py
@@ -7,7 +7,7 @@
 """
 
 fail_macros__failure_sql = """
-{% macro get_catalog(information_schema, schemas) %}
+{% macro get_catalog_relations(information_schema, relations) %}
     {% do exceptions.raise_compiler_error('rejected: no catalogs for you') %}
 {% endmacro %}
 
diff --git a/tests/unit/test_postgres_adapter.py b/tests/unit/test_postgres_adapter.py
--- a/tests/unit/test_postgres_adapter.py
+++ b/tests/unit/test_postgres_adapter.py
@@ -322,8 +322,8 @@ def test_set_zero_keepalive(self, psycopg2):
         )
 
     @mock.patch.object(PostgresAdapter, "execute_macro")
-    @mock.patch.object(PostgresAdapter, "_get_catalog_schemas")
-    def test_get_catalog_various_schemas(self, mock_get_schemas, mock_execute):
+    @mock.patch.object(PostgresAdapter, "_get_catalog_relations_by_info_schema")
+    def test_get_catalog_various_schemas(self, mock_get_relations, mock_execute):
         column_names = ["table_database", "table_schema", "table_name"]
         rows = [
             ("dbt", "foo", "bar"),
@@ -334,9 +334,13 @@ def test_get_catalog_various_schemas(self, mock_get_schemas, mock_execute):
         ]
         mock_execute.return_value = agate.Table(rows=rows, column_names=column_names)
 
-        mock_get_schemas.return_value.items.return_value = [
-            (mock.MagicMock(database="dbt"), {"foo", "FOO", "quux"})
-        ]
+        mock_get_relations.return_value = {
+            mock.MagicMock(database="dbt"): [
+                mock.MagicMock(schema="foo"),
+                mock.MagicMock(schema="FOO"),
+                mock.MagicMock(schema="quux"),
+            ]
+        }
 
         mock_manifest = mock.MagicMock()
         mock_manifest.get_used_schemas.return_value = {("dbt", "foo"), ("dbt", "quux")}