Override list_relations_without_caching

.changes/unreleased/Features-20230407-104723.yaml

@@ -0,0 +1,5 @@
+kind: Features
+time: 2023-04-07T10:47:23.105369-07:00
+custom:
+  Author: jiezhen-chen
+  Issue: 17 179 217

dbt/adapters/redshift/__init__.py

@@ -1,5 +1,6 @@
 from dbt.adapters.redshift.connections import RedshiftConnectionManager  # noqa
 from dbt.adapters.redshift.connections import RedshiftCredentials
+
 from dbt.adapters.redshift.column import RedshiftColumn  # noqa
 from dbt.adapters.redshift.relation import RedshiftRelation  # noqa: F401
 from dbt.adapters.redshift.impl import RedshiftAdapter

dbt/adapters/redshift/connections.py

@@ -19,10 +19,13 @@
 from dbt.dataclass_schema import FieldEncoder, dbtClassMixin, StrEnum
 from dbt.helper_types import Port
 
+
 logger = AdapterLogger("Redshift")
 
+
 drop_lock: Lock = dbt.flags.MP_CONTEXT.Lock()  # type: ignore
 
+
 IAMDuration = NewType("IAMDuration", int)
 
 
@@ -77,6 +80,7 @@ class RedshiftCredentials(Credentials):
     sslmode: Optional[str] = None
     retries: int = 1
     region: Optional[str] = None  # if not provided, will be determined from host
+    current_db_only: Optional[bool] = False
     autocommit: Optional[bool] = False
 
     _ALIASES = {"dbname": "database", "pass": "password"}
@@ -127,6 +131,7 @@ def get_connect_method(self):
             "db_groups": self.credentials.db_groups,
             "region": self.credentials.region,
             "timeout": self.credentials.connect_timeout,
+            "database_metadata_current_db_only": self.credentials.current_db_only,
         }
         if kwargs["region"] is None:
             logger.debug("No region provided, attempting to determine from host.")

dbt/adapters/redshift/impl.py

@@ -1,14 +1,15 @@
 from dataclasses import dataclass
-from typing import Optional, Set, Any, Dict, Type
+from typing import Optional, Set, Any, Dict, Type, List
 from collections import namedtuple
-
 from dbt.adapters.base import PythonJobHelper
 from dbt.adapters.base.impl import AdapterConfig, ConstraintSupport
 from dbt.adapters.sql import SQLAdapter
 from dbt.adapters.base.meta import available
 from dbt.contracts.connection import AdapterResponse
 from dbt.contracts.graph.nodes import ConstraintType
 from dbt.events import AdapterLogger
+from dbt.adapters.base.relation import BaseRelation
+
 
 import dbt.exceptions
 
@@ -30,7 +31,7 @@ class RedshiftConfig(AdapterConfig):
 
 
 class RedshiftAdapter(SQLAdapter):
-    Relation = RedshiftRelation
+    Relation = RedshiftRelation  # type: ignore
     ConnectionManager = RedshiftConnectionManager
     connections: RedshiftConnectionManager
     Column = RedshiftColumn  # type: ignore
@@ -114,6 +115,44 @@ def valid_incremental_strategies(self):
     def timestamp_add_sql(self, add_to: str, number: int = 1, interval: str = "hour") -> str:
         return f"{add_to} + interval '{number} {interval}'"
 
+    def _get_cursor(self):
+        return self.connections.get_thread_connection().handle.cursor()
+
+    def _get_tables(self, database: Optional[str], schema: Optional[str]):
+        cursor = self._get_cursor()
+        results = []
+        for table in cursor.get_tables(
+            catalog=database,
+            schema_pattern=schema,
+            table_name_pattern=None,
+            types=["VIEW", "TABLE"],
+        ):
+            results.append([table[0], table[1], table[2], table[3]])
+            # kinda feel like it might not be a good thing to do this by index
+        return results
+
+    def list_relations_without_caching(  # type: ignore
+        self, schema_relation: BaseRelation
+    ) -> List[RedshiftRelation]:
+        results = self._get_tables(schema_relation.database, schema_relation.schema)
+        relations = []
+        quote_policy = {"database": True, "schema": True, "identifier": True}
+        for _database, _schema, name, _type in results:
+            try:
+                _type = self.Relation.get_relation_type(_type.lower())
+            except ValueError:
+                _type = self.Relation.External
+            relations.append(
+                self.Relation.create(
+                    database=_database,
+                    schema=_schema,
+                    identifier=name,
+                    quote_policy=quote_policy,
+                    type=_type,
+                )
+            )
+        return relations
+
     def _link_cached_database_relations(self, schemas: Set[str]):
         """
         :param schemas: The set of schemas that should have links added.

dbt/include/redshift/macros/adapters.sql

@@ -224,11 +224,6 @@
 {% endmacro %}
 
 
-{% macro redshift__list_relations_without_caching(schema_relation) %}
-  {{ return(postgres__list_relations_without_caching(schema_relation)) }}
-{% endmacro %}
-
-
 {% macro redshift__information_schema_name(database) -%}
   {{ return(postgres__information_schema_name(database)) }}
 {%- endmacro %}

tests/functional/adapter/test_adapter_methods.py

@@ -0,0 +1,125 @@
+import pytest
+
+from dbt.tests.util import run_dbt, check_relations_equal
+from dbt.tests.fixtures.project import write_project_files
+
+
+tests__get_relation_invalid = """
+{% set upstream = ref('upstream') %}
+{% set relations = adapter.get_relation(database=upstream.database, schema=upstream.schema, identifier="doesnotexist") %}
+{% set limit_query = 0 %}
+{% if relations.identifier %}
+    {% set limit_query = 1 %}
+{% endif %}
+
+select 1 as id limit {{ limit_query }}
+
+"""
+
+models__upstream_sql = """
+select 1 as id
+
+"""
+
+models__expected_sql = """
+select 1 as valid_relation
+
+"""
+
+models__model_sql = """
+
+{% set upstream = ref('upstream') %}
+
+select * from {{ upstream }}
+
+"""
+
+models__call_get_relation = """
+
+{% set model = ref('model') %}
+
+{% set relation = adapter.get_relation(database=model.database, schema=model.schema, identifier=model.identifier) %}
+{% if relation.identifier == model.identifier %}
+
+select 1 as valid_relation
+
+{% else %}
+
+select 0 as valid_relation
+
+{% endif %}
+
+"""
+
+models__get_relation_type = """
+
+{% set base_view = ref('base_view') %}
+
+{% set relation = adapter.get_relation(database=base_view.database, schema=base_view.schema, identifier=base_view.identifier) %}
+{% if relation.type == 'view' %}
+
+select 1 as valid_type
+
+{% else %}
+
+select 0 as valid_type
+
+{% endif %}
+
+"""
+
+
+class RedshiftAdapterMethod:
+    @pytest.fixture(scope="class")
+    def tests(self):
+        return {"get_relation_invalid.sql": tests__get_relation_invalid}
+
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "upstream.sql": models__upstream_sql,
+            "expected.sql": models__expected_sql,
+            "model.sql": models__model_sql,
+            "call_get_relation.sql": models__call_get_relation,
+            "base_view.sql": "{{ config(bind=True) }} select * from {{ ref('model') }}",
+            "get_relation_type.sql": models__get_relation_type,
+            "expected_type.sql": "select 1 as valid_type",
+        }
+
+    @pytest.fixture(scope="class")
+    def project_files(
+        self,
+        project_root,
+        tests,
+        models,
+    ):
+        write_project_files(project_root, "tests", tests)
+        write_project_files(project_root, "models", models)
+
+    @pytest.fixture(scope="class")
+    def project_config_update(self):
+        return {
+            "name": "adapter_methods",
+        }
+
+    # snowflake need all tables in CAP name
+    @pytest.fixture(scope="class")
+    def equal_tables(self):
+        return ["call_get_relation", "expected"]
+
+    @pytest.fixture(scope="class")
+    def equal_types(self):
+        return ["get_relation_type", "expected_type"]
+
+    def test_adapter_methods(self, project, equal_tables):
+        run_dbt(["compile"])  # trigger any compile-time issues
+        result = run_dbt()
+        assert len(result) == 7
+
+        run_dbt(["test"])
+        check_relations_equal(project.adapter, equal_tables)
+        check_relations_equal(project.adapter, ["get_relation_type", "expected_type"])
+
+
+class TestRedshiftAdapterMethod(RedshiftAdapterMethod):
+    pass

tests/unit/test_redshift_adapter.py

@@ -74,6 +74,7 @@ def test_implicit_database_conn(self):
             db_groups=[],
             timeout=30,
             region="us-east-1",
+            database_metadata_current_db_only=False,
         )
 
     @mock.patch("redshift_connector.connect", Mock())
@@ -92,6 +93,7 @@ def test_explicit_database_conn(self):
             db_groups=[],
             region="us-east-1",
             timeout=30,
+            database_metadata_current_db_only=False,
         )
 
     @mock.patch("redshift_connector.connect", Mock())
@@ -117,6 +119,7 @@ def test_explicit_iam_conn_without_profile(self):
             profile=None,
             timeout=30,
             port=5439,
+            database_metadata_current_db_only=False,
         )
 
     @mock.patch("redshift_connector.connect", Mock())
@@ -145,6 +148,7 @@ def test_explicit_iam_conn_with_profile(self):
             profile="test",
             timeout=30,
             port=5439,
+            database_metadata_current_db_only=False,
         )
 
     @mock.patch("redshift_connector.connect", Mock())
@@ -171,6 +175,7 @@ def test_explicit_iam_serverless_with_profile(self):
             profile="test",
             timeout=30,
             port=5439,
+            database_metadata_current_db_only=False,
         )
 
     @mock.patch("redshift_connector.connect", Mock())
@@ -199,6 +204,7 @@ def test_explicit_region(self):
             profile="test",
             timeout=30,
             port=5439,
+            database_metadata_current_db_only=False,
         )
 
     @mock.patch("redshift_connector.connect", Mock())
@@ -228,6 +234,7 @@ def test_explicit_region_failure(self):
                 profile="test",
                 timeout=30,
                 port=5439,
+                database_metadata_current_db_only=False,
             )
 
     @mock.patch("redshift_connector.connect", Mock())
@@ -257,6 +264,7 @@ def test_explicit_invalid_region(self):
                 profile="test",
                 timeout=30,
                 port=5439,
+                database_metadata_current_db_only=False,
             )
 
     @mock.patch("redshift_connector.connect", Mock())
@@ -284,6 +292,7 @@ def test_serverless_iam_failure(self):
                 profile="test",
                 port=5439,
                 timeout=30,
+                database_metadata_current_db_only=False,
             )
         self.assertTrue("'host' must be provided" in context.exception.msg)
 
@@ -450,6 +459,35 @@ def test_add_query_success(self):
             "select * from test3", True, bindings=None, abridge_sql_log=False
         )
 
+    @mock.patch.object(
+        dbt.adapters.redshift.connections.SQLConnectionManager, "get_thread_connection"
+    )
+    def mock_cursor(self, mock_get_thread_conn):
+        conn = mock.MagicMock
+        mock_get_thread_conn.return_value = conn
+        mock_handle = mock.MagicMock
+        conn.return_value = mock_handle
+        mock_cursor = mock.MagicMock
+        mock_handle.return_value = mock_cursor
+        return mock_cursor
+
+    @mock.patch("dbt.adapters.redshift.impl.RedshiftAdapter._get_cursor")
+    def test_get_tables(self, mock_cursor):
+        mock_cursor.return_value.get_tables.return_value = [
+            ("apple", "banana", "cherry", "orange")
+        ]
+        results = self.adapter._get_tables(database="somedb", schema="someschema")
+        self.assertTrue(results == [["apple", "banana", "cherry", "orange"]])
+
+    @mock.patch("dbt.adapters.redshift.impl.RedshiftAdapter._get_tables")
+    def test_list_relations_without_caching(self, mock_get_tables):
+        mock_get_tables.return_value = [["somedb", "someschema", "sometb", "VIEW"]]
+        mock_schema = mock.MagicMock(database="somedb", schema="someschema")
+        results = self.adapter.list_relations_without_caching(mock_schema)
+        self.assertTrue(results[0].database == "somedb")
+        self.assertTrue(results[0].schema == "someschema")
+        self.assertTrue(results[0].identifier == "sometb")
+
 
 class TestRedshiftAdapterConversions(TestAdapterConversions):
     def test_convert_text_type(self):