From a5a02b4ce83140fb702040aa2474813803554b4f Mon Sep 17 00:00:00 2001 From: Mike Alfare Date: Wed, 23 Aug 2023 19:43:07 -0400 Subject: [PATCH] updated typing and methods to meet mypy standards --- .../Under the Hood-20230823-194237.yaml | 6 +++++ core/dbt/adapters/base/connections.py | 23 +++++++++++++++++- core/dbt/adapters/base/impl.py | 24 ++++++++++++------- core/dbt/adapters/sql/connections.py | 10 +------- 4 files changed, 45 insertions(+), 18 deletions(-) create mode 100644 .changes/unreleased/Under the Hood-20230823-194237.yaml diff --git a/.changes/unreleased/Under the Hood-20230823-194237.yaml b/.changes/unreleased/Under the Hood-20230823-194237.yaml new file mode 100644 index 00000000000..7bd77acab08 --- /dev/null +++ b/.changes/unreleased/Under the Hood-20230823-194237.yaml @@ -0,0 +1,6 @@ +kind: Under the Hood +body: Update typing to meet mypy standards +time: 2023-08-23T19:42:37.130694-04:00 +custom: + Author: mikealfare + Issue: "8396" diff --git a/core/dbt/adapters/base/connections.py b/core/dbt/adapters/base/connections.py index d449b27e5e6..23a8b6e6588 100644 --- a/core/dbt/adapters/base/connections.py +++ b/core/dbt/adapters/base/connections.py @@ -400,7 +400,7 @@ def _add_query_comment(self, sql: str) -> str: @abc.abstractmethod def execute( - self, sql: str, auto_begin: bool = False, fetch: bool = False + self, sql: str, auto_begin: bool = False, fetch: bool = False, limit: Optional[int] = None ) -> Tuple[AdapterResponse, agate.Table]: """Execute the given SQL. @@ -408,7 +408,28 @@ def execute( :param bool auto_begin: If set, and dbt is not currently inside a transaction, automatically begin one. :param bool fetch: If set, fetch results. + :param int limit: If set, limits the result set :return: A tuple of the query status and results (empty if fetch=False). :rtype: Tuple[AdapterResponse, agate.Table] """ raise dbt.exceptions.NotImplementedError("`execute` is not implemented for this adapter!") + + def add_select_query(self, sql: str) -> Tuple[Connection, Any]: + """ + This was added here because base.impl.BaseAdapter.get_column_schema_from_query expects it to be here. + That method wouldn't work unless the adapter used sql.impl.SQLAdapter, sql.connections.SQLConnectionManager + or defined this method on ConnectionManager before passing it in to Adapter. + + See https://github.com/dbt-labs/dbt-core/issues/8396 for more information. + """ + raise dbt.exceptions.NotImplementedError( + "`add_select_query` is not implemented for this adapter!" + ) + + @classmethod + def data_type_code_to_name(cls, type_code: Union[int, str]) -> str: + """Get the string representation of the data type from the type_code.""" + # https://peps.python.org/pep-0249/#type-objects + raise dbt.exceptions.NotImplementedError( + "`data_type_code_to_name` is not implemented for this adapter!" + ) diff --git a/core/dbt/adapters/base/impl.py b/core/dbt/adapters/base/impl.py index d18c9af7f50..dd147dce845 100644 --- a/core/dbt/adapters/base/impl.py +++ b/core/dbt/adapters/base/impl.py @@ -43,7 +43,7 @@ UnexpectedNullError, ) -from dbt.adapters.protocol import AdapterConfig, ConnectionManagerProtocol +from dbt.adapters.protocol import AdapterConfig from dbt.clients.agate_helper import empty_table, merge_tables, table_from_rows from dbt.clients.jinja import MacroGenerator from dbt.contracts.graph.manifest import Manifest, MacroManifest @@ -60,7 +60,7 @@ ) from dbt.utils import filter_null_values, executor, cast_to_str, AttrDict -from dbt.adapters.base.connections import Connection, AdapterResponse +from dbt.adapters.base.connections import Connection, AdapterResponse, BaseConnectionManager from dbt.adapters.base.meta import AdapterMeta, available from dbt.adapters.base.relation import ( ComponentName, @@ -208,7 +208,7 @@ class BaseAdapter(metaclass=AdapterMeta): Relation: Type[BaseRelation] = BaseRelation Column: Type[BaseColumn] = BaseColumn - ConnectionManager: Type[ConnectionManagerProtocol] + ConnectionManager: Type[BaseConnectionManager] # A set of clobber config fields accepted by this adapter # for use in materializations @@ -315,14 +315,21 @@ def get_column_schema_from_query(self, sql: str) -> List[BaseColumn]: @available.parse(lambda *a, **k: ("", empty_table())) def get_partitions_metadata(self, table: str) -> Tuple[agate.Table]: - """Obtain partitions metadata for a BigQuery partitioned table. + """ + TODO: Can we move this to dbt-bigquery? + Obtain partitions metadata for a BigQuery partitioned table. - :param str table_id: a partitioned table id, in standard SQL format. + :param str table: a partitioned table id, in standard SQL format. :return: a partition metadata tuple, as described in https://cloud.google.com/bigquery/docs/creating-partitioned-tables#getting_partition_metadata_using_meta_tables. :rtype: agate.Table """ - return self.connections.get_partitions_metadata(table=table) + if hasattr(self.connections, "get_partitions_metadata"): + return self.connections.get_partitions_metadata(table=table) + else: + raise NotImplementedError( + "`get_partitions_metadata` is not implemented for this adapter!" + ) ### # Methods that should never be overridden @@ -453,9 +460,10 @@ def _relations_cache_for_schemas( # it's possible that there were no relations in some schemas. We want # to insert the schemas we query into the cache's `.schemas` attribute # so we can check it later - cache_update: Set[Tuple[Optional[str], Optional[str]]] = set() + cache_update: Set[Tuple[Optional[str], str]] = set() for relation in cache_schemas: - cache_update.add((relation.database, relation.schema)) + if relation.schema: + cache_update.add((relation.database, relation.schema)) self.cache.update_schemas(cache_update) def set_relations_cache( diff --git a/core/dbt/adapters/sql/connections.py b/core/dbt/adapters/sql/connections.py index 43463d3a47a..7347d961d15 100644 --- a/core/dbt/adapters/sql/connections.py +++ b/core/dbt/adapters/sql/connections.py @@ -1,6 +1,6 @@ import abc import time -from typing import List, Optional, Tuple, Any, Iterable, Dict, Union +from typing import List, Optional, Tuple, Any, Iterable, Dict import agate @@ -131,14 +131,6 @@ def get_result_from_cursor(cls, cursor: Any, limit: Optional[int]) -> agate.Tabl return dbt.clients.agate_helper.table_from_data_flat(data, column_names) - @classmethod - def data_type_code_to_name(cls, type_code: Union[int, str]) -> str: - """Get the string representation of the data type from the type_code.""" - # https://peps.python.org/pep-0249/#type-objects - raise dbt.exceptions.NotImplementedError( - "`data_type_code_to_name` is not implemented for this adapter!" - ) - def execute( self, sql: str, auto_begin: bool = False, fetch: bool = False, limit: Optional[int] = None ) -> Tuple[AdapterResponse, agate.Table]: