Repository: dbt-labs/dbt-core

Issue:
[CT-3216] [Regression] Calls to `get_catalog` are not all wrapped in the feature flag
### Is this a regression in a recent version of dbt-core?

- [X] I believe this is a regression in dbt-core functionality
- [X] I have searched the existing issues, and I could not find an existing issue for this regression

### Current Behavior

Calls to the new `get_catalog` method fail due to providing too many arguments (the instance plus two arguments). This line causes the error:

https://github.com/dbt-labs/dbt-core/blob/140b655f0d7663468066cfcb2c77e9b8ca9a3d38/core/dbt/task/generate.py#L248

### Expected/Previous Behavior

Calls to `get_catalog` work so that `docs generate` runs properly.

### Steps To Reproduce

Try to run `docs generate` to trigger a call to `get_catalog`.

### Relevant log output

```shell
https://github.com/databricks/dbt-databricks/actions/runs/6499031208/job/17651537811
```


### Environment

```markdown
- OS:
- Python:
- dbt (working version): dbt-core/dbt-databricks 1.7.0rc1
- dbt (regression version): dbt-core/dbt-databricks 1.6
```


### Which database adapter are you using with dbt?

other (mention it in "Additional Context")

### Additional Context

dbt-databricks 1.7.0rc1



Discussion:
@mikealfare do you think the root cause for the particular error you observed is that dbt-databricks overrides `get_catalog`  [here](https://github.com/databricks/dbt-databricks/blob/29b4c20d79b633d2d31fc585d57acff214745164/dbt/adapters/databricks/impl.py#L420)?

Either way, do we have a proposed fix yet?
> @mikealfare do you think the root cause for the particular error you observed is that dbt-databricks overrides `get_catalog` [here](https://github.com/databricks/dbt-databricks/blob/29b4c20d79b633d2d31fc585d57acff214745164/dbt/adapters/databricks/impl.py#L420)?
> 
> Either way, do we have a proposed fix yet?

Yes, the fact that it's overridden is a factor. We added an optional argument to `get_catalog` and then called `get_catalog` somewhere else using the new argument. If an adapter does not override `get_catalog`, it will work fine. But if they do, then we will provide too many arguments to `get_catalog`. On the surface, the solution appears to be wrapping the call to `get_catalog` in the new feature flag framework (Capability) since the reason we added the argument is also covered by that framework. We could also try wrapping it in a simple try/catch with the exception path omitting the extra argument. These would have the same effect. That being said, we still need to discuss it to make sure we're making the right decision.

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
@@ -455,30 +455,16 @@ def _get_catalog_relations_by_info_schema(
 
         return relations_by_info_schema
 
-    def _get_catalog_relations(
-        self, manifest: Manifest, selected_nodes: Optional[Set] = None
-    ) -> List[BaseRelation]:
-        nodes: Iterator[ResultNode]
-        if selected_nodes:
-            selected: List[ResultNode] = []
-            for unique_id in selected_nodes:
-                if unique_id in manifest.nodes:
-                    node = manifest.nodes[unique_id]
-                    if node.is_relational and not node.is_ephemeral_model:
-                        selected.append(node)
-                elif unique_id in manifest.sources:
-                    source = manifest.sources[unique_id]
-                    selected.append(source)
-            nodes = iter(selected)
-        else:
-            nodes = chain(
-                [
-                    node
-                    for node in manifest.nodes.values()
-                    if (node.is_relational and not node.is_ephemeral_model)
-                ],
-                manifest.sources.values(),
-            )
+    def _get_catalog_relations(self, manifest: Manifest) -> List[BaseRelation]:
+
+        nodes = chain(
+            [
+                node
+                for node in manifest.nodes.values()
+                if (node.is_relational and not node.is_ephemeral_model)
+            ],
+            manifest.sources.values(),
+        )
 
         relations = [self.Relation.create_from(self.config, n) for n in nodes]
         return relations
@@ -1166,43 +1152,83 @@ def _get_one_catalog_by_relations(
         results = self._catalog_filter_table(table, manifest)  # type: ignore[arg-type]
         return results
 
-    def get_catalog(
-        self, manifest: Manifest, selected_nodes: Optional[Set] = None
-    ) -> Tuple[agate.Table, List[Exception]]:
+    def get_filtered_catalog(
+        self, manifest: Manifest, relations: Optional[Set[BaseRelation]] = None
+    ):
+        catalogs: agate.Table
+        if (
+            relations is None
+            or len(relations) > 100
+            or not self.supports(Capability.SchemaMetadataByRelations)
+        ):
+            # Do it the traditional way. We get the full catalog.
+            catalogs, exceptions = self.get_catalog(manifest)
+        else:
+            # Do it the new way. We try to save time by selecting information
+            # only for the exact set of relations we are interested in.
+            catalogs, exceptions = self.get_catalog_by_relations(manifest, relations)
+
+        if relations and catalogs:
+            relation_map = {
+                (
+                    r.database.casefold() if r.database else None,
+                    r.schema.casefold() if r.schema else None,
+                    r.identifier.casefold() if r.identifier else None,
+                )
+                for r in relations
+            }
+
+            def in_map(row: agate.Row):
+                d = _expect_row_value("table_database", row)
+                s = _expect_row_value("table_schema", row)
+                i = _expect_row_value("table_name", row)
+                return (d, s, i) in relation_map
+
+            catalogs = catalogs.where(in_map)
+
+        return catalogs, exceptions
+
+    def row_matches_relation(self, row: agate.Row, relations: Set[BaseRelation]):
+        pass
 
+    def get_catalog(self, manifest: Manifest) -> Tuple[agate.Table, List[Exception]]:
         with executor(self.config) as tpe:
             futures: List[Future[agate.Table]] = []
-            catalog_relations = self._get_catalog_relations(manifest, selected_nodes)
-            relation_count = len(catalog_relations)
-            if relation_count <= 100 and self.supports(Capability.SchemaMetadataByRelations):
-                relations_by_schema = self._get_catalog_relations_by_info_schema(catalog_relations)
-                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)
+            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)
         return catalogs, exceptions
 
+    def get_catalog_by_relations(
+        self, manifest: Manifest, relations: Set[BaseRelation]
+    ) -> Tuple[agate.Table, List[Exception]]:
+        with executor(self.config) as tpe:
+            futures: List[Future[agate.Table]] = []
+            relations_by_schema = self._get_catalog_relations_by_info_schema(relations)
+            for info_schema in relations_by_schema:
+                name = ".".join([str(info_schema.database), "information_schema"])
+                relations = set(relations_by_schema[info_schema])
+                fut = tpe.submit_connected(
+                    self,
+                    name,
+                    self._get_one_catalog_by_relations,
+                    info_schema,
+                    relations,
+                    manifest,
+                )
+                futures.append(fut)
+
+            catalogs, exceptions = catch_as_completed(futures)
+            return catalogs, exceptions
+
     def cancel_open_connections(self):
         """Cancel all open connections."""
         return self.connections.cancel_open()
diff --git a/core/dbt/task/generate.py b/core/dbt/task/generate.py
--- a/core/dbt/task/generate.py
+++ b/core/dbt/task/generate.py
@@ -1,7 +1,7 @@
 import os
 import shutil
 from datetime import datetime
-from typing import Dict, List, Any, Optional, Tuple, Set
+from typing import Dict, List, Any, Optional, Tuple, Set, Iterable
 import agate
 
 from dbt.dataclass_schema import ValidationError
@@ -223,11 +223,6 @@ def run(self) -> CatalogArtifact:
             DOCS_INDEX_FILE_PATH, os.path.join(self.config.project_target_path, "index.html")
         )
 
-        # Get the list of nodes that have been selected
-        selected_nodes = None
-        if self.job_queue is not None:
-            selected_nodes = self.job_queue.get_selected_nodes()
-
         for asset_path in self.config.asset_paths:
             to_asset_path = os.path.join(self.config.project_target_path, asset_path)
 
@@ -247,8 +242,18 @@ def run(self) -> CatalogArtifact:
             adapter = get_adapter(self.config)
             with adapter.connection_named("generate_catalog"):
                 fire_event(BuildingCatalog())
+                # Get a list of relations we need from the catalog
+                relations = None
+                if self.job_queue is not None:
+                    selected_node_ids = self.job_queue.get_selected_nodes()
+                    selected_nodes = self._get_nodes_from_ids(self.manifest, selected_node_ids)
+                    relations = {
+                        adapter.Relation.create_from(adapter.config, node_id)
+                        for node_id in selected_nodes
+                    }
+
                 # This generates the catalog as an agate.Table
-                catalog_table, exceptions = adapter.get_catalog(self.manifest, selected_nodes)
+                catalog_table, exceptions = adapter.get_filtered_catalog(self.manifest, relations)
 
         catalog_data: List[PrimitiveDict] = [
             dict(zip(catalog_table.column_names, map(dbt.utils._coerce_decimal, row)))
@@ -298,6 +303,19 @@ def run(self) -> CatalogArtifact:
         fire_event(CatalogWritten(path=os.path.abspath(catalog_path)))
         return results
 
+    @staticmethod
+    def _get_nodes_from_ids(manifest: Manifest, node_ids: Iterable[str]) -> List[ResultNode]:
+        selected: List[ResultNode] = []
+        for unique_id in node_ids:
+            if unique_id in manifest.nodes:
+                node = manifest.nodes[unique_id]
+                if node.is_relational and not node.is_ephemeral_model:
+                    selected.append(node)
+            elif unique_id in manifest.sources:
+                source = manifest.sources[unique_id]
+                selected.append(source)
+        return selected
+
     def get_node_selector(self) -> ResourceTypeSelector:
         if self.manifest is None or self.graph is None:
             raise DbtInternalError("manifest and graph must be set to perform node selection")


Test patch:
diff --git a/tests/functional/docs/test_generate.py b/tests/functional/docs/test_generate.py
--- a/tests/functional/docs/test_generate.py
+++ b/tests/functional/docs/test_generate.py
@@ -28,3 +28,8 @@ def test_select_limits_catalog(self, project):
         catalog = run_dbt(["docs", "generate", "--select", "my_model"])
         assert len(catalog.nodes) == 1
         assert "model.test.my_model" in catalog.nodes
+
+    def test_select_limits_no_match(self, project):
+        run_dbt(["run"])
+        catalog = run_dbt(["docs", "generate", "--select", "my_missing_model"])
+        assert len(catalog.nodes) == 0
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
@@ -1,8 +1,12 @@
+import dataclasses
+
 import agate
 import decimal
 import unittest
 from unittest import mock
 
+from dbt.adapters.base import BaseRelation
+from dbt.contracts.relation import Path
 from dbt.task.debug import DebugTask
 
 from dbt.adapters.base.query_headers import MacroQueryStringSetter
@@ -322,34 +326,45 @@ def test_set_zero_keepalive(self, psycopg2):
         )
 
     @mock.patch.object(PostgresAdapter, "execute_macro")
-    @mock.patch.object(PostgresAdapter, "_get_catalog_relations_by_info_schema")
+    @mock.patch.object(PostgresAdapter, "_get_catalog_relations")
     def test_get_catalog_various_schemas(self, mock_get_relations, mock_execute):
+        self.catalog_test(mock_get_relations, mock_execute, False)
+
+    @mock.patch.object(PostgresAdapter, "execute_macro")
+    @mock.patch.object(PostgresAdapter, "_get_catalog_relations")
+    def test_get_filtered_catalog(self, mock_get_relations, mock_execute):
+        self.catalog_test(mock_get_relations, mock_execute, True)
+
+    def catalog_test(self, mock_get_relations, mock_execute, filtered=False):
         column_names = ["table_database", "table_schema", "table_name"]
-        rows = [
-            ("dbt", "foo", "bar"),
-            ("dbt", "FOO", "baz"),
-            ("dbt", None, "bar"),
-            ("dbt", "quux", "bar"),
-            ("dbt", "skip", "bar"),
+        relations = [
+            BaseRelation(path=Path(database="dbt", schema="foo", identifier="bar")),
+            BaseRelation(path=Path(database="dbt", schema="FOO", identifier="baz")),
+            BaseRelation(path=Path(database="dbt", schema=None, identifier="bar")),
+            BaseRelation(path=Path(database="dbt", schema="quux", identifier="bar")),
+            BaseRelation(path=Path(database="dbt", schema="skip", identifier="bar")),
         ]
+        rows = list(map(lambda x: dataclasses.astuple(x.path), relations))
         mock_execute.return_value = agate.Table(rows=rows, column_names=column_names)
 
-        mock_get_relations.return_value = {
-            mock.MagicMock(database="dbt"): [
-                mock.MagicMock(schema="foo"),
-                mock.MagicMock(schema="FOO"),
-                mock.MagicMock(schema="quux"),
-            ]
-        }
+        mock_get_relations.return_value = relations
 
         mock_manifest = mock.MagicMock()
         mock_manifest.get_used_schemas.return_value = {("dbt", "foo"), ("dbt", "quux")}
 
-        catalog, exceptions = self.adapter.get_catalog(mock_manifest)
-        self.assertEqual(
-            set(map(tuple, catalog)),
-            {("dbt", "foo", "bar"), ("dbt", "FOO", "baz"), ("dbt", "quux", "bar")},
-        )
+        if filtered:
+            catalog, exceptions = self.adapter.get_filtered_catalog(
+                mock_manifest, set([relations[0], relations[3]])
+            )
+        else:
+            catalog, exceptions = self.adapter.get_catalog(mock_manifest)
+
+        tupled_catalog = set(map(tuple, catalog))
+        if filtered:
+            self.assertEqual(tupled_catalog, {rows[0], rows[3]})
+        else:
+            self.assertEqual(tupled_catalog, {rows[0], rows[1], rows[3]})
+
         self.assertEqual(exceptions, [])