feature/decouple adapters from core

.changes/unreleased/Under the Hood-20240111-114806.yaml

@@ -0,0 +1,6 @@
+kind: Under the Hood
+body: Update import paths and list_relations to support decoupling adapters/core
+time: 2024-01-11T11:48:06.120111-08:00
+custom:
+  Author: colin-rogers-dbt
+  Issue: "972"

dagger/run_dbt_spark_tests.py

@@ -2,6 +2,7 @@
 
 import argparse
 import sys
+from typing import Dict
 
 import anyio as anyio
 import dagger as dagger
@@ -19,7 +20,7 @@
 TESTING_ENV_VARS.update({"ODBC_DRIVER": "/opt/simba/spark/lib/64/libsparkodbc_sb64.so"})
 
 
-def env_variables(envs: dict[str, str]):
+def env_variables(envs: Dict[str, str]):
     def env_variables_inner(ctr: dagger.Container):
         for key, value in envs.items():
             ctr = ctr.with_env_variable(key, value)

dbt/adapters/spark/column.py

@@ -2,7 +2,7 @@
 from typing import Any, Dict, Optional, TypeVar, Union
 
 from dbt.adapters.base.column import Column
-from dbt.dataclass_schema import dbtClassMixin
+from dbt.common.dataclass_schema import dbtClassMixin
 
 Self = TypeVar("Self", bound="SparkColumn")
 

dbt/adapters/spark/connections.py

@@ -1,11 +1,13 @@
 from contextlib import contextmanager
 
-import dbt.exceptions
 from dbt.adapters.base import Credentials
+from dbt.adapters.contracts.connection import AdapterResponse, ConnectionState, Connection
+from dbt.adapters.events.logging import AdapterLogger
+from dbt.adapters.exceptions import FailedToConnectError
 from dbt.adapters.sql import SQLConnectionManager
-from dbt.contracts.connection import ConnectionState, AdapterResponse
-from dbt.events import AdapterLogger
-from dbt.utils import DECIMALS
+from dbt.common.exceptions import DbtConfigError, DbtRuntimeError, DbtDatabaseError
+
+from dbt.common.utils.encoding import DECIMALS
 from dbt.adapters.spark import __version__
 
 try:
@@ -22,8 +24,7 @@
     pyodbc = None
 from datetime import datetime
 import sqlparams
-from dbt.contracts.connection import Connection
-from dbt.dataclass_schema import StrEnum
+from dbt.common.dataclass_schema import StrEnum
 from dataclasses import dataclass, field
 from typing import Any, Dict, Optional, Union, Tuple, List, Generator, Iterable, Sequence
 
@@ -92,15 +93,15 @@ def cluster_id(self) -> Optional[str]:
 
     def __post_init__(self) -> None:
         if self.method is None:
-            raise dbt.exceptions.DbtRuntimeError("Must specify `method` in profile")
+            raise DbtRuntimeError("Must specify `method` in profile")
         if self.host is None:
-            raise dbt.exceptions.DbtRuntimeError("Must specify `host` in profile")
+            raise DbtRuntimeError("Must specify `host` in profile")
         if self.schema is None:
-            raise dbt.exceptions.DbtRuntimeError("Must specify `schema` in profile")
+            raise DbtRuntimeError("Must specify `schema` in profile")
 
         # spark classifies database and schema as the same thing
         if self.database is not None and self.database != self.schema:
-            raise dbt.exceptions.DbtRuntimeError(
+            raise DbtRuntimeError(
                 f"    schema: {self.schema} \n"
                 f"    database: {self.database} \n"
                 f"On Spark, database must be omitted or have the same value as"
@@ -112,7 +113,7 @@ def __post_init__(self) -> None:
             try:
                 import pyodbc  # noqa: F401
             except ImportError as e:
-                raise dbt.exceptions.DbtRuntimeError(
+                raise DbtRuntimeError(
                     f"{self.method} connection method requires "
                     "additional dependencies. \n"
                     "Install the additional required dependencies with "
@@ -121,7 +122,7 @@ def __post_init__(self) -> None:
                 ) from e
 
         if self.method == SparkConnectionMethod.ODBC and self.cluster and self.endpoint:
-            raise dbt.exceptions.DbtRuntimeError(
+            raise DbtRuntimeError(
                 "`cluster` and `endpoint` cannot both be set when"
                 f" using {self.method} method to connect to Spark"
             )
@@ -130,7 +131,7 @@ def __post_init__(self) -> None:
             self.method == SparkConnectionMethod.HTTP
             or self.method == SparkConnectionMethod.THRIFT
         ) and not (ThriftState and THttpClient and hive):
-            raise dbt.exceptions.DbtRuntimeError(
+            raise DbtRuntimeError(
                 f"{self.method} connection method requires "
                 "additional dependencies. \n"
                 "Install the additional required dependencies with "
@@ -141,7 +142,7 @@ def __post_init__(self) -> None:
             try:
                 import pyspark  # noqa: F401
             except ImportError as e:
-                raise dbt.exceptions.DbtRuntimeError(
+                raise DbtRuntimeError(
                     f"{self.method} connection method requires "
                     "additional dependencies. \n"
                     "Install the additional required dependencies with "
@@ -291,13 +292,11 @@ def execute(self, sql: str, bindings: Optional[List[Any]] = None) -> None:
         if poll_state.errorMessage:
             logger.debug("Poll response: {}".format(poll_state))
             logger.debug("Poll status: {}".format(state))
-            raise dbt.exceptions.DbtDatabaseError(poll_state.errorMessage)
+            raise DbtDatabaseError(poll_state.errorMessage)
 
         elif state not in STATE_SUCCESS:
             status_type = ThriftState._VALUES_TO_NAMES.get(state, "Unknown<{!r}>".format(state))
-            raise dbt.exceptions.DbtDatabaseError(
-                "Query failed with status: {}".format(status_type)
-            )
+            raise DbtDatabaseError("Query failed with status: {}".format(status_type))
 
         logger.debug("Poll status: {}, query complete".format(state))
 
@@ -358,9 +357,9 @@ def exception_handler(self, sql: str) -> Generator[None, None, None]:
             thrift_resp = exc.args[0]
             if hasattr(thrift_resp, "status"):
                 msg = thrift_resp.status.errorMessage
-                raise dbt.exceptions.DbtRuntimeError(msg)
+                raise DbtRuntimeError(msg)
             else:
-                raise dbt.exceptions.DbtRuntimeError(str(exc))
+                raise DbtRuntimeError(str(exc))
 
     def cancel(self, connection: Connection) -> None:
         connection.handle.cancel()
@@ -390,7 +389,7 @@ def validate_creds(cls, creds: Any, required: Iterable[str]) -> None:
 
         for key in required:
             if not hasattr(creds, key):
-                raise dbt.exceptions.DbtProfileError(
+                raise DbtConfigError(
                     "The config '{}' is required when using the {} method"
                     " to connect to Spark".format(key, method)
                 )
@@ -481,7 +480,7 @@ def open(cls, connection: Connection) -> Connection:
                             endpoint=creds.endpoint
                         )
                     else:
-                        raise dbt.exceptions.DbtProfileError(
+                        raise DbtConfigError(
                             "Either `cluster` or `endpoint` must set when"
                             " using the odbc method to connect to Spark"
                         )
@@ -525,9 +524,7 @@ def open(cls, connection: Connection) -> Connection:
                         Connection(server_side_parameters=creds.server_side_parameters)
                     )
                 else:
-                    raise dbt.exceptions.DbtProfileError(
-                        f"invalid credential method: {creds.method}"
-                    )
+                    raise DbtConfigError(f"invalid credential method: {creds.method}")
                 break
             except Exception as e:
                 exc = e
@@ -537,7 +534,7 @@ def open(cls, connection: Connection) -> Connection:
                     msg = "Failed to connect"
                     if creds.token is not None:
                         msg += ", is your token valid?"
-                    raise dbt.exceptions.FailedToConnectError(msg) from e
+                    raise FailedToConnectError(msg) from e
                 retryable_message = _is_retryable_error(e)
                 if retryable_message and creds.connect_retries > 0:
                     msg = (
@@ -558,7 +555,7 @@ def open(cls, connection: Connection) -> Connection:
                     logger.warning(msg)
                     time.sleep(creds.connect_timeout)
                 else:
-                    raise dbt.exceptions.FailedToConnectError("failed to connect") from e
+                    raise FailedToConnectError("failed to connect") from e
         else:
             raise exc  # type: ignore
 

dbt/adapters/spark/impl.py

@@ -1,18 +1,31 @@
+import os
 import re
 from concurrent.futures import Future
 from dataclasses import dataclass
-from typing import Any, Dict, Iterable, List, Optional, Union, Type, Tuple, Callable, Set
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Union,
+    Type,
+    Tuple,
+    Callable,
+    Set,
+    FrozenSet,
+)
 
 from dbt.adapters.base.relation import InformationSchema
-from dbt.contracts.graph.manifest import Manifest
+from dbt.adapters.contracts.connection import AdapterResponse
+from dbt.adapters.events.logging import AdapterLogger
+from dbt.common.exceptions import DbtRuntimeError, CompilationError
+from dbt.common.utils import AttrDict, executor
 
 from typing_extensions import TypeAlias
 
 import agate
 
-import dbt
-import dbt.exceptions
-
 from dbt.adapters.base import AdapterConfig, PythonJobHelper
 from dbt.adapters.base.impl import catch_as_completed, ConstraintSupport
 from dbt.adapters.sql import SQLAdapter
@@ -24,14 +37,16 @@
     AllPurposeClusterPythonJobHelper,
 )
 from dbt.adapters.base import BaseRelation
-from dbt.clients.agate_helper import DEFAULT_TYPE_TESTER
-from dbt.contracts.connection import AdapterResponse
-from dbt.contracts.graph.nodes import ConstraintType
-from dbt.contracts.relation import RelationType
-from dbt.events import AdapterLogger
-from dbt.utils import executor, AttrDict
+from dbt.adapters.contracts.relation import RelationType, RelationConfig
+from dbt.common.clients.agate_helper import DEFAULT_TYPE_TESTER
+from dbt.common.contracts.constraints import ConstraintType
 
 logger = AdapterLogger("Spark")
+packages = ["pyhive.hive", "thrift.transport", "thrift.protocol"]
+log_level = os.getenv("DBT_SPARK_LOG_LEVEL", "ERROR")
+for package in packages:
+    logger.debug(f"Setting {package} logging to {log_level}")
+    logger.set_adapter_dependency_log_level(package, log_level)
 
 GET_COLUMNS_IN_RELATION_RAW_MACRO_NAME = "get_columns_in_relation_raw"
 LIST_SCHEMAS_MACRO_NAME = "list_schemas"
@@ -144,7 +159,7 @@ def _get_relation_information(self, row: agate.Row) -> RelationInfo:
         try:
             _schema, name, _, information = row
         except ValueError:
-            raise dbt.exceptions.DbtRuntimeError(
+            raise DbtRuntimeError(
                 f'Invalid value from "show tables extended ...", got {len(row)} values, expected 4'
             )
 
@@ -155,7 +170,7 @@ def _get_relation_information_using_describe(self, row: agate.Row) -> RelationIn
         try:
             _schema, name, _ = row
         except ValueError:
-            raise dbt.exceptions.DbtRuntimeError(
+            raise DbtRuntimeError(
                 f'Invalid value from "show tables ...", got {len(row)} values, expected 3'
             )
 
@@ -164,7 +179,7 @@ def _get_relation_information_using_describe(self, row: agate.Row) -> RelationIn
             table_results = self.execute_macro(
                 DESCRIBE_TABLE_EXTENDED_MACRO_NAME, kwargs={"table_name": table_name}
             )
-        except dbt.exceptions.DbtRuntimeError as e:
+        except DbtRuntimeError as e:
             logger.debug(f"Error while retrieving information about {table_name}: {e.msg}")
             table_results = AttrDict()
 
@@ -219,7 +234,7 @@ def list_relations_without_caching(self, schema_relation: BaseRelation) -> List[
                 row_list=show_table_extended_rows,
                 relation_info_func=self._get_relation_information,
             )
-        except dbt.exceptions.DbtRuntimeError as e:
+        except DbtRuntimeError as e:
             errmsg = getattr(e, "msg", "")
             if f"Database '{schema_relation}' not found" in errmsg:
                 return []
@@ -236,7 +251,7 @@ def list_relations_without_caching(self, schema_relation: BaseRelation) -> List[
                         row_list=show_table_rows,
                         relation_info_func=self._get_relation_information_using_describe,
                     )
-                except dbt.exceptions.DbtRuntimeError as e:
+                except DbtRuntimeError as e:
                     description = "Error while retrieving information about"
                     logger.debug(f"{description} {schema_relation}: {e.msg}")
                     return []
@@ -298,7 +313,7 @@ def get_columns_in_relation(self, relation: BaseRelation) -> List[SparkColumn]:
                 GET_COLUMNS_IN_RELATION_RAW_MACRO_NAME, kwargs={"relation": relation}
             )
             columns = self.parse_describe_extended(relation, rows)
-        except dbt.exceptions.DbtRuntimeError as e:
+        except DbtRuntimeError as e:
             # spark would throw error when table doesn't exist, where other
             # CDW would just return and empty list, normalizing the behavior here
             errmsg = getattr(e, "msg", "")
@@ -352,11 +367,13 @@ def _get_columns_for_catalog(self, relation: BaseRelation) -> Iterable[Dict[str,
             yield as_dict
 
     def get_catalog(
-        self, manifest: Manifest, selected_nodes: Optional[Set] = None
+        self,
+        relation_configs: Iterable[RelationConfig],
+        used_schemas: FrozenSet[Tuple[str, str]],
     ) -> Tuple[agate.Table, List[Exception]]:
-        schema_map = self._get_catalog_schemas(manifest)
+        schema_map = self._get_catalog_schemas(relation_configs)
         if len(schema_map) > 1:
-            raise dbt.exceptions.CompilationError(
+            raise CompilationError(
                 f"Expected only one database in get_catalog, found " f"{list(schema_map)}"
             )
 
@@ -371,7 +388,7 @@ def get_catalog(
                             self._get_one_catalog,
                             info,
                             [schema],
-                            manifest,
+                            relation_configs,
                         )
                     )
             catalogs, exceptions = catch_as_completed(futures)
@@ -381,10 +398,10 @@ def _get_one_catalog(
         self,
         information_schema: InformationSchema,
         schemas: Set[str],
-        manifest: Manifest,
+        used_schemas: FrozenSet[Tuple[str, str]],
     ) -> agate.Table:
         if len(schemas) != 1:
-            raise dbt.exceptions.CompilationError(
+            raise CompilationError(
                 f"Expected only one schema in spark _get_one_catalog, found " f"{schemas}"
             )