dbt-adapters 0.1.0a1__py3-none-any.whl

dbt/adapters/protocol.py

@@ -0,0 +1,173 @@
+from dataclasses import dataclass
+from typing import (
+    Any,
+    ContextManager,
+    Dict,
+    Generic,
+    Hashable,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Tuple,
+)
+from typing_extensions import Protocol
+
+import agate
+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
+
+
+@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(  # type: ignore[misc]
+    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/reference_keys.py

@@ -0,0 +1,39 @@
+# this module exists to resolve circular imports with the events module
+from collections import namedtuple
+from typing import Any, Optional
+
+
+_ReferenceKey = namedtuple("_ReferenceKey", "database schema identifier")
+
+
+def lowercase(value: Optional[str]) -> Optional[str]:
+    if value is None:
+        return None
+    else:
+        return value.lower()
+
+
+# For backwards compatibility. New code should use _make_ref_key
+def _make_key(relation: Any) -> _ReferenceKey:
+    return _make_ref_key(relation)
+
+
+def _make_ref_key(relation: Any) -> _ReferenceKey:
+    """
+    Make _ReferenceKeys with lowercase values for the cache,
+    so we don't have to keep track of quoting
+    """
+    # databases and schemas can both be None
+    return _ReferenceKey(
+        lowercase(relation.database),
+        lowercase(relation.schema),
+        lowercase(relation.identifier),
+    )
+
+
+def _make_ref_key_dict(relation: Any):
+    return {
+        "database": relation.database,
+        "schema": relation.schema,
+        "identifier": relation.identifier,
+    }

dbt/adapters/relation_configs/README.md

@@ -0,0 +1,25 @@
+# RelationConfig
+This package serves as an initial abstraction for managing the inspection of existing relations and determining
+changes on those relations. It arose from the materialized view work and is currently only supporting 
+materialized views for Postgres and Redshift as well as dynamic tables for Snowflake. There are three main
+classes in this package.
+
+## RelationConfigBase
+This is a very small class that only has a `from_dict()` method and a default `NotImplementedError()`. At some
+point this could be replaced by a more robust framework, like `mashumaro` or `pydantic`.
+
+## RelationConfigChange
+This class inherits from `RelationConfigBase` ; however, this can be thought of as a separate class. The subclassing
+merely points to the idea that both classes would likely inherit from the same class in a `mashumaro` or
+`pydantic` implementation. This class is much more restricted in attribution. It should really only
+ever need an `action` and a `context`. This can be though of as being analogous to a web request. You need to
+know what you're doing (`action`: 'create' = GET, 'drop' = DELETE, etc.) and the information (`context`) needed
+to make the change. In our scenarios, the context tends to be an instance of `RelationConfigBase` corresponding
+to the new state.
+
+## RelationConfigValidationMixin
+This mixin provides optional validation mechanics that can be applied to either `RelationConfigBase` or
+`RelationConfigChange` subclasses. A validation rule is a combination of a `validation_check`, something
+that should evaluate to `True`, and an optional `validation_error`, an instance of `DbtRuntimeError`
+that should be raised in the event the `validation_check` fails. While optional, it's recommended that
+the `validation_error` be provided for clearer transparency to the end user.

dbt/adapters/relation_configs/__init__.py

@@ -0,0 +1,12 @@
+from dbt.adapters.relation_configs.config_base import (
+    RelationConfigBase,
+    RelationResults,
+)
+from dbt.adapters.relation_configs.config_change import (
+    RelationConfigChange,
+    RelationConfigChangeAction,
+)
+from dbt.adapters.relation_configs.config_validation import (
+    RelationConfigValidationMixin,
+    RelationConfigValidationRule,
+)

dbt/adapters/relation_configs/config_base.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass
+from typing import Dict, Union
+
+import agate
+from dbt_common.utils import filter_null_values
+
+
+"""
+This is what relation metadata from the database looks like. It's a dictionary because there will be
+multiple grains of data for a single object. For example, a materialized view in Postgres has base level information,
+like name. But it also can have multiple indexes, which needs to be a separate query. It might look like this:
+
+{
+    "base": agate.Row({"table_name": "table_abc", "query": "select * from table_def"})
+    "indexes": agate.Table("rows": [
+        agate.Row({"name": "index_a", "columns": ["column_a"], "type": "hash", "unique": False}),
+        agate.Row({"name": "index_b", "columns": ["time_dim_a"], "type": "btree", "unique": False}),
+    ])
+}
+"""
+RelationResults = Dict[str, Union[agate.Row, agate.Table]]
+
+
+@dataclass(frozen=True)
+class RelationConfigBase:
+    @classmethod
+    def from_dict(cls, kwargs_dict) -> "RelationConfigBase":
+        """
+        This assumes the subclass of `RelationConfigBase` is flat, in the sense that no attribute is
+        itself another subclass of `RelationConfigBase`. If that's not the case, this should be overriden
+        to manually manage that complexity.
+
+        Args:
+            kwargs_dict: the dict representation of this instance
+
+        Returns: the `RelationConfigBase` representation associated with the provided dict
+        """
+        return cls(**filter_null_values(kwargs_dict))  # type: ignore
+
+    @classmethod
+    def _not_implemented_error(cls) -> NotImplementedError:
+        return NotImplementedError(
+            "This relation type has not been fully configured for this adapter."
+        )

dbt/adapters/relation_configs/config_change.py

@@ -0,0 +1,24 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Hashable
+
+from dbt_common.dataclass_schema import StrEnum
+
+from dbt.adapters.relation_configs.config_base import RelationConfigBase
+
+
+class RelationConfigChangeAction(StrEnum):
+    alter = "alter"
+    create = "create"
+    drop = "drop"
+
+
+@dataclass(frozen=True, eq=True, unsafe_hash=True)
+class RelationConfigChange(RelationConfigBase, ABC):
+    action: RelationConfigChangeAction
+    context: Hashable  # this is usually a RelationConfig, e.g. IndexConfig, but shouldn't be limited
+
+    @property
+    @abstractmethod
+    def requires_full_refresh(self) -> bool:
+        raise self._not_implemented_error()

dbt/adapters/relation_configs/config_validation.py

@@ -0,0 +1,57 @@
+from dataclasses import dataclass
+from typing import Optional, Set
+
+from dbt_common.exceptions import DbtRuntimeError
+
+
+@dataclass(frozen=True, eq=True, unsafe_hash=True)
+class RelationConfigValidationRule:
+    validation_check: bool
+    validation_error: Optional[DbtRuntimeError]
+
+    @property
+    def default_error(self):
+        return DbtRuntimeError(
+            "There was a validation error in preparing this relation config."
+            "No additional context was provided by this adapter."
+        )
+
+
+@dataclass(frozen=True)
+class RelationConfigValidationMixin:
+    def __post_init__(self):
+        self.run_validation_rules()
+
+    @property
+    def validation_rules(self) -> Set[RelationConfigValidationRule]:
+        """
+        A set of validation rules to run against the object upon creation.
+
+        A validation rule is a combination of a validation check (bool) and an optional error message.
+
+        This defaults to no validation rules if not implemented. It's recommended to override this with values,
+        but that may not always be necessary.
+
+        Returns: a set of validation rules
+        """
+        return set()
+
+    def run_validation_rules(self):
+        for validation_rule in self.validation_rules:
+            try:
+                assert validation_rule.validation_check
+            except AssertionError:
+                if validation_rule.validation_error:
+                    raise validation_rule.validation_error
+                else:
+                    raise validation_rule.default_error
+        self.run_child_validation_rules()
+
+    def run_child_validation_rules(self):
+        for attr_value in vars(self).values():
+            if hasattr(attr_value, "validation_rules"):
+                attr_value.run_validation_rules()
+            if isinstance(attr_value, set):
+                for member in attr_value:
+                    if hasattr(member, "validation_rules"):
+                        member.run_validation_rules()

dbt/adapters/sql/__init__.py

@@ -0,0 +1,2 @@
+from dbt.adapters.sql.connections import SQLConnectionManager
+from dbt.adapters.sql.impl import SQLAdapter

dbt/adapters/sql/connections.py

@@ -0,0 +1,195 @@
+import abc
+import time
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+
+import agate
+from dbt_common.clients.agate_helper import empty_table, table_from_data_flat
+from dbt_common.events.contextvars import get_node_info
+from dbt_common.events.functions import fire_event
+from dbt_common.exceptions import DbtInternalError, NotImplementedError
+from dbt_common.utils import cast_to_str
+
+from dbt.adapters.base import BaseConnectionManager
+from dbt.adapters.contracts.connection import (
+    AdapterResponse,
+    Connection,
+    ConnectionState,
+)
+from dbt.adapters.events.types import (
+    ConnectionUsed,
+    SQLCommit,
+    SQLQuery,
+    SQLQueryStatus,
+)
+
+
+class SQLConnectionManager(BaseConnectionManager):
+    """The default connection manager with some common SQL methods implemented.
+
+    Methods to implement:
+        - exception_handler
+        - cancel
+        - get_response
+        - open
+    """
+
+    @abc.abstractmethod
+    def cancel(self, connection: Connection):
+        """Cancel the given connection."""
+        raise NotImplementedError("`cancel` is not implemented for this adapter!")
+
+    def cancel_open(self) -> List[str]:
+        names = []
+        this_connection = self.get_if_exists()
+        with self.lock:
+            for connection in self.thread_connections.values():
+                if connection is this_connection:
+                    continue
+
+                # if the connection failed, the handle will be None so we have
+                # nothing to cancel.
+                if connection.handle is not None and connection.state == ConnectionState.OPEN:
+                    self.cancel(connection)
+                if connection.name is not None:
+                    names.append(connection.name)
+        return names
+
+    def add_query(
+        self,
+        sql: str,
+        auto_begin: bool = True,
+        bindings: Optional[Any] = None,
+        abridge_sql_log: bool = False,
+    ) -> Tuple[Connection, Any]:
+        connection = self.get_thread_connection()
+        if auto_begin and connection.transaction_open is False:
+            self.begin()
+        fire_event(
+            ConnectionUsed(
+                conn_type=self.TYPE,
+                conn_name=cast_to_str(connection.name),
+                node_info=get_node_info(),
+            )
+        )
+
+        with self.exception_handler(sql):
+            if abridge_sql_log:
+                log_sql = "{}...".format(sql[:512])
+            else:
+                log_sql = sql
+
+            fire_event(
+                SQLQuery(
+                    conn_name=cast_to_str(connection.name),
+                    sql=log_sql,
+                    node_info=get_node_info(),
+                )
+            )
+            pre = time.time()
+
+            cursor = connection.handle.cursor()
+            cursor.execute(sql, bindings)
+
+            fire_event(
+                SQLQueryStatus(
+                    status=str(self.get_response(cursor)),
+                    elapsed=round((time.time() - pre)),
+                    node_info=get_node_info(),
+                )
+            )
+
+            return connection, cursor
+
+    @classmethod
+    @abc.abstractmethod
+    def get_response(cls, cursor: Any) -> AdapterResponse:
+        """Get the status of the cursor."""
+        raise NotImplementedError("`get_response` is not implemented for this adapter!")
+
+    @classmethod
+    def process_results(
+        cls, column_names: Iterable[str], rows: Iterable[Any]
+    ) -> List[Dict[str, Any]]:
+        # TODO CT-211
+        unique_col_names = dict()  # type: ignore[var-annotated]
+        # TODO CT-211
+        for idx in range(len(column_names)):  # type: ignore[arg-type]
+            # TODO CT-211
+            col_name = column_names[idx]  # type: ignore[index]
+            if col_name in unique_col_names:
+                unique_col_names[col_name] += 1
+                # TODO CT-211
+                column_names[idx] = f"{col_name}_{unique_col_names[col_name]}"  # type: ignore[index] # noqa
+            else:
+                # TODO CT-211
+                unique_col_names[column_names[idx]] = 1  # type: ignore[index]
+        return [dict(zip(column_names, row)) for row in rows]
+
+    @classmethod
+    def get_result_from_cursor(cls, cursor: Any, limit: Optional[int]) -> agate.Table:
+        data: List[Any] = []
+        column_names: List[str] = []
+
+        if cursor.description is not None:
+            column_names = [col[0] for col in cursor.description]
+            if limit:
+                rows = cursor.fetchmany(limit)
+            else:
+                rows = cursor.fetchall()
+            data = cls.process_results(column_names, rows)
+
+        return table_from_data_flat(data, column_names)
+
+    def execute(
+        self,
+        sql: str,
+        auto_begin: bool = False,
+        fetch: bool = False,
+        limit: Optional[int] = None,
+    ) -> Tuple[AdapterResponse, agate.Table]:
+        sql = self._add_query_comment(sql)
+        _, cursor = self.add_query(sql, auto_begin)
+        response = self.get_response(cursor)
+        if fetch:
+            table = self.get_result_from_cursor(cursor, limit)
+        else:
+            table = empty_table()
+        return response, table
+
+    def add_begin_query(self):
+        return self.add_query("BEGIN", auto_begin=False)
+
+    def add_commit_query(self):
+        return self.add_query("COMMIT", auto_begin=False)
+
+    def add_select_query(self, sql: str) -> Tuple[Connection, Any]:
+        sql = self._add_query_comment(sql)
+        return self.add_query(sql, auto_begin=False)
+
+    def begin(self):
+        connection = self.get_thread_connection()
+        if connection.transaction_open is True:
+            raise DbtInternalError(
+                'Tried to begin a new transaction on connection "{}", but '
+                "it already had one open!".format(connection.name)
+            )
+
+        self.add_begin_query()
+
+        connection.transaction_open = True
+        return connection
+
+    def commit(self):
+        connection = self.get_thread_connection()
+        if connection.transaction_open is False:
+            raise DbtInternalError(
+                'Tried to commit transaction on connection "{}", but '
+                "it does not have one open!".format(connection.name)
+            )
+
+        fire_event(SQLCommit(conn_name=connection.name, node_info=get_node_info()))
+        self.add_commit_query()
+
+        connection.transaction_open = False
+
+        return connection