dbt-adapters 1.22.2__py3-none-any.whl

dbt/adapters/protocol.py

@@ -0,0 +1,150 @@
+from dataclasses import dataclass
+from typing import (
+    Any,
+    ContextManager,
+    Dict,
+    Generic,
+    Hashable,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Tuple,
+    TYPE_CHECKING,
+)
+from typing_extensions import Protocol
+
+from dbt_common.clients.jinja import MacroProtocol
+from dbt_common.contracts.config.base import BaseConfig
+
+from dbt.adapters.contracts.connection import (
+    AdapterRequiredConfig,
+    AdapterResponse,
+    Connection,
+)
+from dbt.adapters.contracts.macros import MacroResolverProtocol
+from dbt.adapters.contracts.relation import HasQuoting, Policy, RelationConfig
+
+if TYPE_CHECKING:
+    import agate
+
+
+@dataclass
+class AdapterConfig(BaseConfig):
+    pass
+
+
+class ConnectionManagerProtocol(Protocol):
+    TYPE: str
+
+
+class ColumnProtocol(Protocol):
+    pass
+
+
+Self = TypeVar("Self", bound="RelationProtocol")
+
+
+class RelationProtocol(Protocol):
+    @classmethod
+    def get_default_quote_policy(cls) -> Policy: ...
+
+    @classmethod
+    def create_from(
+        cls: Type[Self],
+        quoting: HasQuoting,
+        relation_config: RelationConfig,
+        **kwargs: Any,
+    ) -> Self: ...
+
+
+AdapterConfig_T = TypeVar("AdapterConfig_T", bound=AdapterConfig)
+ConnectionManager_T = TypeVar("ConnectionManager_T", bound=ConnectionManagerProtocol)
+Relation_T = TypeVar("Relation_T", bound=RelationProtocol)
+Column_T = TypeVar("Column_T", bound=ColumnProtocol)
+
+
+class MacroContextGeneratorCallable(Protocol):
+    def __call__(
+        self,
+        macro_protocol: MacroProtocol,
+        config: AdapterRequiredConfig,
+        macro_resolver: MacroResolverProtocol,
+        package_name: Optional[str],
+    ) -> Dict[str, Any]: ...
+
+
+# TODO CT-211
+class AdapterProtocol(
+    Protocol,
+    Generic[
+        AdapterConfig_T,
+        ConnectionManager_T,
+        Relation_T,
+        Column_T,
+    ],
+):
+    # N.B. Technically these are ClassVars, but mypy doesn't support putting type vars in a
+    # ClassVar due to the restrictiveness of PEP-526
+    # See: https://github.com/python/mypy/issues/5144
+    AdapterSpecificConfigs: Type[AdapterConfig_T]
+    Column: Type[Column_T]
+    Relation: Type[Relation_T]
+    ConnectionManager: Type[ConnectionManager_T]
+    connections: ConnectionManager_T
+
+    def __init__(self, config: AdapterRequiredConfig) -> None: ...
+
+    def set_macro_resolver(self, macro_resolver: MacroResolverProtocol) -> None: ...
+
+    def get_macro_resolver(self) -> Optional[MacroResolverProtocol]: ...
+
+    def clear_macro_resolver(self) -> None: ...
+
+    def set_macro_context_generator(
+        self,
+        macro_context_generator: MacroContextGeneratorCallable,
+    ) -> None: ...
+
+    @classmethod
+    def type(cls) -> str:
+        pass
+
+    def set_query_header(self, query_header_context: Dict[str, Any]) -> None: ...
+
+    @staticmethod
+    def get_thread_identifier() -> Hashable: ...
+
+    def get_thread_connection(self) -> Connection: ...
+
+    def set_thread_connection(self, conn: Connection) -> None: ...
+
+    def get_if_exists(self) -> Optional[Connection]: ...
+
+    def clear_thread_connection(self) -> None: ...
+
+    def clear_transaction(self) -> None: ...
+
+    def exception_handler(self, sql: str) -> ContextManager: ...
+
+    def set_connection_name(self, name: Optional[str] = None) -> Connection: ...
+
+    def cancel_open(self) -> Optional[List[str]]: ...
+
+    def open(cls, connection: Connection) -> Connection: ...
+
+    def release(self) -> None: ...
+
+    def cleanup_all(self) -> None: ...
+
+    def begin(self) -> None: ...
+
+    def commit(self) -> None: ...
+
+    def close(cls, connection: Connection) -> Connection: ...
+
+    def commit_if_has_connection(self) -> None: ...
+
+    def execute(
+        self, sql: str, auto_begin: bool = False, fetch: bool = False
+    ) -> Tuple[AdapterResponse, "agate.Table"]: ...

dbt/adapters/record/__init__.py

@@ -0,0 +1,2 @@
+from dbt.adapters.record.handle import RecordReplayHandle
+from dbt.adapters.record.cursor.cursor import RecordReplayCursor

dbt/adapters/record/base.py

@@ -0,0 +1,291 @@
+"""Implementations of record/replay classes for the base adapter implementation."""
+
+import dataclasses
+
+from typing import Optional, Tuple, Dict, Any, TYPE_CHECKING, List
+
+from dbt.adapters.contracts.connection import AdapterResponse
+from dbt.adapters.record.serialization import serialize_agate_table, serialize_bindings
+from dbt_common.record import Record, Recorder
+
+if TYPE_CHECKING:
+    from agate import Table
+    from dbt.adapters.base.relation import BaseRelation
+    from dbt.adapters.base.column import Column as BaseColumn
+
+
+@dataclasses.dataclass
+class AdapterExecuteParams:
+    thread_id: str
+    sql: str
+    auto_begin: bool = False
+    fetch: bool = False
+    limit: Optional[int] = None
+
+
+@dataclasses.dataclass
+class AdapterExecuteResult:
+    return_val: Tuple[AdapterResponse, "Table"]
+
+    def _to_dict(self):
+        adapter_response = self.return_val[0]
+        table = self.return_val[1]
+        return {
+            "return_val": {
+                "adapter_response": adapter_response.to_dict(),
+                "table": serialize_agate_table(table),
+            }
+        }
+
+    def _from_dict(self, data: Dict[str, Any]):
+        # We will need this for replay, but it is not a priority at time of writing.
+        raise NotImplementedError()
+
+
+@Recorder.register_record_type
+class AdapterExecuteRecord(Record):
+    """Implements record/replay support for the BaseAdapter.execute() method."""
+
+    params_cls = AdapterExecuteParams
+    result_cls = AdapterExecuteResult
+    group = "Available"
+
+
+@dataclasses.dataclass
+class AdapterTestSqlResult:
+    return_val: str
+
+
+@dataclasses.dataclass
+class AdapterTestSqlParams:
+    thread_id: str
+    sql: str
+    fetch: str
+    conn: Any
+
+    def _to_dict(self):
+        return {
+            "thread_id": self.thread_id,
+            "sql": self.sql,
+            "fetch": self.fetch,
+            "conn": "conn",
+        }
+
+
+@Recorder.register_record_type
+class AdapterTestSqlRecord(Record):
+    """Implements record/replay support for the BaseAdapter.execute() method."""
+
+    params_cls = AdapterTestSqlParams
+    result_cls = AdapterTestSqlResult
+    group = "Available"
+
+
+@dataclasses.dataclass
+class AdapterGetPartitionsMetadataParams:
+    thread_id: str
+    table: str
+
+
+@dataclasses.dataclass
+class AdapterGetPartitionsMetadataResult:
+    return_val: tuple["Table"]
+
+    def _to_dict(self):
+        return list(map(serialize_agate_table, self.return_val))
+
+    def _from_dict(self, data: Dict[str, Any]):
+        # We will need this for replay, but it is not a priority at time of writing.
+        raise NotImplementedError()
+
+
+@Recorder.register_record_type
+class AdapterGetPartitionsMetadataRecord(Record):
+    """Implements record/replay support for the BaseAdapter.get_partitions_metadata() method."""
+
+    params_cls = AdapterGetPartitionsMetadataParams
+    result_cls = AdapterGetPartitionsMetadataResult
+    group = "Available"
+
+
+@dataclasses.dataclass
+class AdapterConvertTypeParams:
+    thread_id: str
+    table: "Table"
+    col_idx: int
+
+    def _to_dict(self):
+        return {
+            "thread_id": self.thread_id,
+            "table": serialize_agate_table(self.table),
+            "col_idx": self.col_idx,
+        }
+
+    def _from_dict(self, data: Dict[str, Any]):
+        # We will need this for replay, but it is not a priority at time of writing.
+        raise NotImplementedError()
+
+
+@dataclasses.dataclass
+class AdapterConvertTypeResult:
+    return_val: Optional[str]
+
+
+@Recorder.register_record_type
+class AdapterConvertTypeRecord(Record):
+    """Implements record/replay support for the BaseAdapter.convert_type() method."""
+
+    params_cls = AdapterConvertTypeParams
+    result_cls = AdapterConvertTypeResult
+    group = "Available"
+
+
+@dataclasses.dataclass
+class AdapterStandardizeGrantsDictParams:
+    thread_id: str
+    table: "Table"
+
+    def _to_dict(self):
+        return {"thread_id": self.thread_id, "table": serialize_agate_table(self.table)}
+
+    def _from_dict(self, data: Dict[str, Any]):
+        # We will need this for replay, but it is not a priority at time of writing.
+        raise NotImplementedError()
+
+
+@dataclasses.dataclass
+class AdapterStandardizeGrantsDictResult:
+    return_val: dict
+
+
+@Recorder.register_record_type
+class AdapterStandardizeGrantsDictRecord(Record):
+    params_cls = AdapterStandardizeGrantsDictParams
+    result_cls = AdapterStandardizeGrantsDictResult
+    group = "Available"
+
+
+@dataclasses.dataclass
+class AdapterAddQueryParams:
+    thread_id: str
+    sql: str
+    auto_begin: bool = True
+    bindings: Optional[Any] = None
+    abridge_sql_log: bool = False
+
+    def _to_dict(self):
+        return {
+            "thread_id": self.thread_id,
+            "sql": self.sql,
+            "auto_begin": self.auto_begin,
+            "bindings": serialize_bindings(self.bindings),
+            "abridge_sql_log": self.abridge_sql_log,
+        }
+
+
+@dataclasses.dataclass
+class AdapterAddQueryResult:
+    return_val: Tuple[str, str]
+
+    def _to_dict(self):
+        return {
+            "return_val": {
+                "conn": "conn",
+                "cursor": "cursor",
+            }
+        }
+
+
+@Recorder.register_record_type
+class AdapterAddQueryRecord(Record):
+    params_cls = AdapterAddQueryParams
+    result_cls = AdapterAddQueryResult
+    group = "Available"
+
+
+@dataclasses.dataclass
+class AdapterListRelationsWithoutCachingParams:
+    thread_id: str
+    schema_relation: "BaseRelation"
+
+    def _to_dict(self):
+        from dbt.adapters.record.serialization import serialize_base_relation
+
+        return {
+            "thread_id": self.thread_id,
+            "schema_relation": serialize_base_relation(self.schema_relation),
+        }
+
+    def _from_dict(self, data: Dict[str, Any]):
+        from dbt.adapters.record.serialization import deserialize_base_relation
+
+        self.thread_id = data["thread_id"]
+        self.schema_relation = deserialize_base_relation(data["schema_relation"])
+
+
+@dataclasses.dataclass
+class AdapterListRelationsWithoutCachingResult:
+    return_val: List["BaseRelation"]
+
+    def _to_dict(self):
+        from dbt.adapters.record.serialization import serialize_base_relation_list
+
+        return {"return_val": serialize_base_relation_list(self.return_val)}
+
+    def _from_dict(self, data: Dict[str, Any]):
+        from dbt.adapters.record.serialization import deserialize_base_relation_list
+
+        self.return_val = deserialize_base_relation_list(data["return_val"])
+
+
+@Recorder.register_record_type
+class AdapterListRelationsWithoutCachingRecord(Record):
+    """Implements record/replay support for the BaseAdapter.list_relations_without_caching() method."""
+
+    params_cls = AdapterListRelationsWithoutCachingParams
+    result_cls = AdapterListRelationsWithoutCachingResult
+    group = "Available"
+
+
+@dataclasses.dataclass
+class AdapterGetColumnsInRelationParams:
+    thread_id: str
+    relation: "BaseRelation"
+
+    def _to_dict(self):
+        from dbt.adapters.record.serialization import serialize_base_relation
+
+        return {
+            "thread_id": self.thread_id,
+            "relation": serialize_base_relation(self.relation),
+        }
+
+    def _from_dict(self, data: Dict[str, Any]):
+        from dbt.adapters.record.serialization import deserialize_base_relation
+
+        self.thread_id = data["thread_id"]
+        self.relation = deserialize_base_relation(data["relation"])
+
+
+@dataclasses.dataclass
+class AdapterGetColumnsInRelationResult:
+    return_val: List["BaseColumn"]
+
+    def _to_dict(self):
+        from dbt.adapters.record.serialization import serialize_base_column_list
+
+        return {"return_val": serialize_base_column_list(self.return_val)}
+
+    def _from_dict(self, data: Dict[str, Any]):
+        from dbt.adapters.record.serialization import deserialize_base_column_list
+
+        self.return_val = deserialize_base_column_list(data["return_val"])
+
+
+@Recorder.register_record_type
+class AdapterGetColumnsInRelationRecord(Record):
+    """Implements record/replay support for the BaseAdapter.get_columns_in_relation() method."""
+
+    params_cls = AdapterGetColumnsInRelationParams
+    result_cls = AdapterGetColumnsInRelationResult
+    group = "Available"

dbt/adapters/record/cursor/cursor.py

@@ -0,0 +1,69 @@
+from typing import Any, Optional
+
+from dbt_common.events.base_types import BaseEvent
+from dbt_common.events.functions import fire_event
+from dbt_common.events.types import RecordReplayIssue
+from dbt_common.record import record_function
+
+from dbt.adapters.contracts.connection import Connection
+from dbt.adapters.record.cursor.description import CursorGetDescriptionRecord
+from dbt.adapters.record.cursor.execute import CursorExecuteRecord
+from dbt.adapters.record.cursor.fetchone import CursorFetchOneRecord
+from dbt.adapters.record.cursor.fetchmany import CursorFetchManyRecord
+from dbt.adapters.record.cursor.fetchall import CursorFetchAllRecord
+from dbt.adapters.record.cursor.rowcount import CursorGetRowCountRecord
+
+
+class RecordReplayCursor:
+    """A proxy object used to wrap native database cursors under record/replay
+    modes. In record mode, this proxy notes the parameters and return values
+    of the methods and properties it implements, which closely match the Python
+    DB API 2.0 cursor methods used by many dbt adapters to interact with the
+    database or DWH. In replay mode, it mocks out those calls using previously
+    recorded calls, so that no interaction with a database actually occurs."""
+
+    def __init__(self, native_cursor: Any, connection: Connection) -> None:
+        self.native_cursor = native_cursor
+        self.connection = connection
+
+    @record_function(CursorExecuteRecord, method=True, id_field_name="connection_name")
+    def execute(self, operation, parameters=None) -> None:
+        self.native_cursor.execute(operation, parameters)
+
+    @record_function(CursorFetchOneRecord, method=True, id_field_name="connection_name")
+    def fetchone(self) -> Any:
+        return self.native_cursor.fetchone()
+
+    @record_function(CursorFetchManyRecord, method=True, id_field_name="connection_name")
+    def fetchmany(self, size: int) -> Any:
+        return self.native_cursor.fetchmany(size)
+
+    @record_function(CursorFetchAllRecord, method=True, id_field_name="connection_name")
+    def fetchall(self) -> Any:
+        return self.native_cursor.fetchall()
+
+    @property
+    def connection_name(self) -> Optional[str]:
+        return self.connection.name
+
+    @property
+    @record_function(CursorGetRowCountRecord, method=True, id_field_name="connection_name")
+    def rowcount(self) -> int:
+        return self.native_cursor.rowcount
+
+    @property
+    @record_function(CursorGetDescriptionRecord, method=True, id_field_name="connection_name")
+    def description(self) -> str:
+        return self.native_cursor.description
+
+    def _fire_event(self, evt: BaseEvent) -> None:
+        """Wraps fire_event for easier test mocking."""
+        fire_event(evt)
+
+    def __getattr__(self, name: str) -> Any:
+        self._fire_event(
+            RecordReplayIssue(
+                msg=f"Unexpected attribute '{name}' accessed on {self.__class__.__name__}"
+            )
+        )
+        return getattr(self.native_cursor, name)

dbt/adapters/record/cursor/description.py

@@ -0,0 +1,37 @@
+import dataclasses
+from typing import Any, Iterable, Mapping
+
+from dbt_common.record import Record, Recorder
+
+
+@dataclasses.dataclass
+class CursorGetDescriptionParams:
+    connection_name: str
+
+
+@dataclasses.dataclass
+class CursorGetDescriptionResult:
+    columns: Iterable[Any]
+
+    def _to_dict(self) -> Any:
+        column_dicts = []
+        for c in self.columns:
+            # This captures the mandatory column information, but we might need
+            # more for some adapters.
+            # See https://peps.python.org/pep-0249/#description
+            column_dicts.append((c[0], c[1]))
+
+        return {"columns": column_dicts}
+
+    @classmethod
+    def _from_dict(cls, dct: Mapping) -> "CursorGetDescriptionResult":
+        return CursorGetDescriptionResult(columns=dct["columns"])
+
+
+@Recorder.register_record_type
+class CursorGetDescriptionRecord(Record):
+    """Implements record/replay support for the cursor.description property."""
+
+    params_cls = CursorGetDescriptionParams
+    result_cls = CursorGetDescriptionResult
+    group = "Database"

dbt/adapters/record/cursor/execute.py

@@ -0,0 +1,39 @@
+import dataclasses
+from typing import Any, Iterable, Union, Mapping, Optional
+
+from dbt.adapters.record.cursor.fetchall import CursorFetchAllResult
+from dbt_common.record import Record, Recorder
+
+
+@dataclasses.dataclass
+class CursorExecuteParams:
+    connection_name: str
+    operation: str
+    parameters: Optional[Union[Iterable[Any], Mapping[str, Any]]] = None
+
+    def _to_dict(self):
+        p = self.parameters
+        if isinstance(self.parameters, dict):
+            p = {(k, CursorFetchAllResult._process_value(v)) for k, v in self.parameters.items()}
+        elif isinstance(self.parameters, list) or isinstance(self.parameters, tuple):
+            p = [CursorFetchAllResult._process_value(v) for v in self.parameters]
+
+        return {
+            "connection_name": self.connection_name,
+            "operation": self.operation,
+            "parameters": p,
+        }
+
+    def _from_dict(cls, data):
+        # NOTE: This will be needed for replay, but is not needed at time
+        # of writing.
+        raise NotImplementedError()
+
+
+@Recorder.register_record_type
+class CursorExecuteRecord(Record):
+    """Implements record/replay support for the cursor.execute() method."""
+
+    params_cls = CursorExecuteParams
+    result_cls = None
+    group = "Database"

dbt/adapters/record/cursor/fetchall.py

@@ -0,0 +1,69 @@
+import dataclasses
+import datetime
+import decimal
+from typing import Any, Dict, List, Mapping
+
+from dbt_common.record import Record, Recorder
+
+
+@dataclasses.dataclass
+class CursorFetchAllParams:
+    connection_name: str
+
+
+@dataclasses.dataclass
+class CursorFetchAllResult:
+    results: List[Any]
+
+    def _to_dict(self) -> Dict[str, Any]:
+        processed_results = []
+        for result in self.results:
+            result = tuple(map(self._process_value, result))
+            processed_results.append(result)
+
+        return {"results": processed_results}
+
+    @classmethod
+    def _from_dict(cls, dct: Mapping) -> "CursorFetchAllResult":
+        unprocessed_results = []
+        for result in dct["results"]:
+            result = tuple(map(cls._unprocess_value, result))
+            unprocessed_results.append(result)
+
+        return CursorFetchAllResult(unprocessed_results)
+
+    @classmethod
+    def _process_value(cls, value: Any) -> Any:
+        if type(value) is datetime.date:
+            return {"type": "date", "value": value.isoformat()}
+        elif type(value) is datetime.datetime:
+            return {"type": "datetime", "value": value.isoformat()}
+        elif type(value) is decimal.Decimal:
+            return float(value)
+        else:
+            return value
+
+    @classmethod
+    def _unprocess_value(cls, value: Any) -> Any:
+        if type(value) is dict:
+            value_type = value.get("type")
+            if value_type == "date":
+                date_string = value.get("value")
+                assert isinstance(date_string, str)
+                return datetime.date.fromisoformat(date_string)
+            elif value_type == "datetime":
+                date_string = value.get("value")
+                assert isinstance(date_string, str)
+                return datetime.datetime.fromisoformat(date_string)
+            return value
+        else:
+            return value
+
+
+@Recorder.register_record_type
+class CursorFetchAllRecord(Record):
+    """Implements record/replay support for the cursor.fetchall() method."""
+
+    params_cls = CursorFetchAllParams
+    result_cls = CursorFetchAllResult
+    group = "Database"

dbt/adapters/record/cursor/fetchmany.py

@@ -0,0 +1,23 @@
+import dataclasses
+from typing import Any, List
+
+from dbt_common.record import Record, Recorder
+
+
+@dataclasses.dataclass
+class CursorFetchManyParams:
+    connection_name: str
+
+
+@dataclasses.dataclass
+class CursorFetchManyResult:
+    results: List[Any]
+
+
+@Recorder.register_record_type
+class CursorFetchManyRecord(Record):
+    """Implements record/replay support for the cursor.fetchmany() method."""
+
+    params_cls = CursorFetchManyParams
+    result_cls = CursorFetchManyResult
+    group = "Database"