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

@@ -4,6 +4,7 @@
     RedshiftConnectionManager,
     RedshiftCredentials,
 )
+
 from dbt.adapters.redshift.relation import RedshiftRelation  # noqa: F401
 from dbt.adapters.redshift.impl import RedshiftAdapter
 from dbt.include import redshift

dbt/adapters/redshift/connections.py

@@ -32,8 +32,10 @@ def get_message(self) -> str:
 
 logger = AdapterLogger("Redshift")
 
+
 drop_lock: Lock = dbt.flags.MP_CONTEXT.Lock()  # type: ignore
 
+
 IAMDuration = NewType("IAMDuration", int)
 
 

dbt/adapters/redshift/impl.py

@@ -1,14 +1,15 @@
 from dataclasses import dataclass
 from typing import Optional, Set, Any, Dict, Type
 from collections import namedtuple
-
 from dbt.adapters.base import PythonJobHelper
 from dbt.adapters.base.impl import AdapterConfig, ConstraintSupport
 from dbt.adapters.base.meta import available
 from dbt.adapters.sql import SQLAdapter
 from dbt.contracts.connection import AdapterResponse
 from dbt.contracts.graph.nodes import ConstraintType
 from dbt.events import AdapterLogger
+
+
 import dbt.exceptions
 
 from dbt.adapters.redshift import RedshiftConnectionManager, RedshiftRelation
@@ -31,7 +32,7 @@ class RedshiftConfig(AdapterConfig):
 
 
 class RedshiftAdapter(SQLAdapter):
-    Relation = RedshiftRelation
+    Relation = RedshiftRelation  # type: ignore
     ConnectionManager = RedshiftConnectionManager
     connections: RedshiftConnectionManager
 

dbt/include/redshift/macros/adapters.sql

@@ -223,33 +223,33 @@
   {{ return(sql_convert_columns_in_relation(table)) }}
 {% endmacro %}
 
-
 {% macro redshift__list_relations_without_caching(schema_relation) %}
+
   {% call statement('list_relations_without_caching', fetch_result=True) -%}
     select
-      '{{ schema_relation.database }}' as database,
-      tablename as name,
-      schemaname as schema,
-      'table' as type
-    from pg_tables
-    where schemaname ilike '{{ schema_relation.schema }}'
+        table_catalog as database,
+        table_name as name,
+        table_schema as schema,
+        'table' as type
+    from information_schema.tables
+    where table_schema ilike '{{ schema_relation.schema }}'
+    and table_type = 'BASE TABLE'
     union all
     select
-      '{{ schema_relation.database }}' as database,
-      viewname as name,
-      schemaname as schema,
+      table_catalog as database,
+      table_name as name,
+      table_schema as schema,
       case
-        when definition ilike '%create materialized view%'
+        when view_definition ilike '%create materialized view%'
           then 'materialized_view'
         else 'view'
       end as type
-    from pg_views
-    where schemaname ilike '{{ schema_relation.schema }}'
+    from information_schema.views
+    where table_schema ilike '{{ schema_relation.schema }}'
   {% endcall %}
   {{ return(load_result('list_relations_without_caching').table) }}
 {% endmacro %}
 
-
 {% macro redshift__information_schema_name(database) -%}
   {{ return(postgres__information_schema_name(database)) }}
 {%- endmacro %}
@@ -259,7 +259,6 @@
   {{ return(postgres__list_schemas(database)) }}
 {%- endmacro %}
 
-
 {% macro redshift__check_schema_exists(information_schema, schema) -%}
   {{ return(postgres__check_schema_exists(information_schema, schema)) }}
 {%- endmacro %}

tests/functional/adapter/test_adapter_methods.py

@@ -0,0 +1,115 @@
+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",
+        }
+
+    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",
+        }
+
+    def test_adapter_methods(self, project):
+        run_dbt(["compile"])  # trigger any compile-time issues
+        result = run_dbt()
+        assert len(result) == 7
+
+        run_dbt(["test"])
+        check_relations_equal(project.adapter, ["call_get_relation", "expected"])
+        check_relations_equal(project.adapter, ["get_relation_type", "expected_type"])
+
+
+class TestRedshiftAdapterMethod(RedshiftAdapterMethod):
+    pass

tests/unit/test_redshift_adapter.py

@@ -63,6 +63,23 @@ def adapter(self):
             inject_adapter(self._adapter, RedshiftPlugin)
         return self._adapter
 
+    @mock.patch("redshift_connector.connect", Mock())
+    def test_implicit_database_conn(self):
+        connection = self.adapter.acquire_connection("dummy")
+        connection.handle
+        redshift_connector.connect.assert_called_once_with(
+            host="thishostshouldnotexist.test.us-east-1",
+            database="redshift",
+            user="root",
+            password="password",
+            port=5439,
+            auto_create=False,
+            db_groups=[],
+            timeout=None,
+            region=None,
+            **DEFAULT_SSL_CONFIG,
+        )
+
     @mock.patch("redshift_connector.connect", Mock())
     def test_explicit_region_with_database_conn(self):
         self.config.method = "database"
@@ -211,6 +228,66 @@ def test_explicit_region(self):
             **DEFAULT_SSL_CONFIG,
         )
 
+    @mock.patch("redshift_connector.connect", Mock())
+    @mock.patch("boto3.Session", Mock())
+    def test_explicit_region_failure(self):
+        # Failure test with no region
+        self.config.credentials = self.config.credentials.replace(
+            method="iam",
+            iam_profile="test",
+            host="doesnotexist.1233_no_region",
+            region=None,
+        )
+
+        with self.assertRaises(dbt.exceptions.FailedToConnectError):
+            connection = self.adapter.acquire_connection("dummy")
+            connection.handle
+            redshift_connector.connect.assert_called_once_with(
+                iam=True,
+                host="doesnotexist.1233_no_region",
+                database="redshift",
+                cluster_identifier=None,
+                auto_create=False,
+                db_groups=[],
+                db_user="root",
+                password="",
+                user="",
+                profile="test",
+                timeout=None,
+                port=5439,
+                **DEFAULT_SSL_CONFIG,
+            )
+
+    @mock.patch("redshift_connector.connect", Mock())
+    @mock.patch("boto3.Session", Mock())
+    def test_explicit_invalid_region(self):
+        # Invalid region test
+        self.config.credentials = self.config.credentials.replace(
+            method="iam",
+            iam_profile="test",
+            host="doesnotexist.1233_no_region.us-not-a-region-1",
+            region=None,
+        )
+
+        with self.assertRaises(dbt.exceptions.FailedToConnectError):
+            connection = self.adapter.acquire_connection("dummy")
+            connection.handle
+            redshift_connector.connect.assert_called_once_with(
+                iam=True,
+                host="doesnotexist.1233_no_region",
+                database="redshift",
+                cluster_identifier=None,
+                auto_create=False,
+                db_groups=[],
+                db_user="root",
+                password="",
+                user="",
+                profile="test",
+                timeout=None,
+                port=5439,
+                **DEFAULT_SSL_CONFIG,
+            )
+
     @mock.patch("redshift_connector.connect", Mock())
     def test_sslmode_disable(self):
         self.config.credentials.sslmode = "disable"
@@ -498,6 +575,18 @@ 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
+
 
 class TestRedshiftAdapterConversions(TestAdapterConversions):
     def test_convert_text_type(self):