From e86609a1e15a766eb764a535f277100ca10ee67b Mon Sep 17 00:00:00 2001 From: colin-rogers-dbt <111200756+colin-rogers-dbt@users.noreply.github.com> Date: Mon, 22 Jan 2024 13:59:04 -0800 Subject: [PATCH] Migrate to dbt-adapter and common (#1071) * use dynamic schema in test_grant_access_to.py * use dynamic schema in test_grant_access_to.py * revert setup * replace dbt.common with dbt_common * add dbt-adapters * delete dbt/adapters * fix Credentials import and test fixtures * remove global exceptions import --- .../Under the Hood-20240116-154305.yaml | 6 ++++ dbt/adapters/bigquery/connections.py | 18 +++++----- dbt/adapters/bigquery/gcloud.py | 7 ++-- dbt/adapters/bigquery/impl.py | 30 ++++++++-------- dbt/adapters/bigquery/relation.py | 4 +-- .../bigquery/relation_configs/_partition.py | 8 ++--- dbt/adapters/bigquery/utility.py | 4 +-- setup.py | 3 +- .../adapter/column_types/fixtures.py | 4 +-- tests/functional/adapter/test_aliases.py | 4 +-- tests/unit/test_bigquery_adapter.py | 36 ++++++++++--------- .../unit/test_bigquery_connection_manager.py | 7 ++-- tests/unit/utils.py | 4 +-- 13 files changed, 73 insertions(+), 62 deletions(-) create mode 100644 .changes/unreleased/Under the Hood-20240116-154305.yaml diff --git a/.changes/unreleased/Under the Hood-20240116-154305.yaml b/.changes/unreleased/Under the Hood-20240116-154305.yaml new file mode 100644 index 000000000..bb115abd6 --- /dev/null +++ b/.changes/unreleased/Under the Hood-20240116-154305.yaml @@ -0,0 +1,6 @@ +kind: Under the Hood +body: Migrate to dbt-common and dbt-adapters package +time: 2024-01-16T15:43:05.046735-08:00 +custom: + Author: colin-rogers-dbt + Issue: "1071" diff --git a/dbt/adapters/bigquery/connections.py b/dbt/adapters/bigquery/connections.py index cb933baed..c74effcdc 100644 --- a/dbt/adapters/bigquery/connections.py +++ b/dbt/adapters/bigquery/connections.py @@ -5,9 +5,9 @@ from contextlib import contextmanager from dataclasses import dataclass, field -from dbt.common.invocation import get_invocation_id +from dbt_common.invocation import get_invocation_id -from dbt.common.events.contextvars import get_node_info +from dbt_common.events.contextvars import get_node_info from mashumaro.helper import pass_through from functools import lru_cache @@ -27,21 +27,21 @@ ) from dbt.adapters.bigquery import gcloud -from dbt.common.clients import agate_helper -from dbt.adapters.contracts.connection import ConnectionState, AdapterResponse -from dbt.common.exceptions import ( +from dbt_common.clients import agate_helper +from dbt.adapters.contracts.connection import ConnectionState, AdapterResponse, Credentials +from dbt_common.exceptions import ( DbtRuntimeError, DbtConfigError, ) -from dbt.common.exceptions import DbtDatabaseError +from dbt_common.exceptions import DbtDatabaseError from dbt.adapters.exceptions.connection import FailedToConnectError -from dbt.adapters.base import BaseConnectionManager, Credentials +from dbt.adapters.base import BaseConnectionManager from dbt.adapters.events.logging import AdapterLogger from dbt.adapters.events.types import SQLQuery -from dbt.common.events.functions import fire_event +from dbt_common.events.functions import fire_event from dbt.adapters.bigquery import __version__ as dbt_version -from dbt.common.dataclass_schema import ExtensibleDbtClassMixin, StrEnum +from dbt_common.dataclass_schema import ExtensibleDbtClassMixin, StrEnum logger = AdapterLogger("BigQuery") diff --git a/dbt/adapters/bigquery/gcloud.py b/dbt/adapters/bigquery/gcloud.py index 0a08f734d..ea1f644ba 100644 --- a/dbt/adapters/bigquery/gcloud.py +++ b/dbt/adapters/bigquery/gcloud.py @@ -1,6 +1,7 @@ +from dbt_common.exceptions import DbtRuntimeError + from dbt.adapters.events.logging import AdapterLogger -import dbt.common.exceptions -from dbt.common.clients.system import run_cmd +from dbt_common.clients.system import run_cmd NOT_INSTALLED_MSG = """ dbt requires the gcloud SDK to be installed to authenticate with BigQuery. @@ -25,4 +26,4 @@ def setup_default_credentials(): if gcloud_installed(): run_cmd(".", ["gcloud", "auth", "application-default", "login"]) else: - raise dbt.common.exceptions.DbtRuntimeError(NOT_INSTALLED_MSG) + raise DbtRuntimeError(NOT_INSTALLED_MSG) diff --git a/dbt/adapters/bigquery/impl.py b/dbt/adapters/bigquery/impl.py index 03cfd3561..3ca7435c4 100644 --- a/dbt/adapters/bigquery/impl.py +++ b/dbt/adapters/bigquery/impl.py @@ -9,7 +9,7 @@ import agate from dbt.adapters.contracts.relation import RelationConfig -import dbt.common.exceptions.base +import dbt_common.exceptions.base from dbt.adapters.base import ( # type: ignore AdapterConfig, BaseAdapter, @@ -21,15 +21,15 @@ available, ) from dbt.adapters.cache import _make_ref_key_dict # type: ignore -import dbt.common.clients.agate_helper +import dbt_common.clients.agate_helper from dbt.adapters.contracts.connection import AdapterResponse -from dbt.common.contracts.constraints import ColumnLevelConstraint, ConstraintType, ModelLevelConstraint # type: ignore -from dbt.common.dataclass_schema import dbtClassMixin +from dbt_common.contracts.constraints import ColumnLevelConstraint, ConstraintType, ModelLevelConstraint # type: ignore +from dbt_common.dataclass_schema import dbtClassMixin from dbt.adapters.events.logging import AdapterLogger -from dbt.common.events.functions import fire_event +from dbt_common.events.functions import fire_event from dbt.adapters.events.types import SchemaCreation, SchemaDrop -import dbt.common.exceptions -from dbt.common.utils import filter_null_values +import dbt_common.exceptions +from dbt_common.utils import filter_null_values import google.api_core import google.auth import google.oauth2 @@ -147,7 +147,7 @@ def drop_relation(self, relation: BigQueryRelation) -> None: conn.handle.delete_table(table_ref, not_found_ok=True) def truncate_relation(self, relation: BigQueryRelation) -> None: - raise dbt.common.exceptions.base.NotImplementedError( + raise dbt_common.exceptions.base.NotImplementedError( "`truncate` is not implemented for this adapter!" ) @@ -164,7 +164,7 @@ def rename_relation( or from_relation.type == RelationType.View or to_relation.type == RelationType.View ): - raise dbt.common.exceptions.DbtRuntimeError( + raise dbt_common.exceptions.DbtRuntimeError( "Renaming of views is not currently supported in BigQuery" ) @@ -390,7 +390,7 @@ def copy_table(self, source, destination, materialization): elif materialization == "table": write_disposition = WRITE_TRUNCATE else: - raise dbt.common.exceptions.CompilationError( + raise dbt_common.exceptions.CompilationError( 'Copy table materialization must be "copy" or "table", but ' f"config.get('copy_materialization', 'table') was " f"{materialization}" @@ -437,11 +437,11 @@ def poll_until_job_completes(cls, job, timeout): job.reload() if job.state != "DONE": - raise dbt.common.exceptions.DbtRuntimeError("BigQuery Timeout Exceeded") + raise dbt_common.exceptions.DbtRuntimeError("BigQuery Timeout Exceeded") elif job.error_result: message = "\n".join(error["message"].strip() for error in job.errors) - raise dbt.common.exceptions.DbtRuntimeError(message) + raise dbt_common.exceptions.DbtRuntimeError(message) def _bq_table_to_relation(self, bq_table) -> Union[BigQueryRelation, None]: if bq_table is None: @@ -465,7 +465,7 @@ def add_query(self, sql, auto_begin=True, bindings=None, abridge_sql_log=False): if self.nice_connection_name() in ["on-run-start", "on-run-end"]: self.warning_on_hooks(self.nice_connection_name()) else: - raise dbt.common.exceptions.base.NotImplementedError( + raise dbt_common.exceptions.base.NotImplementedError( "`add_query` is not implemented for this adapter!" ) @@ -777,7 +777,7 @@ def describe_relation( bq_table = self.get_bq_table(relation) parser = BigQueryMaterializedViewConfig else: - raise dbt.common.exceptions.DbtRuntimeError( + raise dbt_common.exceptions.DbtRuntimeError( f"The method `BigQueryAdapter.describe_relation` is not implemented " f"for the relation type: {relation.type}" ) @@ -843,7 +843,7 @@ def string_add_sql( elif location == "prepend": return f"concat('{value}', {add_to})" else: - raise dbt.common.exceptions.DbtRuntimeError( + raise dbt_common.exceptions.DbtRuntimeError( f'Got an unexpected location value of "{location}"' ) diff --git a/dbt/adapters/bigquery/relation.py b/dbt/adapters/bigquery/relation.py index c25ef0a67..8abda577b 100644 --- a/dbt/adapters/bigquery/relation.py +++ b/dbt/adapters/bigquery/relation.py @@ -12,8 +12,8 @@ BigQueryPartitionConfigChange, ) from dbt.adapters.contracts.relation import RelationType, RelationConfig -from dbt.common.exceptions import CompilationError -from dbt.common.utils.dict import filter_null_values +from dbt_common.exceptions import CompilationError +from dbt_common.utils.dict import filter_null_values Self = TypeVar("Self", bound="BigQueryRelation") diff --git a/dbt/adapters/bigquery/relation_configs/_partition.py b/dbt/adapters/bigquery/relation_configs/_partition.py index 0fe816359..8fe8bf5d6 100644 --- a/dbt/adapters/bigquery/relation_configs/_partition.py +++ b/dbt/adapters/bigquery/relation_configs/_partition.py @@ -1,10 +1,10 @@ from dataclasses import dataclass from typing import Any, Dict, List, Optional -import dbt.common.exceptions +import dbt_common.exceptions from dbt.adapters.relation_configs import RelationConfigChange from dbt.adapters.contracts.relation import RelationConfig -from dbt.common.dataclass_schema import dbtClassMixin, ValidationError +from dbt_common.dataclass_schema import dbtClassMixin, ValidationError from google.cloud.bigquery.table import Table as BigQueryTable @@ -92,11 +92,11 @@ def parse(cls, raw_partition_by) -> Optional["PartitionConfig"]: } ) except ValidationError as exc: - raise dbt.common.exceptions.base.DbtValidationError( + raise dbt_common.exceptions.base.DbtValidationError( "Could not parse partition config" ) from exc except TypeError: - raise dbt.common.exceptions.CompilationError( + raise dbt_common.exceptions.CompilationError( f"Invalid partition_by config:\n" f" Got: {raw_partition_by}\n" f' Expected a dictionary with "field" and "data_type" keys' diff --git a/dbt/adapters/bigquery/utility.py b/dbt/adapters/bigquery/utility.py index 5d9c3de12..5914280a3 100644 --- a/dbt/adapters/bigquery/utility.py +++ b/dbt/adapters/bigquery/utility.py @@ -1,7 +1,7 @@ import json from typing import Any, Optional -import dbt.common.exceptions +import dbt_common.exceptions def bool_setting(value: Optional[Any] = None) -> Optional[bool]: @@ -41,5 +41,5 @@ def float_setting(value: Optional[Any] = None) -> Optional[float]: def sql_escape(string): if not isinstance(string, str): - raise dbt.common.exceptions.CompilationError(f"cannot escape a non-string: {string}") + raise dbt_common.exceptions.CompilationError(f"cannot escape a non-string: {string}") return json.dumps(string)[1:-1] diff --git a/setup.py b/setup.py index 56c271fec..b3a1065cb 100644 --- a/setup.py +++ b/setup.py @@ -74,7 +74,8 @@ def _dbt_core_version(plugin_version: str) -> str: packages=find_namespace_packages(include=["dbt", "dbt.*"]), include_package_data=True, install_requires=[ - f"dbt-core~={_dbt_core_version(_dbt_bigquery_version())}", + "dbt-common<1.0", + "dbt-adapters~=0.1.0a1", "google-cloud-bigquery~=3.0", "google-cloud-storage~=2.4", "google-cloud-dataproc~=5.0", diff --git a/tests/functional/adapter/column_types/fixtures.py b/tests/functional/adapter/column_types/fixtures.py index b7be1e646..88175a88b 100644 --- a/tests/functional/adapter/column_types/fixtures.py +++ b/tests/functional/adapter/column_types/fixtures.py @@ -26,7 +26,7 @@ version: 2 models: - name: model - tests: + data_tests: - is_type: column_map: int64_col: ['integer', 'number'] @@ -39,7 +39,7 @@ version: 2 models: - name: model - tests: + data_tests: - is_type: column_map: int64_col: ['string', 'not number'] diff --git a/tests/functional/adapter/test_aliases.py b/tests/functional/adapter/test_aliases.py index fa28ce5d9..5ce13c8bc 100644 --- a/tests/functional/adapter/test_aliases.py +++ b/tests/functional/adapter/test_aliases.py @@ -32,12 +32,12 @@ version: 2 models: - name: model_a - tests: + data_tests: - expect_value: field: tablename value: duped_alias - name: model_b - tests: + data_tests: - expect_value: field: tablename value: duped_alias diff --git a/tests/unit/test_bigquery_adapter.py b/tests/unit/test_bigquery_adapter.py index a324c431a..34abd0caf 100644 --- a/tests/unit/test_bigquery_adapter.py +++ b/tests/unit/test_bigquery_adapter.py @@ -10,14 +10,16 @@ import unittest from unittest.mock import patch, MagicMock, create_autospec -import dbt.common.dataclass_schema -import dbt.common.exceptions.base +import dbt_common.dataclass_schema +import dbt_common.exceptions.base + +import dbt.adapters from dbt.adapters.bigquery.relation_configs import PartitionConfig from dbt.adapters.bigquery import BigQueryAdapter, BigQueryRelation from google.cloud.bigquery.table import Table from dbt.adapters.bigquery.connections import _sanitize_label, _VALIDATE_LABEL_LENGTH_LIMIT -from dbt.common.clients import agate_helper -import dbt.common.exceptions +from dbt_common.clients import agate_helper +import dbt_common.exceptions from dbt.context.manifest import generate_query_header_context from dbt.contracts.files import FileHash from dbt.contracts.graph.manifest import ManifestStateCheck @@ -214,7 +216,7 @@ def test_acquire_connection_oauth_no_project_validations( connection = adapter.acquire_connection("dummy") self.assertEqual(connection.type, "bigquery") - except dbt.common.exceptions.base.DbtValidationError as e: + except dbt_common.exceptions.base.DbtValidationError as e: self.fail("got DbtValidationError: {}".format(str(e))) except BaseException: @@ -231,7 +233,7 @@ def test_acquire_connection_oauth_validations(self, mock_open_connection): connection = adapter.acquire_connection("dummy") self.assertEqual(connection.type, "bigquery") - except dbt.common.exceptions.base.DbtValidationError as e: + except dbt_common.exceptions.base.DbtValidationError as e: self.fail("got DbtValidationError: {}".format(str(e))) except BaseException: @@ -255,7 +257,7 @@ def test_acquire_connection_dataproc_serverless( connection = adapter.acquire_connection("dummy") self.assertEqual(connection.type, "bigquery") - except dbt.common.exceptions.ValidationException as e: + except dbt_common.exceptions.ValidationException as e: self.fail("got ValidationException: {}".format(str(e))) except BaseException: @@ -272,7 +274,7 @@ def test_acquire_connection_service_account_validations(self, mock_open_connecti connection = adapter.acquire_connection("dummy") self.assertEqual(connection.type, "bigquery") - except dbt.common.exceptions.base.DbtValidationError as e: + except dbt_common.exceptions.base.DbtValidationError as e: self.fail("got DbtValidationError: {}".format(str(e))) except BaseException: @@ -289,7 +291,7 @@ def test_acquire_connection_oauth_token_validations(self, mock_open_connection): connection = adapter.acquire_connection("dummy") self.assertEqual(connection.type, "bigquery") - except dbt.common.exceptions.base.DbtValidationError as e: + except dbt_common.exceptions.base.DbtValidationError as e: self.fail("got DbtValidationError: {}".format(str(e))) except BaseException: @@ -306,7 +308,7 @@ def test_acquire_connection_oauth_credentials_validations(self, mock_open_connec connection = adapter.acquire_connection("dummy") self.assertEqual(connection.type, "bigquery") - except dbt.common.exceptions.base.DbtValidationError as e: + except dbt_common.exceptions.base.DbtValidationError as e: self.fail("got DbtValidationError: {}".format(str(e))) except BaseException: @@ -325,7 +327,7 @@ def test_acquire_connection_impersonated_service_account_validations( connection = adapter.acquire_connection("dummy") self.assertEqual(connection.type, "bigquery") - except dbt.common.exceptions.base.DbtValidationError as e: + except dbt_common.exceptions.base.DbtValidationError as e: self.fail("got DbtValidationError: {}".format(str(e))) except BaseException: @@ -343,7 +345,7 @@ def test_acquire_connection_priority(self, mock_open_connection): self.assertEqual(connection.type, "bigquery") self.assertEqual(connection.credentials.priority, "batch") - except dbt.common.exceptions.base.DbtValidationError as e: + except dbt_common.exceptions.base.DbtValidationError as e: self.fail("got DbtValidationError: {}".format(str(e))) mock_open_connection.assert_not_called() @@ -358,7 +360,7 @@ def test_acquire_connection_maximum_bytes_billed(self, mock_open_connection): self.assertEqual(connection.type, "bigquery") self.assertEqual(connection.credentials.maximum_bytes_billed, 0) - except dbt.common.exceptions.base.DbtValidationError as e: + except dbt_common.exceptions.base.DbtValidationError as e: self.fail("got DbtValidationError: {}".format(str(e))) mock_open_connection.assert_not_called() @@ -509,7 +511,7 @@ def test_invalid_relation(self): }, "quote_policy": {"identifier": False, "schema": True}, } - with self.assertRaises(dbt.common.dataclass_schema.ValidationError): + with self.assertRaises(dbt_common.dataclass_schema.ValidationError): BigQueryRelation.validate(kwargs) @@ -581,10 +583,10 @@ def test_copy_table_materialization_incremental(self): def test_parse_partition_by(self): adapter = self.get_adapter("oauth") - with self.assertRaises(dbt.common.exceptions.base.DbtValidationError): + with self.assertRaises(dbt_common.exceptions.base.DbtValidationError): adapter.parse_partition_by("date(ts)") - with self.assertRaises(dbt.common.exceptions.base.DbtValidationError): + with self.assertRaises(dbt_common.exceptions.base.DbtValidationError): adapter.parse_partition_by("ts") self.assertEqual( @@ -736,7 +738,7 @@ def test_parse_partition_by(self): ) # Invalid, should raise an error - with self.assertRaises(dbt.common.exceptions.base.DbtValidationError): + with self.assertRaises(dbt_common.exceptions.base.DbtValidationError): adapter.parse_partition_by({}) # passthrough diff --git a/tests/unit/test_bigquery_connection_manager.py b/tests/unit/test_bigquery_connection_manager.py index 04e6d1352..6bb89ed36 100644 --- a/tests/unit/test_bigquery_connection_manager.py +++ b/tests/unit/test_bigquery_connection_manager.py @@ -6,12 +6,13 @@ from requests.exceptions import ConnectionError from unittest.mock import patch, MagicMock, Mock, ANY -import dbt.common.dataclass_schema +import dbt.adapters +import dbt_common.dataclass_schema from dbt.adapters.bigquery import BigQueryCredentials from dbt.adapters.bigquery import BigQueryRelation from dbt.adapters.bigquery.connections import BigQueryConnectionManager -import dbt.common.exceptions +import dbt_common.exceptions from dbt.logger import GLOBAL_LOGGER as logger # noqa @@ -127,7 +128,7 @@ 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.common.exceptions.DbtRuntimeError) as exc: + with pytest.raises(dbt_common.exceptions.DbtRuntimeError) as exc: self.connections._query_and_results( self.mock_client, "sql", diff --git a/tests/unit/utils.py b/tests/unit/utils.py index 6d21828b3..88b09ce60 100644 --- a/tests/unit/utils.py +++ b/tests/unit/utils.py @@ -10,7 +10,7 @@ import agate import pytest -from dbt.common.dataclass_schema import ValidationError +from dbt_common.dataclass_schema import ValidationError from dbt.config.project import PartialProject @@ -256,7 +256,7 @@ def generate_name_macros(package): class TestAdapterConversions(TestCase): def _get_tester_for(self, column_type): - from dbt.common.clients import agate_helper + from dbt_common.clients import agate_helper if column_type is agate.TimeDelta: # dbt never makes this! return agate.TimeDelta()