Repository: dbt-labs/dbt-bigquery
Pull request #933 (2023-09-21 22:55:44)

Address type annotation issues and clean up impl
- pull in changes from #935 
- update impl ConnectionManager typing and move list_datasets into BigQueryConnectionManager
- refactor unit tests and add one to cover list_datasets

resolves #
[docs](https://github.com/dbt-labs/docs.getdbt.com/issues/new/choose) dbt-labs/docs.getdbt.com/#

<!---
  Include the number of the issue addressed by this PR above if applicable.
  PRs for code changes without an associated issue *will not be merged*.
  See CONTRIBUTING.md for more information.

  Include the number of the docs issue that was opened for this PR. If
  this change has no user-facing implications, "N/A" suffices instead. New
  docs tickets can be created by clicking the link above or by going to
  https://github.com/dbt-labs/docs.getdbt.com/issues/new/choose.
-->

### Problem

<!---
  Describe the problem this PR is solving. What is the application state
  before this PR is merged?
-->

### Solution

<!---
  Describe the way this PR solves the above problem. Add as much detail as you
  can to help reviewers understand your changes. Include any alternatives and
  tradeoffs you considered.
-->

### Checklist

- [ ] I have read [the contributing guide](https://github.com/dbt-labs/dbt-core/blob/main/CONTRIBUTING.md) and understand what's expected of me
- [ ] I have run this code in development and it appears to resolve the stated issue
- [ ] This PR includes tests, or tests are not required/relevant for this PR
- [ ] This PR has no interface changes (e.g. macros, cli, logs, json artifacts, config files, adapter interface, etc) or this PR has already received feedback and approval from Product or DX


github-actions[bot]: Thank you for your pull request! We could not find a changelog entry for this change. For details on how to document a change, see the [dbt-bigquery contributing guide](https://github.com/dbt-labs/dbt-bigquery/blob/main/CONTRIBUTING.md).

Patch:
diff --git a/.changes/unreleased/Under the Hood-20230921-155645.yaml b/.changes/unreleased/Under the Hood-20230921-155645.yaml
new file mode 100644
index 0000000..12cd663
--- /dev/null
+++ b/.changes/unreleased/Under the Hood-20230921-155645.yaml
@@ -0,0 +1,6 @@
+kind: Under the Hood
+body: Address type annotation issues and remove protected method ref from impl
+time: 2023-09-21T15:56:45.329798-07:00
+custom:
+  Author: colin-rogers-dbt
+  Issue: "933"
diff --git a/.changes/unreleased/Under the Hood-20230922-125327.yaml b/.changes/unreleased/Under the Hood-20230922-125327.yaml
new file mode 100644
index 0000000..9ce8713
--- /dev/null
+++ b/.changes/unreleased/Under the Hood-20230922-125327.yaml
@@ -0,0 +1,6 @@
+kind: Under the Hood
+body: Fixed a mypy failure by reworking BigQueryAdapter constructor.
+time: 2023-09-22T12:53:27.339599-04:00
+custom:
+  Author: peterallenwebb
+  Issue: "934"
diff --git a/dbt/adapters/bigquery/connections.py b/dbt/adapters/bigquery/connections.py
index c136042..7799ecb 100644
--- a/dbt/adapters/bigquery/connections.py
+++ b/dbt/adapters/bigquery/connections.py
@@ -700,6 +700,20 @@ class BigQueryConnectionManager(BaseConnectionManager):
 
         self._retry_and_handle(msg="create dataset", conn=conn, fn=fn)
 
+    def list_dataset(self, database: str):
+        # the database string we get here is potentially quoted. Strip that off
+        # for the API call.
+        database = database.strip("`")
+        conn = self.get_thread_connection()
+        client = conn.handle
+
+        def query_schemas():
+            # this is similar to how we have to deal with listing tables
+            all_datasets = client.list_datasets(project=database, max_results=10000)
+            return [ds.dataset_id for ds in all_datasets]
+
+        return self._retry_and_handle(msg="list dataset", conn=conn, fn=query_schemas)
+
     def _query_and_results(
         self,
         client,
diff --git a/dbt/adapters/bigquery/impl.py b/dbt/adapters/bigquery/impl.py
index bb04c78..8fc1b69 100644
--- a/dbt/adapters/bigquery/impl.py
+++ b/dbt/adapters/bigquery/impl.py
@@ -217,6 +217,10 @@ class BigQueryAdapter(BaseAdapter):
         ConstraintType.foreign_key: ConstraintSupport.ENFORCED,
     }
 
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.connections: BigQueryConnectionManager = self.connections
+
     ###
     # Implementations of abstract methods
     ###
@@ -267,18 +271,7 @@ class BigQueryAdapter(BaseAdapter):
 
     @available
     def list_schemas(self, database: str) -> List[str]:
-        # the database string we get here is potentially quoted. Strip that off
-        # for the API call.
-        database = database.strip("`")
-        conn = self.connections.get_thread_connection()
-        client = conn.handle
-
-        def query_schemas():
-            # this is similar to how we have to deal with listing tables
-            all_datasets = client.list_datasets(project=database, max_results=10000)
-            return [ds.dataset_id for ds in all_datasets]
-
-        return self.connections._retry_and_handle(msg="list dataset", conn=conn, fn=query_schemas)
+        return self.connections.list_dataset(database)
 
     @available.parse(lambda *a, **k: False)
     def check_schema_exists(self, database: str, schema: str) -> bool:
@@ -481,40 +474,6 @@ class BigQueryAdapter(BaseAdapter):
             bq_schema.append(SchemaField(col_name, type_))  # type: ignore[arg-type]
         return bq_schema
 
-    def _materialize_as_view(self, model: Dict[str, Any]) -> str:
-        model_database = model.get("database")
-        model_schema = model.get("schema")
-        model_alias = model.get("alias")
-        model_code = model.get("compiled_code")
-
-        logger.debug("Model SQL ({}):\n{}".format(model_alias, model_code))
-        self.connections.create_view(
-            database=model_database, schema=model_schema, table_name=model_alias, sql=model_code
-        )
-        return "CREATE VIEW"
-
-    def _materialize_as_table(
-        self,
-        model: Dict[str, Any],
-        model_sql: str,
-        decorator: Optional[str] = None,
-    ) -> str:
-        model_database = model.get("database")
-        model_schema = model.get("schema")
-        model_alias = model.get("alias")
-
-        if decorator is None:
-            table_name = model_alias
-        else:
-            table_name = "{}${}".format(model_alias, decorator)
-
-        logger.debug("Model SQL ({}):\n{}".format(table_name, model_sql))
-        self.connections.create_table(
-            database=model_database, schema=model_schema, table_name=table_name, sql=model_sql
-        )
-
-        return "CREATE TABLE"
-
     @available.parse(lambda *a, **k: "")
     def copy_table(self, source, destination, materialization):
         if materialization == "incremental":


Test patch:
diff --git a/tests/unit/test_bigquery_adapter.py b/tests/unit/test_bigquery_adapter.py
index 10cb3f5..4db2ce8 100644
--- a/tests/unit/test_bigquery_adapter.py
+++ b/tests/unit/test_bigquery_adapter.py
@@ -1,26 +1,19 @@
-import time
-
 import agate
 import decimal
-import json
 import string
 import random
 import re
 import pytest
 import unittest
-from contextlib import contextmanager
-from requests.exceptions import ConnectionError
-from unittest.mock import patch, MagicMock, Mock, create_autospec, ANY
+from unittest.mock import patch, MagicMock, create_autospec
 
 import dbt.dataclass_schema
 
 from dbt.adapters.bigquery import PartitionConfig
-from dbt.adapters.bigquery import BigQueryCredentials
 from dbt.adapters.bigquery import BigQueryAdapter
 from dbt.adapters.bigquery import BigQueryRelation
 from dbt.adapters.bigquery import Plugin as BigQueryPlugin
 from google.cloud.bigquery.table import Table
-from dbt.adapters.bigquery.connections import BigQueryConnectionManager
 from dbt.adapters.bigquery.connections import _sanitize_label, _VALIDATE_LABEL_LENGTH_LIMIT
 from dbt.adapters.base.query_headers import MacroQueryStringSetter
 from dbt.clients import agate_helper
@@ -543,179 +536,6 @@ class TestBigQueryInformationSchema(unittest.TestCase):
         assert other_schema.quote_policy.database is False
 
 
-class TestBigQueryConnectionManager(unittest.TestCase):
-    def setUp(self):
-        credentials = Mock(BigQueryCredentials)
-        profile = Mock(query_comment=None, credentials=credentials)
-        self.connections = BigQueryConnectionManager(profile=profile)
-
-        self.mock_client = Mock(dbt.adapters.bigquery.impl.google.cloud.bigquery.Client)
-        self.mock_connection = MagicMock()
-
-        self.mock_connection.handle = self.mock_client
-
-        self.connections.get_thread_connection = lambda: self.mock_connection
-        self.connections.get_job_retry_deadline_seconds = lambda x: None
-        self.connections.get_job_retries = lambda x: 1
-
-    @patch("dbt.adapters.bigquery.connections._is_retryable", return_value=True)
-    def test_retry_and_handle(self, is_retryable):
-        self.connections.DEFAULT_MAXIMUM_DELAY = 2.0
-
-        @contextmanager
-        def dummy_handler(msg):
-            yield
-
-        self.connections.exception_handler = dummy_handler
-
-        class DummyException(Exception):
-            """Count how many times this exception is raised"""
-
-            count = 0
-
-            def __init__(self):
-                DummyException.count += 1
-
-        def raiseDummyException():
-            raise DummyException()
-
-        with self.assertRaises(DummyException):
-            self.connections._retry_and_handle(
-                "some sql", Mock(credentials=Mock(retries=8)), raiseDummyException
-            )
-            self.assertEqual(DummyException.count, 9)
-
-    @patch("dbt.adapters.bigquery.connections._is_retryable", return_value=True)
-    def test_retry_connection_reset(self, is_retryable):
-        self.connections.open = MagicMock()
-        self.connections.close = MagicMock()
-        self.connections.DEFAULT_MAXIMUM_DELAY = 2.0
-
-        @contextmanager
-        def dummy_handler(msg):
-            yield
-
-        self.connections.exception_handler = dummy_handler
-
-        def raiseConnectionResetError():
-            raise ConnectionResetError("Connection broke")
-
-        mock_conn = Mock(credentials=Mock(retries=1))
-        with self.assertRaises(ConnectionResetError):
-            self.connections._retry_and_handle("some sql", mock_conn, raiseConnectionResetError)
-        self.connections.close.assert_called_once_with(mock_conn)
-        self.connections.open.assert_called_once_with(mock_conn)
-
-    def test_is_retryable(self):
-        _is_retryable = dbt.adapters.bigquery.connections._is_retryable
-        exceptions = dbt.adapters.bigquery.impl.google.cloud.exceptions
-        internal_server_error = exceptions.InternalServerError("code broke")
-        bad_request_error = exceptions.BadRequest("code broke")
-        connection_error = ConnectionError("code broke")
-        client_error = exceptions.ClientError("bad code")
-        rate_limit_error = exceptions.Forbidden(
-            "code broke", errors=[{"reason": "rateLimitExceeded"}]
-        )
-
-        self.assertTrue(_is_retryable(internal_server_error))
-        self.assertTrue(_is_retryable(bad_request_error))
-        self.assertTrue(_is_retryable(connection_error))
-        self.assertFalse(_is_retryable(client_error))
-        self.assertTrue(_is_retryable(rate_limit_error))
-
-    def test_drop_dataset(self):
-        mock_table = Mock()
-        mock_table.reference = "table1"
-        self.mock_client.list_tables.return_value = [mock_table]
-
-        self.connections.drop_dataset("project", "dataset")
-
-        self.mock_client.list_tables.assert_not_called()
-        self.mock_client.delete_table.assert_not_called()
-        self.mock_client.delete_dataset.assert_called_once()
-
-    @patch("dbt.adapters.bigquery.impl.google.cloud.bigquery")
-    def test_query_and_results(self, mock_bq):
-        self.mock_client.query = Mock(return_value=Mock(state="DONE"))
-        self.connections._query_and_results(
-            self.mock_client,
-            "sql",
-            {"job_param_1": "blah"},
-            job_creation_timeout=15,
-            job_execution_timeout=3,
-        )
-
-        mock_bq.QueryJobConfig.assert_called_once()
-        self.mock_client.query.assert_called_once_with(
-            query="sql", job_config=mock_bq.QueryJobConfig(), timeout=15
-        )
-
-    @patch("dbt.adapters.bigquery.impl.google.cloud.bigquery")
-    def test_query_and_results_timeout(self, mock_bq):
-        self.mock_client.query = Mock(
-            return_value=Mock(result=lambda *args, **kwargs: time.sleep(4))
-        )
-        with pytest.raises(dbt.exceptions.DbtRuntimeError) as exc:
-            self.connections._query_and_results(
-                self.mock_client,
-                "sql",
-                {"job_param_1": "blah"},
-                job_creation_timeout=15,
-                job_execution_timeout=1,
-            )
-
-        mock_bq.QueryJobConfig.assert_called_once()
-        self.mock_client.query.assert_called_once_with(
-            query="sql", job_config=mock_bq.QueryJobConfig(), timeout=15
-        )
-        assert "Query exceeded configured timeout of 1s" in str(exc.value)
-
-    def test_copy_bq_table_appends(self):
-        self._copy_table(write_disposition=dbt.adapters.bigquery.impl.WRITE_APPEND)
-        args, kwargs = self.mock_client.copy_table.call_args
-        self.mock_client.copy_table.assert_called_once_with(
-            [self._table_ref("project", "dataset", "table1")],
-            self._table_ref("project", "dataset", "table2"),
-            job_config=ANY,
-        )
-        args, kwargs = self.mock_client.copy_table.call_args
-        self.assertEqual(
-            kwargs["job_config"].write_disposition, dbt.adapters.bigquery.impl.WRITE_APPEND
-        )
-
-    def test_copy_bq_table_truncates(self):
-        self._copy_table(write_disposition=dbt.adapters.bigquery.impl.WRITE_TRUNCATE)
-        args, kwargs = self.mock_client.copy_table.call_args
-        self.mock_client.copy_table.assert_called_once_with(
-            [self._table_ref("project", "dataset", "table1")],
-            self._table_ref("project", "dataset", "table2"),
-            job_config=ANY,
-        )
-        args, kwargs = self.mock_client.copy_table.call_args
-        self.assertEqual(
-            kwargs["job_config"].write_disposition, dbt.adapters.bigquery.impl.WRITE_TRUNCATE
-        )
-
-    def test_job_labels_valid_json(self):
-        expected = {"key": "value"}
-        labels = self.connections._labels_from_query_comment(json.dumps(expected))
-        self.assertEqual(labels, expected)
-
-    def test_job_labels_invalid_json(self):
-        labels = self.connections._labels_from_query_comment("not json")
-        self.assertEqual(labels, {"query_comment": "not_json"})
-
-    def _table_ref(self, proj, ds, table):
-        return self.connections.table_ref(proj, ds, table)
-
-    def _copy_table(self, write_disposition):
-        source = BigQueryRelation.create(database="project", schema="dataset", identifier="table1")
-        destination = BigQueryRelation.create(
-            database="project", schema="dataset", identifier="table2"
-        )
-        self.connections.copy_bq_table(source, destination, write_disposition)
-
-
 class TestBigQueryAdapter(BaseTestBigQueryAdapter):
     def test_copy_table_materialization_table(self):
         adapter = self.get_adapter("oauth")
diff --git a/tests/unit/test_bigquery_connection_manager.py b/tests/unit/test_bigquery_connection_manager.py
new file mode 100644
index 0000000..d6c3f64
--- /dev/null
+++ b/tests/unit/test_bigquery_connection_manager.py
@@ -0,0 +1,198 @@
+import time
+import json
+import pytest
+import unittest
+from contextlib import contextmanager
+from requests.exceptions import ConnectionError
+from unittest.mock import patch, MagicMock, Mock, ANY
+
+import dbt.dataclass_schema
+
+from dbt.adapters.bigquery import BigQueryCredentials
+from dbt.adapters.bigquery import BigQueryRelation
+from dbt.adapters.bigquery.connections import BigQueryConnectionManager
+import dbt.exceptions
+from dbt.logger import GLOBAL_LOGGER as logger  # noqa
+
+
+class TestBigQueryConnectionManager(unittest.TestCase):
+    def setUp(self):
+        credentials = Mock(BigQueryCredentials)
+        profile = Mock(query_comment=None, credentials=credentials)
+        self.connections = BigQueryConnectionManager(profile=profile)
+
+        self.mock_client = Mock(dbt.adapters.bigquery.impl.google.cloud.bigquery.Client)
+        self.mock_connection = MagicMock()
+
+        self.mock_connection.handle = self.mock_client
+
+        self.connections.get_thread_connection = lambda: self.mock_connection
+        self.connections.get_job_retry_deadline_seconds = lambda x: None
+        self.connections.get_job_retries = lambda x: 1
+
+    @patch("dbt.adapters.bigquery.connections._is_retryable", return_value=True)
+    def test_retry_and_handle(self, is_retryable):
+        self.connections.DEFAULT_MAXIMUM_DELAY = 2.0
+
+        @contextmanager
+        def dummy_handler(msg):
+            yield
+
+        self.connections.exception_handler = dummy_handler
+
+        class DummyException(Exception):
+            """Count how many times this exception is raised"""
+
+            count = 0
+
+            def __init__(self):
+                DummyException.count += 1
+
+        def raiseDummyException():
+            raise DummyException()
+
+        with self.assertRaises(DummyException):
+            self.connections._retry_and_handle(
+                "some sql", Mock(credentials=Mock(retries=8)), raiseDummyException
+            )
+            self.assertEqual(DummyException.count, 9)
+
+    @patch("dbt.adapters.bigquery.connections._is_retryable", return_value=True)
+    def test_retry_connection_reset(self, is_retryable):
+        self.connections.open = MagicMock()
+        self.connections.close = MagicMock()
+        self.connections.DEFAULT_MAXIMUM_DELAY = 2.0
+
+        @contextmanager
+        def dummy_handler(msg):
+            yield
+
+        self.connections.exception_handler = dummy_handler
+
+        def raiseConnectionResetError():
+            raise ConnectionResetError("Connection broke")
+
+        mock_conn = Mock(credentials=Mock(retries=1))
+        with self.assertRaises(ConnectionResetError):
+            self.connections._retry_and_handle("some sql", mock_conn, raiseConnectionResetError)
+        self.connections.close.assert_called_once_with(mock_conn)
+        self.connections.open.assert_called_once_with(mock_conn)
+
+    def test_is_retryable(self):
+        _is_retryable = dbt.adapters.bigquery.connections._is_retryable
+        exceptions = dbt.adapters.bigquery.impl.google.cloud.exceptions
+        internal_server_error = exceptions.InternalServerError("code broke")
+        bad_request_error = exceptions.BadRequest("code broke")
+        connection_error = ConnectionError("code broke")
+        client_error = exceptions.ClientError("bad code")
+        rate_limit_error = exceptions.Forbidden(
+            "code broke", errors=[{"reason": "rateLimitExceeded"}]
+        )
+
+        self.assertTrue(_is_retryable(internal_server_error))
+        self.assertTrue(_is_retryable(bad_request_error))
+        self.assertTrue(_is_retryable(connection_error))
+        self.assertFalse(_is_retryable(client_error))
+        self.assertTrue(_is_retryable(rate_limit_error))
+
+    def test_drop_dataset(self):
+        mock_table = Mock()
+        mock_table.reference = "table1"
+        self.mock_client.list_tables.return_value = [mock_table]
+
+        self.connections.drop_dataset("project", "dataset")
+
+        self.mock_client.list_tables.assert_not_called()
+        self.mock_client.delete_table.assert_not_called()
+        self.mock_client.delete_dataset.assert_called_once()
+
+    @patch("dbt.adapters.bigquery.impl.google.cloud.bigquery")
+    def test_query_and_results(self, mock_bq):
+        self.mock_client.query = Mock(return_value=Mock(state="DONE"))
+        self.connections._query_and_results(
+            self.mock_client,
+            "sql",
+            {"job_param_1": "blah"},
+            job_creation_timeout=15,
+            job_execution_timeout=3,
+        )
+
+        mock_bq.QueryJobConfig.assert_called_once()
+        self.mock_client.query.assert_called_once_with(
+            query="sql", job_config=mock_bq.QueryJobConfig(), timeout=15
+        )
+
+    @patch("dbt.adapters.bigquery.impl.google.cloud.bigquery")
+    def test_query_and_results_timeout(self, mock_bq):
+        self.mock_client.query = Mock(
+            return_value=Mock(result=lambda *args, **kwargs: time.sleep(4))
+        )
+        with pytest.raises(dbt.exceptions.DbtRuntimeError) as exc:
+            self.connections._query_and_results(
+                self.mock_client,
+                "sql",
+                {"job_param_1": "blah"},
+                job_creation_timeout=15,
+                job_execution_timeout=1,
+            )
+
+        mock_bq.QueryJobConfig.assert_called_once()
+        self.mock_client.query.assert_called_once_with(
+            query="sql", job_config=mock_bq.QueryJobConfig(), timeout=15
+        )
+        assert "Query exceeded configured timeout of 1s" in str(exc.value)
+
+    def test_copy_bq_table_appends(self):
+        self._copy_table(write_disposition=dbt.adapters.bigquery.impl.WRITE_APPEND)
+        args, kwargs = self.mock_client.copy_table.call_args
+        self.mock_client.copy_table.assert_called_once_with(
+            [self._table_ref("project", "dataset", "table1")],
+            self._table_ref("project", "dataset", "table2"),
+            job_config=ANY,
+        )
+        args, kwargs = self.mock_client.copy_table.call_args
+        self.assertEqual(
+            kwargs["job_config"].write_disposition, dbt.adapters.bigquery.impl.WRITE_APPEND
+        )
+
+    def test_copy_bq_table_truncates(self):
+        self._copy_table(write_disposition=dbt.adapters.bigquery.impl.WRITE_TRUNCATE)
+        args, kwargs = self.mock_client.copy_table.call_args
+        self.mock_client.copy_table.assert_called_once_with(
+            [self._table_ref("project", "dataset", "table1")],
+            self._table_ref("project", "dataset", "table2"),
+            job_config=ANY,
+        )
+        args, kwargs = self.mock_client.copy_table.call_args
+        self.assertEqual(
+            kwargs["job_config"].write_disposition, dbt.adapters.bigquery.impl.WRITE_TRUNCATE
+        )
+
+    def test_job_labels_valid_json(self):
+        expected = {"key": "value"}
+        labels = self.connections._labels_from_query_comment(json.dumps(expected))
+        self.assertEqual(labels, expected)
+
+    def test_job_labels_invalid_json(self):
+        labels = self.connections._labels_from_query_comment("not json")
+        self.assertEqual(labels, {"query_comment": "not_json"})
+
+    def test_list_dataset_correctly_calls_lists_datasets(self):
+        mock_dataset = Mock(dataset_id="d1")
+        mock_list_dataset = Mock(return_value=[mock_dataset])
+        self.mock_client.list_datasets = mock_list_dataset
+        result = self.connections.list_dataset("project")
+        self.mock_client.list_datasets.assert_called_once_with(
+            project="project", max_results=10000
+        )
+        assert result == ["d1"]
+
+    def _table_ref(self, proj, ds, table):
+        return self.connections.table_ref(proj, ds, table)
+
+    def _copy_table(self, write_disposition):
+        source = BigQueryRelation.create(database="project", schema="dataset", identifier="table1")
+        destination = BigQueryRelation.create(
+            database="project", schema="dataset", identifier="table2"
+        )
+        self.connections.copy_bq_table(source, destination, write_disposition)