From 947adf6e66c36bd94749dcb350d079b97af7fc96 Mon Sep 17 00:00:00 2001 From: dkruh1 <101973979+dkruh1@users.noreply.github.com> Date: Wed, 27 Nov 2024 12:20:18 +0200 Subject: [PATCH 1/9] support session method --- dbt/adapters/databricks/__version__.py | 2 +- dbt/adapters/databricks/connections.py | 65 ++++++++++++++++++- dbt/adapters/databricks/impl.py | 16 +++-- dbt/adapters/databricks/session_connection.py | 41 ++++++++++++ .../databricks/macros/adapters/metadata.sql | 34 +++++++++- 5 files changed, 146 insertions(+), 12 deletions(-) create mode 100644 dbt/adapters/databricks/session_connection.py diff --git a/dbt/adapters/databricks/__version__.py b/dbt/adapters/databricks/__version__.py index 8fe277395..624cb06a2 100644 --- a/dbt/adapters/databricks/__version__.py +++ b/dbt/adapters/databricks/__version__.py @@ -1 +1 @@ -version: str = "1.8.6" +version: str = "1.8.6a14" diff --git a/dbt/adapters/databricks/connections.py b/dbt/adapters/databricks/connections.py index cf6f600a7..dd6aad8ef 100644 --- a/dbt/adapters/databricks/connections.py +++ b/dbt/adapters/databricks/connections.py @@ -10,7 +10,7 @@ from multiprocessing.context import SpawnContext from numbers import Number from threading import get_ident -from typing import Any +from typing import Any, Generator from typing import Callable from typing import cast from typing import Dict @@ -70,7 +70,9 @@ from dbt.adapters.events.types import NewConnection from dbt.adapters.events.types import SQLQuery from dbt.adapters.events.types import SQLQueryStatus -from dbt.adapters.spark.connections import SparkConnectionManager +from dbt.adapters.spark.connections import ( + SparkConnectionManager, +) from dbt_common.events.contextvars import get_node_info from dbt_common.events.functions import fire_event from dbt_common.exceptions import DbtInternalError @@ -93,6 +95,8 @@ # toggle for session managements that minimizes the number of sessions opened/closed USE_LONG_SESSIONS = os.getenv("DBT_DATABRICKS_LONG_SESSIONS", "True").upper() == "TRUE" +# toggle for session managements that assumes the adapter is running in a Databricks session +USE_SESSION_CONNECTION = os.getenv("DBT_DATABRICKS_SESSION_CONNECTION", "False").upper() == "TRUE" # Number of idle seconds before a connection is automatically closed. Only applicable if # USE_LONG_SESSIONS is true. @@ -1086,6 +1090,63 @@ def exponential_backoff(attempt: int) -> int: ) +class DatabricksSessionConnectionManager(DatabricksConnectionManager): + + def cancel_open(self) -> List[str]: + SparkConnectionManager.cancel_open(self) + + def compare_dbr_version(self, major: int, minor: int) -> int: + version = (major, minor) + connection = ( + self.get_thread_connection().handle + ) + dbr_version = connection.dbr_version + return (dbr_version > version) - (dbr_version < version) + + def set_query_header(self, query_header_context: Dict[str, Any]) -> None: + SparkConnectionManager.set_query_header(self, query_header_context) + + def set_connection_name( + self, name: Optional[str] = None, query_header_context: Any = None + ) -> Connection: + SparkConnectionManager.set_connection_name(self, name) + + def add_query( + self, + sql: str, + auto_begin: bool = True, + bindings: Optional[Any] = None, + abridge_sql_log: bool = False, + *, + close_cursor: bool = False, + ) -> Tuple[Connection, Any]: + return SparkConnectionManager.add_query( + self, sql, auto_begin, bindings, abridge_sql_log + ) + + def list_schemas(self, database: str, schema: Optional[str] = None) -> "Table": + raise NotImplementedError("list_schemas is not implemented for DatabricksSessionConnectionManager - should call the list_schemas macro instead") + + def list_tables(self, database: str, schema: str, identifier: Optional[str] = None) -> "Table": + raise NotImplementedError("list_tables is not implemented for DatabricksSessionConnectionManager - should call the list_tables macro instead") + + @classmethod + def open(cls, connection: Connection) -> Connection: + from dbt.adapters.spark.session import Connection + from dbt.adapters.databricks.session_connection import DatabricksSessionConnectionWrapper + handle = DatabricksSessionConnectionWrapper( + Connection() + ) + connection.handle = handle + connection.state = ConnectionState.OPEN + return connection + + @classmethod + def get_response(cls, cursor) -> DatabricksAdapterResponse: + response = SparkConnectionManager.get_response(cursor) + return DatabricksAdapterResponse(_message=response._message, query_id=None) + + def _get_pipeline_state(session: Session, host: str, pipeline_id: str) -> dict: pipeline_url = f"https://{host}/api/2.0/pipelines/{pipeline_id}" diff --git a/dbt/adapters/databricks/impl.py b/dbt/adapters/databricks/impl.py index 46cb11dd6..93d9a8362 100644 --- a/dbt/adapters/databricks/impl.py +++ b/dbt/adapters/databricks/impl.py @@ -38,11 +38,11 @@ from dbt.adapters.contracts.relation import RelationConfig from dbt.adapters.contracts.relation import RelationType from dbt.adapters.databricks.column import DatabricksColumn -from dbt.adapters.databricks.connections import DatabricksConnectionManager +from dbt.adapters.databricks.connections import DatabricksConnectionManager, DatabricksSessionConnectionManager from dbt.adapters.databricks.connections import DatabricksDBTConnection from dbt.adapters.databricks.connections import DatabricksSQLConnectionWrapper from dbt.adapters.databricks.connections import ExtendedSessionConnectionManager -from dbt.adapters.databricks.connections import USE_LONG_SESSIONS +from dbt.adapters.databricks.connections import USE_LONG_SESSIONS, USE_SESSION_CONNECTION from dbt.adapters.databricks.python_submissions import ( DbtDatabricksAllPurposeClusterPythonJobHelper, ) @@ -148,7 +148,9 @@ class DatabricksAdapter(SparkAdapter): Relation = DatabricksRelation Column = DatabricksColumn - if USE_LONG_SESSIONS: + if USE_SESSION_CONNECTION: + ConnectionManager: Type[DatabricksConnectionManager] = DatabricksSessionConnectionManager + elif USE_LONG_SESSIONS: ConnectionManager: Type[DatabricksConnectionManager] = ExtendedSessionConnectionManager else: ConnectionManager = DatabricksConnectionManager @@ -207,10 +209,12 @@ def list_schemas(self, database: Optional[str]) -> List[str]: If `database` is `None`, fallback to executing `show databases` because `list_schemas` tries to collect schemas from all catalogs when `database` is `None`. """ - if database is not None: + if database is not None and not USE_SESSION_CONNECTION: results = self.connections.list_schemas(database=database) else: - results = self.execute_macro(LIST_SCHEMAS_MACRO_NAME, kwargs={"database": database}) + results = self.execute_macro( + LIST_SCHEMAS_MACRO_NAME, kwargs={"database": database} + ) return [row[0] for row in results] def check_schema_exists(self, database: Optional[str], schema: str) -> bool: @@ -281,7 +285,7 @@ def _get_hive_relations( kwargs = {"relation": relation} new_rows: List[Tuple[str, Optional[str]]] - if all([relation.database, relation.schema]): + if all([relation.database, relation.schema]) and not USE_SESSION_CONNECTION: tables = self.connections.list_tables( database=relation.database, schema=relation.schema # type: ignore[arg-type] ) diff --git a/dbt/adapters/databricks/session_connection.py b/dbt/adapters/databricks/session_connection.py new file mode 100644 index 000000000..a42e565fc --- /dev/null +++ b/dbt/adapters/databricks/session_connection.py @@ -0,0 +1,41 @@ +import re +import sys +from typing import Tuple +from dbt.adapters.spark.session import SessionConnectionWrapper, Connection + + +DBR_VERSION_REGEX = re.compile(r"([1-9][0-9]*)\.(x|0|[1-9][0-9]*)") + + +class DatabricksSessionConnectionWrapper(SessionConnectionWrapper): + + _is_cluster: bool + + def __init__(self, handle: Connection) -> None: + super().__init__(handle) + self._is_cluster = True + self.cursor() + + @property + def dbr_version(self) -> Tuple[int, int]: + if not hasattr(self, "_dbr_version"): + if self._is_cluster: + with self._cursor() as cursor: + cursor.execute("SET spark.databricks.clusterUsageTags.sparkVersion") + results = cursor.fetchone() + if results: + dbr_version: str = results[1] + + m = DBR_VERSION_REGEX.search(dbr_version) + assert m, f"Unknown DBR version: {dbr_version}" + major = int(m.group(1)) + try: + minor = int(m.group(2)) + except ValueError: + minor = sys.maxsize + self._dbr_version = (major, minor) + else: + # Assuming SQL Warehouse uses the latest version. + self._dbr_version = (sys.maxsize, sys.maxsize) + + return self._dbr_version diff --git a/dbt/include/databricks/macros/adapters/metadata.sql b/dbt/include/databricks/macros/adapters/metadata.sql index fd79a58ce..53f382e25 100644 --- a/dbt/include/databricks/macros/adapters/metadata.sql +++ b/dbt/include/databricks/macros/adapters/metadata.sql @@ -19,9 +19,18 @@ {% endmacro %} {% macro databricks__show_tables(relation) %} - {% call statement('show_tables', fetch_result=True) -%} - show tables in {{ relation|lower }} - {% endcall %} + {% set database = (relation.database | default(''))| lower | replace('`', '') %} + {% set schema = relation.schema | lower | replace('`', '') %} + + {% if database and schema -%} + {% call statement('show_tables', fetch_result=True) -%} + SHOW TABLES IN {{ database }}.{{ schema }} + {% endcall %} + {% else -%} + {% call statement('show_tables', fetch_result=True) -%} + SHOW TABLES IN {{ relation | lower }} + {% endcall %} + {% endif %} {% do return(load_result('show_tables').table) %} {% endmacro %} @@ -103,4 +112,23 @@ {% endcall %} {% do return(load_result('get_uc_tables').table) %} +{% endmacro %} + +{% macro list_schemas(database) %} + {{ return(adapter.dispatch('list_schemas', 'dbt')(database)) }} +{% endmacro %} + +{% macro databricks__list_schemas(database) -%} + {% set database_clean = (database | default('')) | replace('`', '') %} + {% if database_clean -%} + {% call statement('list_schemas', fetch_result=True, auto_begin=False) %} + SHOW DATABASES IN {{ database_clean }} + {% endcall %} + {% else -%} + {% call statement('list_schemas', fetch_result=True, auto_begin=False) %} + SHOW DATABASES + {% endcall %} + {% endif -%} + + {{ return(load_result('list_schemas').table) }} {% endmacro %} \ No newline at end of file From 9ff849ae6b8e1d0dfe06a9948b545b11ea172861 Mon Sep 17 00:00:00 2001 From: dkruh1 <101973979+dkruh1@users.noreply.github.com> Date: Wed, 27 Nov 2024 17:10:43 +0200 Subject: [PATCH 2/9] configure tests --- .gitignore | 1 + dbt/adapters/databricks/connections.py | 2 +- tests/unit/test_adapter.py | 10 ++++++++++ tox.ini | 14 ++++++++++++++ 4 files changed, 26 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 3d42d669e..7edc3ffe4 100644 --- a/.gitignore +++ b/.gitignore @@ -19,3 +19,4 @@ logs/ .venv* *.sublime* .python-version +venv/ \ No newline at end of file diff --git a/dbt/adapters/databricks/connections.py b/dbt/adapters/databricks/connections.py index dd6aad8ef..fe82cfad9 100644 --- a/dbt/adapters/databricks/connections.py +++ b/dbt/adapters/databricks/connections.py @@ -1109,7 +1109,7 @@ def set_query_header(self, query_header_context: Dict[str, Any]) -> None: def set_connection_name( self, name: Optional[str] = None, query_header_context: Any = None ) -> Connection: - SparkConnectionManager.set_connection_name(self, name) + return SparkConnectionManager.set_connection_name(self, name) def add_query( self, diff --git a/tests/unit/test_adapter.py b/tests/unit/test_adapter.py index 5364cb158..3d1b72f2f 100644 --- a/tests/unit/test_adapter.py +++ b/tests/unit/test_adapter.py @@ -113,6 +113,7 @@ def test_http_headers(http_header): test_http_headers(["a", "b"]) test_http_headers({"a": 1, "b": 2}) + @pytest.mark.skip_profile("session_connection") def test_invalid_custom_user_agent(self): with pytest.raises(DbtValidationError) as excinfo: config = self._get_config() @@ -123,6 +124,7 @@ def test_invalid_custom_user_agent(self): assert "Invalid invocation environment" in str(excinfo.value) + @pytest.mark.skip_profile("session_connection") def test_custom_user_agent(self): config = self._get_config() adapter = DatabricksAdapter(config, get_context("spawn")) @@ -137,12 +139,14 @@ def test_custom_user_agent(self): connection = adapter.acquire_connection("dummy") connection.handle # trigger lazy-load + @pytest.mark.skip_profile("session_connection") def test_environment_single_http_header(self): self._test_environment_http_headers( http_headers_str='{"test":{"jobId":1,"runId":12123}}', expected_http_headers=[("test", '{"jobId": 1, "runId": 12123}')], ) + @pytest.mark.skip_profile("session_connection") def test_environment_multiple_http_headers(self): self._test_environment_http_headers( http_headers_str='{"test":{"jobId":1,"runId":12123},"dummy":{"jobId":1,"runId":12123}}', @@ -152,6 +156,7 @@ def test_environment_multiple_http_headers(self): ], ) + @pytest.mark.skip_profile("session_connection") def test_environment_users_http_headers_intersection_error(self): with pytest.raises(DbtValidationError) as excinfo: self._test_environment_http_headers( @@ -162,6 +167,7 @@ def test_environment_users_http_headers_intersection_error(self): assert "Intersection with reserved http_headers in keys: {'t'}" in str(excinfo.value) + @pytest.mark.skip_profile("session_connection") def test_environment_users_http_headers_union_success(self): self._test_environment_http_headers( http_headers_str='{"t":{"jobId":1,"runId":12123},"d":{"jobId":1,"runId":12123}}', @@ -173,6 +179,7 @@ def test_environment_users_http_headers_union_success(self): ], ) + @pytest.mark.skip_profile("session_connection") def test_environment_http_headers_string(self): self._test_environment_http_headers( http_headers_str='{"string":"some-string"}', @@ -272,6 +279,7 @@ def connect( return connect + @pytest.mark.skip_profile("session_connection") def test_databricks_sql_connector_connection(self): self._test_databricks_sql_connector_connection(self._connect_func()) @@ -294,6 +302,7 @@ def _test_databricks_sql_connector_connection(self, connect): assert len(connection.credentials.session_properties) == 1 assert connection.credentials.session_properties["spark.sql.ansi.enabled"] == "true" + @pytest.mark.skip_profile("session_connection") def test_databricks_sql_connector_catalog_connection(self): self._test_databricks_sql_connector_catalog_connection( self._connect_func(expected_catalog="main") @@ -317,6 +326,7 @@ def _test_databricks_sql_connector_catalog_connection(self, connect): assert connection.credentials.schema == "analytics" assert connection.credentials.database == "main" + @pytest.mark.skip_profile("session_connection") def test_databricks_sql_connector_http_header_connection(self): self._test_databricks_sql_connector_http_header_connection( {"aaa": "xxx"}, self._connect_func(expected_http_headers=[("aaa", "xxx")]) diff --git a/tox.ini b/tox.ini index 17fcde6de..7dac38b0b 100644 --- a/tox.ini +++ b/tox.ini @@ -27,6 +27,20 @@ deps = [testenv:unit] basepython = python3 commands = {envpython} -m pytest --color=yes -v {posargs} tests/unit +setenv = + DBT_DATABRICKS_SESSION_CONNECTION = False +passenv = + DBT_* + PYTEST_ADDOPTS +deps = + -r{toxinidir}/dev-requirements.txt + -r{toxinidir}/requirements.txt + +[testenv:unit-session] +basepython = python3 +commands = {envpython} -m pytest --color=yes -v {posargs} --profile session_connection tests/unit +setenv = + DBT_DATABRICKS_SESSION_CONNECTION = True passenv = DBT_* PYTEST_ADDOPTS From 891e91601ba4cd1c5aabbc70f102e1927490ba04 Mon Sep 17 00:00:00 2001 From: dkruh1 <101973979+dkruh1@users.noreply.github.com> Date: Thu, 28 Nov 2024 14:41:23 +0200 Subject: [PATCH 3/9] reformat files --- dbt/adapters/databricks/connections.py | 21 ++++++++++----------- dbt/adapters/databricks/impl.py | 9 +++++---- 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/dbt/adapters/databricks/connections.py b/dbt/adapters/databricks/connections.py index fe82cfad9..0817e4b58 100644 --- a/dbt/adapters/databricks/connections.py +++ b/dbt/adapters/databricks/connections.py @@ -1097,9 +1097,7 @@ def cancel_open(self) -> List[str]: def compare_dbr_version(self, major: int, minor: int) -> int: version = (major, minor) - connection = ( - self.get_thread_connection().handle - ) + connection = self.get_thread_connection().handle dbr_version = connection.dbr_version return (dbr_version > version) - (dbr_version < version) @@ -1120,23 +1118,24 @@ def add_query( *, close_cursor: bool = False, ) -> Tuple[Connection, Any]: - return SparkConnectionManager.add_query( - self, sql, auto_begin, bindings, abridge_sql_log - ) + return SparkConnectionManager.add_query(self, sql, auto_begin, bindings, abridge_sql_log) def list_schemas(self, database: str, schema: Optional[str] = None) -> "Table": - raise NotImplementedError("list_schemas is not implemented for DatabricksSessionConnectionManager - should call the list_schemas macro instead") + raise NotImplementedError( + "list_schemas is not implemented for DatabricksSessionConnectionManager - should call the list_schemas macro instead" + ) def list_tables(self, database: str, schema: str, identifier: Optional[str] = None) -> "Table": - raise NotImplementedError("list_tables is not implemented for DatabricksSessionConnectionManager - should call the list_tables macro instead") + raise NotImplementedError( + "list_tables is not implemented for DatabricksSessionConnectionManager - should call the list_tables macro instead" + ) @classmethod def open(cls, connection: Connection) -> Connection: from dbt.adapters.spark.session import Connection from dbt.adapters.databricks.session_connection import DatabricksSessionConnectionWrapper - handle = DatabricksSessionConnectionWrapper( - Connection() - ) + + handle = DatabricksSessionConnectionWrapper(Connection()) connection.handle = handle connection.state = ConnectionState.OPEN return connection diff --git a/dbt/adapters/databricks/impl.py b/dbt/adapters/databricks/impl.py index 93d9a8362..7218b881c 100644 --- a/dbt/adapters/databricks/impl.py +++ b/dbt/adapters/databricks/impl.py @@ -38,7 +38,10 @@ from dbt.adapters.contracts.relation import RelationConfig from dbt.adapters.contracts.relation import RelationType from dbt.adapters.databricks.column import DatabricksColumn -from dbt.adapters.databricks.connections import DatabricksConnectionManager, DatabricksSessionConnectionManager +from dbt.adapters.databricks.connections import ( + DatabricksConnectionManager, + DatabricksSessionConnectionManager, +) from dbt.adapters.databricks.connections import DatabricksDBTConnection from dbt.adapters.databricks.connections import DatabricksSQLConnectionWrapper from dbt.adapters.databricks.connections import ExtendedSessionConnectionManager @@ -212,9 +215,7 @@ def list_schemas(self, database: Optional[str]) -> List[str]: if database is not None and not USE_SESSION_CONNECTION: results = self.connections.list_schemas(database=database) else: - results = self.execute_macro( - LIST_SCHEMAS_MACRO_NAME, kwargs={"database": database} - ) + results = self.execute_macro(LIST_SCHEMAS_MACRO_NAME, kwargs={"database": database}) return [row[0] for row in results] def check_schema_exists(self, database: Optional[str], schema: str) -> bool: From 91d0cfc6228415a5f089643a3920fe472320720a Mon Sep 17 00:00:00 2001 From: dkruh1 <101973979+dkruh1@users.noreply.github.com> Date: Thu, 28 Nov 2024 18:53:46 +0200 Subject: [PATCH 4/9] fix pr comments --- dbt/adapters/databricks/connections.py | 4 +--- dbt/adapters/databricks/session_connection.py | 1 + 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/dbt/adapters/databricks/connections.py b/dbt/adapters/databricks/connections.py index 0817e4b58..3b391758c 100644 --- a/dbt/adapters/databricks/connections.py +++ b/dbt/adapters/databricks/connections.py @@ -70,9 +70,7 @@ from dbt.adapters.events.types import NewConnection from dbt.adapters.events.types import SQLQuery from dbt.adapters.events.types import SQLQueryStatus -from dbt.adapters.spark.connections import ( - SparkConnectionManager, -) +from dbt.adapters.spark.connections import SparkConnectionManager from dbt_common.events.contextvars import get_node_info from dbt_common.events.functions import fire_event from dbt_common.exceptions import DbtInternalError diff --git a/dbt/adapters/databricks/session_connection.py b/dbt/adapters/databricks/session_connection.py index a42e565fc..215f912c1 100644 --- a/dbt/adapters/databricks/session_connection.py +++ b/dbt/adapters/databricks/session_connection.py @@ -10,6 +10,7 @@ class DatabricksSessionConnectionWrapper(SessionConnectionWrapper): _is_cluster: bool + _dbr_version: Tuple[int, int] def __init__(self, handle: Connection) -> None: super().__init__(handle) From 987016c0ed62b0dbdfec23b1f5a18b22d2995f45 Mon Sep 17 00:00:00 2001 From: dkruh1 <101973979+dkruh1@users.noreply.github.com> Date: Sun, 1 Dec 2024 18:21:55 +0200 Subject: [PATCH 5/9] merge from upstream --- dbt/adapters/databricks/api_client.py | 9 +- dbt/adapters/databricks/auth.py | 15 ++- dbt/adapters/databricks/behaviors/columns.py | 3 +- dbt/adapters/databricks/column.py | 3 +- dbt/adapters/databricks/connections.py | 108 +++++++++--------- dbt/adapters/databricks/credentials.py | 24 ++-- .../databricks/events/connection_events.py | 3 +- dbt/adapters/databricks/impl.py | 101 ++++++++-------- dbt/adapters/databricks/logging.py | 4 +- .../databricks/python_models/python_config.py | 6 +- .../python_models/python_submissions.py | 10 +- .../databricks/python_models/run_tracking.py | 3 +- dbt/adapters/databricks/relation.py | 14 ++- .../databricks/relation_configs/base.py | 14 ++- .../databricks/relation_configs/comment.py | 9 +- .../relation_configs/incremental.py | 8 +- .../relation_configs/materialized_view.py | 8 +- .../relation_configs/partitioning.py | 9 +- .../databricks/relation_configs/query.py | 6 +- .../databricks/relation_configs/refresh.py | 9 +- .../relation_configs/streaming_table.py | 11 +- .../databricks/relation_configs/tags.py | 9 +- .../relation_configs/tblproperties.py | 10 +- dbt/adapters/databricks/session_connection.py | 3 +- dbt/adapters/databricks/utils.py | 4 +- .../databricks/macros/adapters/columns.sql | 2 +- .../databricks/macros/adapters/metadata.sql | 2 +- .../incremental/strategies.sql | 4 +- .../macros/materializations/seeds/helpers.sql | 2 +- .../macros/relations/constraints.sql | 6 +- .../macros/relations/liquid_clustering.sql | 2 +- .../databricks/macros/relations/tags.sql | 2 +- .../databricks/macros/utils/split_part.sql | 2 +- pyproject.toml | 24 +++- .../adapter/aliases/test_aliases.py | 10 +- .../adapter/basic/test_docs_generate.py | 6 +- .../adapter/basic/test_incremental.py | 3 +- .../adapter/caching/test_caching.py | 10 +- .../adapter/concurrency/test_concurrency.py | 3 +- .../adapter/constraints/test_constraints.py | 18 ++- .../adapter/dbt_clone/test_dbt_clone.py | 3 +- .../adapter/dbt_show/test_dbt_show.py | 3 +- .../adapter/ephemeral/test_ephemeral.py | 3 +- .../adapter/hooks/test_run_hooks.py | 3 +- .../materialized_view_tests/test_changes.py | 10 +- .../test_persist_constraints.py | 3 +- .../adapter/python_model/test_python_model.py | 6 +- .../query_comments/test_query_comment.py | 14 +-- .../adapter/simple_copy/test_simple_copy.py | 3 +- .../adapter/simple_seed/test_seeds.py | 24 ++-- .../adapter/simple_snapshot/test_snapshot.py | 3 +- .../test_store_failures_as.py | 14 +-- .../adapter/streaming_tables/test_st_basic.py | 4 +- .../adapter/utils/test_source_freshness.py | 3 +- tests/functional/adapter/utils/test_utils.py | 3 +- tests/profiles.py | 3 +- tests/unit/api_client/api_test_base.py | 3 +- tests/unit/api_client/test_command_api.py | 6 +- tests/unit/api_client/test_user_folder_api.py | 3 +- tests/unit/events/test_base_events.py | 3 +- tests/unit/events/test_connection_events.py | 8 +- tests/unit/events/test_cursor_events.py | 3 +- tests/unit/macros/base.py | 5 +- tests/unit/python/test_python_config.py | 8 +- tests/unit/python/test_python_helpers.py | 16 ++- tests/unit/python/test_python_job_support.py | 8 +- tests/unit/python/test_python_submitters.py | 20 ++-- .../python/test_python_workflow_support.py | 6 +- tests/unit/relation_configs/test_comment.py | 3 +- .../test_materialized_view_config.py | 3 +- .../relation_configs/test_partitioning.py | 6 +- tests/unit/relation_configs/test_query.py | 3 +- tests/unit/relation_configs/test_refresh.py | 3 +- tests/unit/relation_configs/test_tags.py | 3 +- .../relation_configs/test_tblproperties.py | 6 +- tests/unit/test_adapter.py | 23 ++-- tests/unit/test_idle_config.py | 3 +- tests/unit/test_relation.py | 3 +- tests/unit/test_utils.py | 3 +- tests/unit/utils.py | 7 +- 80 files changed, 412 insertions(+), 344 deletions(-) diff --git a/dbt/adapters/databricks/api_client.py b/dbt/adapters/databricks/api_client.py index c14af9f53..aa86d7402 100644 --- a/dbt/adapters/databricks/api_client.py +++ b/dbt/adapters/databricks/api_client.py @@ -1,13 +1,16 @@ import base64 import re import time -from abc import ABC, abstractmethod +from abc import ABC +from abc import abstractmethod from collections.abc import Callable from dataclasses import dataclass -from typing import Any, Optional +from typing import Any +from typing import Optional from dbt_common.exceptions import DbtRuntimeError -from requests import Response, Session +from requests import Response +from requests import Session from requests.adapters import HTTPAdapter from urllib3.util.retry import Retry diff --git a/dbt/adapters/databricks/auth.py b/dbt/adapters/databricks/auth.py index 439fa23fd..8662f794d 100644 --- a/dbt/adapters/databricks/auth.py +++ b/dbt/adapters/databricks/auth.py @@ -1,11 +1,16 @@ -from typing import Any, Optional - +from typing import Any +from typing import Optional + +from databricks.sdk.core import Config +from databricks.sdk.core import credentials_provider +from databricks.sdk.core import CredentialsProvider +from databricks.sdk.core import HeaderFactory +from databricks.sdk.oauth import ClientCredentials +from databricks.sdk.oauth import Token +from databricks.sdk.oauth import TokenSource from requests import PreparedRequest from requests.auth import AuthBase -from databricks.sdk.core import Config, CredentialsProvider, HeaderFactory, credentials_provider -from databricks.sdk.oauth import ClientCredentials, Token, TokenSource - class token_auth(CredentialsProvider): _token: str diff --git a/dbt/adapters/databricks/behaviors/columns.py b/dbt/adapters/databricks/behaviors/columns.py index 8f3d1eca7..fc72e66ec 100644 --- a/dbt/adapters/databricks/behaviors/columns.py +++ b/dbt/adapters/databricks/behaviors/columns.py @@ -1,4 +1,5 @@ -from abc import ABC, abstractmethod +from abc import ABC +from abc import abstractmethod from dbt_common.utils.dict import AttrDict diff --git a/dbt/adapters/databricks/column.py b/dbt/adapters/databricks/column.py index 4d08ad4d9..f9511201c 100644 --- a/dbt/adapters/databricks/column.py +++ b/dbt/adapters/databricks/column.py @@ -1,5 +1,6 @@ from dataclasses import dataclass -from typing import ClassVar, Optional +from typing import ClassVar +from typing import Optional from dbt.adapters.spark.column import SparkColumn diff --git a/dbt/adapters/databricks/connections.py b/dbt/adapters/databricks/connections.py index eb6660576..06636c1ec 100644 --- a/dbt/adapters/databricks/connections.py +++ b/dbt/adapters/databricks/connections.py @@ -5,75 +5,77 @@ import time import uuid import warnings -from collections.abc import Callable, Iterator, Sequence +from collections.abc import Callable +from collections.abc import Iterator +from collections.abc import Sequence from contextlib import contextmanager from dataclasses import dataclass from multiprocessing.context import SpawnContext from numbers import Number from threading import get_ident -from typing import TYPE_CHECKING, Any, Hashable, Optional, cast +from typing import Any +from typing import cast +from typing import Hashable +from typing import Optional +from typing import TYPE_CHECKING +import databricks.sql as dbsql +from databricks.sql.client import Connection as DatabricksSQLConnection +from databricks.sql.client import Cursor as DatabricksSQLCursor +from databricks.sql.exc import Error from dbt_common.events.contextvars import get_node_info from dbt_common.events.functions import fire_event -from dbt_common.exceptions import DbtDatabaseError, DbtInternalError, DbtRuntimeError +from dbt_common.exceptions import DbtDatabaseError +from dbt_common.exceptions import DbtInternalError +from dbt_common.exceptions import DbtRuntimeError from dbt_common.utils import cast_to_str from requests import Session -import databricks.sql as dbsql -from databricks.sql.client import Connection as DatabricksSQLConnection -from databricks.sql.client import Cursor as DatabricksSQLCursor -from databricks.sql.exc import Error from dbt.adapters.base.query_headers import MacroQueryStringSetter -from dbt.adapters.contracts.connection import ( - DEFAULT_QUERY_COMMENT, - AdapterRequiredConfig, - AdapterResponse, - Connection, - ConnectionState, - Identifier, - LazyHandle, -) +from dbt.adapters.contracts.connection import AdapterRequiredConfig +from dbt.adapters.contracts.connection import AdapterResponse +from dbt.adapters.contracts.connection import Connection +from dbt.adapters.contracts.connection import ConnectionState +from dbt.adapters.contracts.connection import DEFAULT_QUERY_COMMENT +from dbt.adapters.contracts.connection import Identifier +from dbt.adapters.contracts.connection import LazyHandle from dbt.adapters.databricks.__version__ import version as __version__ from dbt.adapters.databricks.api_client import DatabricksApiClient from dbt.adapters.databricks.auth import BearerAuth -from dbt.adapters.databricks.credentials import DatabricksCredentials, TCredentialProvider -from dbt.adapters.databricks.events.connection_events import ( - ConnectionAcquire, - ConnectionCancel, - ConnectionCancelError, - ConnectionClose, - ConnectionCloseError, - ConnectionCreate, - ConnectionCreated, - ConnectionCreateError, - ConnectionIdleCheck, - ConnectionIdleClose, - ConnectionRelease, - ConnectionReset, - ConnectionRetrieve, - ConnectionReuse, -) -from dbt.adapters.databricks.events.cursor_events import ( - CursorCancel, - CursorCancelError, - CursorClose, - CursorCloseError, - CursorCreate, -) +from dbt.adapters.databricks.credentials import DatabricksCredentials +from dbt.adapters.databricks.credentials import TCredentialProvider +from dbt.adapters.databricks.events.connection_events import ConnectionAcquire +from dbt.adapters.databricks.events.connection_events import ConnectionCancel +from dbt.adapters.databricks.events.connection_events import ConnectionCancelError +from dbt.adapters.databricks.events.connection_events import ConnectionClose +from dbt.adapters.databricks.events.connection_events import ConnectionCloseError +from dbt.adapters.databricks.events.connection_events import ConnectionCreate +from dbt.adapters.databricks.events.connection_events import ConnectionCreated +from dbt.adapters.databricks.events.connection_events import ConnectionCreateError +from dbt.adapters.databricks.events.connection_events import ConnectionIdleCheck +from dbt.adapters.databricks.events.connection_events import ConnectionIdleClose +from dbt.adapters.databricks.events.connection_events import ConnectionRelease +from dbt.adapters.databricks.events.connection_events import ConnectionReset +from dbt.adapters.databricks.events.connection_events import ConnectionRetrieve +from dbt.adapters.databricks.events.connection_events import ConnectionReuse +from dbt.adapters.databricks.events.cursor_events import CursorCancel +from dbt.adapters.databricks.events.cursor_events import CursorCancelError +from dbt.adapters.databricks.events.cursor_events import CursorClose +from dbt.adapters.databricks.events.cursor_events import CursorCloseError +from dbt.adapters.databricks.events.cursor_events import CursorCreate from dbt.adapters.databricks.events.other_events import QueryError -from dbt.adapters.databricks.events.pipeline_events import PipelineRefresh, PipelineRefreshError +from dbt.adapters.databricks.events.pipeline_events import PipelineRefresh +from dbt.adapters.databricks.events.pipeline_events import PipelineRefreshError from dbt.adapters.databricks.logging import logger from dbt.adapters.databricks.python_models.run_tracking import PythonRunTracker from dbt.adapters.databricks.utils import redact_credentials -from dbt.adapters.events.types import ( - ConnectionClosedInCleanup, - ConnectionLeftOpenInCleanup, - ConnectionReused, - ConnectionUsed, - NewConnection, - SQLQuery, - SQLQueryStatus, -) +from dbt.adapters.events.types import ConnectionClosedInCleanup +from dbt.adapters.events.types import ConnectionLeftOpenInCleanup +from dbt.adapters.events.types import ConnectionReused +from dbt.adapters.events.types import ConnectionUsed +from dbt.adapters.events.types import NewConnection +from dbt.adapters.events.types import SQLQuery +from dbt.adapters.events.types import SQLQueryStatus from dbt.adapters.spark.connections import SparkConnectionManager if TYPE_CHECKING: @@ -1112,12 +1114,14 @@ def add_query( def list_schemas(self, database: str, schema: Optional[str] = None) -> "Table": raise NotImplementedError( - "list_schemas is not implemented for DatabricksSessionConnectionManager - should call the list_schemas macro instead" + "list_schemas is not implemented for DatabricksSessionConnectionManager - " + + "should call the list_schemas macro instead" ) def list_tables(self, database: str, schema: str, identifier: Optional[str] = None) -> "Table": raise NotImplementedError( - "list_tables is not implemented for DatabricksSessionConnectionManager - should call the list_tables macro instead" + "list_tables is not implemented for DatabricksSessionConnectionManager - " + + "should call the list_tables macro instead" ) @classmethod diff --git a/dbt/adapters/databricks/credentials.py b/dbt/adapters/databricks/credentials.py index 0f729a5b1..30f058091 100644 --- a/dbt/adapters/databricks/credentials.py +++ b/dbt/adapters/databricks/credentials.py @@ -5,20 +5,24 @@ import threading from collections.abc import Iterable from dataclasses import dataclass -from typing import Any, Optional, Union, cast +from typing import Any +from typing import cast +from typing import Optional +from typing import Union import keyring -from dbt_common.exceptions import DbtConfigError, DbtValidationError - from databricks.sdk.core import CredentialsProvider -from databricks.sdk.oauth import OAuthClient, SessionCredentials +from databricks.sdk.oauth import OAuthClient +from databricks.sdk.oauth import SessionCredentials +from dbt_common.exceptions import DbtConfigError +from dbt_common.exceptions import DbtValidationError + from dbt.adapters.contracts.connection import Credentials -from dbt.adapters.databricks.auth import m2m_auth, token_auth -from dbt.adapters.databricks.events.credential_events import ( - CredentialLoadError, - CredentialSaveError, - CredentialShardEvent, -) +from dbt.adapters.databricks.auth import m2m_auth +from dbt.adapters.databricks.auth import token_auth +from dbt.adapters.databricks.events.credential_events import CredentialLoadError +from dbt.adapters.databricks.events.credential_events import CredentialSaveError +from dbt.adapters.databricks.events.credential_events import CredentialShardEvent from dbt.adapters.databricks.logging import logger CATALOG_KEY_IN_SESSION_PROPERTIES = "databricks.catalog" diff --git a/dbt/adapters/databricks/events/connection_events.py b/dbt/adapters/databricks/events/connection_events.py index 9f8ec8c1b..0b77c2a16 100644 --- a/dbt/adapters/databricks/events/connection_events.py +++ b/dbt/adapters/databricks/events/connection_events.py @@ -1,5 +1,6 @@ from abc import ABC -from typing import Any, Optional +from typing import Any +from typing import Optional from databricks.sql.client import Connection diff --git a/dbt/adapters/databricks/impl.py b/dbt/adapters/databricks/impl.py index 4819eff69..8ce1de93c 100644 --- a/dbt/adapters/databricks/impl.py +++ b/dbt/adapters/databricks/impl.py @@ -1,60 +1,67 @@ import os import re -from abc import ABC, abstractmethod +from abc import ABC +from abc import abstractmethod from collections import defaultdict -from collections.abc import Iterable, Iterator +from collections.abc import Iterable +from collections.abc import Iterator from concurrent.futures import Future from contextlib import contextmanager from dataclasses import dataclass from importlib import metadata from multiprocessing.context import SpawnContext -from typing import TYPE_CHECKING, Any, ClassVar, Generic, Optional, Union, cast +from typing import Any +from typing import cast +from typing import ClassVar +from typing import Generic +from typing import Optional +from typing import TYPE_CHECKING +from typing import Union from uuid import uuid4 from dbt_common.behavior_flags import BehaviorFlag from dbt_common.contracts.config.base import BaseConfig -from dbt_common.exceptions import CompilationError, DbtConfigError, DbtInternalError +from dbt_common.exceptions import CompilationError +from dbt_common.exceptions import DbtConfigError +from dbt_common.exceptions import DbtInternalError from dbt_common.utils import executor from dbt_common.utils.dict import AttrDict from packaging import version -from dbt.adapters.base import AdapterConfig, PythonJobHelper -from dbt.adapters.base.impl import catch_as_completed, log_code_execution +from dbt.adapters.base import AdapterConfig +from dbt.adapters.base import PythonJobHelper +from dbt.adapters.base.impl import catch_as_completed +from dbt.adapters.base.impl import log_code_execution from dbt.adapters.base.meta import available from dbt.adapters.base.relation import BaseRelation -from dbt.adapters.capability import Capability, CapabilityDict, CapabilitySupport, Support -from dbt.adapters.contracts.connection import AdapterResponse, Connection -from dbt.adapters.contracts.relation import RelationConfig, RelationType -from dbt.adapters.databricks.behaviors.columns import ( - GetColumnsBehavior, - GetColumnsByDescribe, - GetColumnsByInformationSchema, -) +from dbt.adapters.capability import Capability +from dbt.adapters.capability import CapabilityDict +from dbt.adapters.capability import CapabilitySupport +from dbt.adapters.capability import Support +from dbt.adapters.contracts.connection import AdapterResponse +from dbt.adapters.contracts.connection import Connection +from dbt.adapters.contracts.relation import RelationConfig +from dbt.adapters.contracts.relation import RelationType +from dbt.adapters.databricks.behaviors.columns import GetColumnsBehavior +from dbt.adapters.databricks.behaviors.columns import GetColumnsByDescribe +from dbt.adapters.databricks.behaviors.columns import GetColumnsByInformationSchema from dbt.adapters.databricks.column import DatabricksColumn -from dbt.adapters.databricks.connections import ( - USE_LONG_SESSIONS, - USE_SESSION_CONNECTION, - DatabricksConnectionManager, - DatabricksDBTConnection, - DatabricksSessionConnectionManager, - DatabricksSQLConnectionWrapper, - ExtendedSessionConnectionManager, -) -from dbt.adapters.databricks.python_models.python_submissions import ( - AllPurposeClusterPythonJobHelper, - JobClusterPythonJobHelper, - ServerlessClusterPythonJobHelper, - WorkflowPythonJobHelper, -) -from dbt.adapters.databricks.relation import ( - KEY_TABLE_PROVIDER, - DatabricksRelation, - DatabricksRelationType, -) -from dbt.adapters.databricks.relation_configs.base import ( - DatabricksRelationConfig, - DatabricksRelationConfigBase, -) +from dbt.adapters.databricks.connections import DatabricksConnectionManager +from dbt.adapters.databricks.connections import DatabricksDBTConnection +from dbt.adapters.databricks.connections import DatabricksSessionConnectionManager +from dbt.adapters.databricks.connections import DatabricksSQLConnectionWrapper +from dbt.adapters.databricks.connections import ExtendedSessionConnectionManager +from dbt.adapters.databricks.connections import USE_LONG_SESSIONS +from dbt.adapters.databricks.connections import USE_SESSION_CONNECTION +from dbt.adapters.databricks.python_models.python_submissions import AllPurposeClusterPythonJobHelper +from dbt.adapters.databricks.python_models.python_submissions import JobClusterPythonJobHelper +from dbt.adapters.databricks.python_models.python_submissions import ServerlessClusterPythonJobHelper +from dbt.adapters.databricks.python_models.python_submissions import WorkflowPythonJobHelper +from dbt.adapters.databricks.relation import DatabricksRelation +from dbt.adapters.databricks.relation import DatabricksRelationType +from dbt.adapters.databricks.relation import KEY_TABLE_PROVIDER +from dbt.adapters.databricks.relation_configs.base import DatabricksRelationConfig +from dbt.adapters.databricks.relation_configs.base import DatabricksRelationConfigBase from dbt.adapters.databricks.relation_configs.incremental import IncrementalTableConfig from dbt.adapters.databricks.relation_configs.materialized_view import ( MaterializedViewConfig, @@ -64,16 +71,16 @@ ) from dbt.adapters.databricks.relation_configs.table_format import TableFormat from dbt.adapters.databricks.relation_configs.tblproperties import TblPropertiesConfig -from dbt.adapters.databricks.utils import get_first_row, handle_missing_objects, redact_credentials +from dbt.adapters.databricks.utils import get_first_row +from dbt.adapters.databricks.utils import handle_missing_objects +from dbt.adapters.databricks.utils import redact_credentials from dbt.adapters.relation_configs import RelationResults -from dbt.adapters.spark.impl import ( - DESCRIBE_TABLE_EXTENDED_MACRO_NAME, - GET_COLUMNS_IN_RELATION_RAW_MACRO_NAME, - KEY_TABLE_OWNER, - KEY_TABLE_STATISTICS, - LIST_SCHEMAS_MACRO_NAME, - SparkAdapter, -) +from dbt.adapters.spark.impl import DESCRIBE_TABLE_EXTENDED_MACRO_NAME +from dbt.adapters.spark.impl import GET_COLUMNS_IN_RELATION_RAW_MACRO_NAME +from dbt.adapters.spark.impl import KEY_TABLE_OWNER +from dbt.adapters.spark.impl import KEY_TABLE_STATISTICS +from dbt.adapters.spark.impl import LIST_SCHEMAS_MACRO_NAME +from dbt.adapters.spark.impl import SparkAdapter if TYPE_CHECKING: from agate import Row, Table diff --git a/dbt/adapters/databricks/logging.py b/dbt/adapters/databricks/logging.py index d0f1d42ba..c309e12c7 100644 --- a/dbt/adapters/databricks/logging.py +++ b/dbt/adapters/databricks/logging.py @@ -1,5 +1,7 @@ import os -from logging import Handler, LogRecord, getLogger +from logging import getLogger +from logging import Handler +from logging import LogRecord from typing import Union from dbt.adapters.events.logging import AdapterLogger diff --git a/dbt/adapters/databricks/python_models/python_config.py b/dbt/adapters/databricks/python_models/python_config.py index 29aa44efa..761b0ef65 100644 --- a/dbt/adapters/databricks/python_models/python_config.py +++ b/dbt/adapters/databricks/python_models/python_config.py @@ -1,7 +1,9 @@ import uuid -from typing import Any, Optional +from typing import Any +from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field DEFAULT_TIMEOUT = 60 * 60 * 24 diff --git a/dbt/adapters/databricks/python_models/python_submissions.py b/dbt/adapters/databricks/python_models/python_submissions.py index afcb383c2..b241fb358 100644 --- a/dbt/adapters/databricks/python_models/python_submissions.py +++ b/dbt/adapters/databricks/python_models/python_submissions.py @@ -1,12 +1,16 @@ -from abc import ABC, abstractmethod -from typing import Any, Optional +from abc import ABC +from abc import abstractmethod +from typing import Any +from typing import Optional from dbt_common.exceptions import DbtRuntimeError from pydantic import BaseModel from typing_extensions import override from dbt.adapters.base import PythonJobHelper -from dbt.adapters.databricks.api_client import CommandExecution, DatabricksApiClient, WorkflowJobApi +from dbt.adapters.databricks.api_client import CommandExecution +from dbt.adapters.databricks.api_client import DatabricksApiClient +from dbt.adapters.databricks.api_client import WorkflowJobApi from dbt.adapters.databricks.credentials import DatabricksCredentials from dbt.adapters.databricks.logging import logger from dbt.adapters.databricks.python_models.python_config import ParsedPythonModel diff --git a/dbt/adapters/databricks/python_models/run_tracking.py b/dbt/adapters/databricks/python_models/run_tracking.py index e48dae7d4..ae703cdd1 100644 --- a/dbt/adapters/databricks/python_models/run_tracking.py +++ b/dbt/adapters/databricks/python_models/run_tracking.py @@ -2,7 +2,8 @@ from dbt_common.exceptions import DbtRuntimeError -from dbt.adapters.databricks.api_client import CommandExecution, DatabricksApiClient +from dbt.adapters.databricks.api_client import CommandExecution +from dbt.adapters.databricks.api_client import DatabricksApiClient from dbt.adapters.databricks.logging import logger diff --git a/dbt/adapters/databricks/relation.py b/dbt/adapters/databricks/relation.py index bc1dc991f..ac6189384 100644 --- a/dbt/adapters/databricks/relation.py +++ b/dbt/adapters/databricks/relation.py @@ -1,17 +1,23 @@ from collections.abc import Iterable -from dataclasses import dataclass, field -from typing import Any, Optional, Type +from dataclasses import dataclass +from dataclasses import field +from typing import Any +from typing import Optional +from typing import Type from dbt_common.dataclass_schema import StrEnum from dbt_common.exceptions import DbtRuntimeError from dbt_common.utils import filter_null_values -from dbt.adapters.base.relation import BaseRelation, InformationSchema, Policy +from dbt.adapters.base.relation import BaseRelation +from dbt.adapters.base.relation import InformationSchema +from dbt.adapters.base.relation import Policy from dbt.adapters.contracts.relation import ( ComponentName, ) from dbt.adapters.databricks.utils import remove_undefined -from dbt.adapters.spark.impl import KEY_TABLE_OWNER, KEY_TABLE_STATISTICS +from dbt.adapters.spark.impl import KEY_TABLE_OWNER +from dbt.adapters.spark.impl import KEY_TABLE_STATISTICS from dbt.adapters.utils import classproperty KEY_TABLE_PROVIDER = "Provider" diff --git a/dbt/adapters/databricks/relation_configs/base.py b/dbt/adapters/databricks/relation_configs/base.py index 53e5bdb8d..9366226c3 100644 --- a/dbt/adapters/databricks/relation_configs/base.py +++ b/dbt/adapters/databricks/relation_configs/base.py @@ -1,7 +1,13 @@ -from abc import ABC, abstractmethod -from typing import Any, ClassVar, Generic, Optional, TypeVar - -from pydantic import BaseModel, ConfigDict +from abc import ABC +from abc import abstractmethod +from typing import Any +from typing import ClassVar +from typing import Generic +from typing import Optional +from typing import TypeVar + +from pydantic import BaseModel +from pydantic import ConfigDict from typing_extensions import Self from dbt.adapters.contracts.relation import RelationConfig diff --git a/dbt/adapters/databricks/relation_configs/comment.py b/dbt/adapters/databricks/relation_configs/comment.py index 3dfa07bc1..0af362d1a 100644 --- a/dbt/adapters/databricks/relation_configs/comment.py +++ b/dbt/adapters/databricks/relation_configs/comment.py @@ -1,10 +1,9 @@ -from typing import ClassVar, Optional +from typing import ClassVar +from typing import Optional from dbt.adapters.contracts.relation import RelationConfig -from dbt.adapters.databricks.relation_configs.base import ( - DatabricksComponentConfig, - DatabricksComponentProcessor, -) +from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig +from dbt.adapters.databricks.relation_configs.base import DatabricksComponentProcessor from dbt.adapters.relation_configs.config_base import RelationResults diff --git a/dbt/adapters/databricks/relation_configs/incremental.py b/dbt/adapters/databricks/relation_configs/incremental.py index 156f5bdf7..f31baf581 100644 --- a/dbt/adapters/databricks/relation_configs/incremental.py +++ b/dbt/adapters/databricks/relation_configs/incremental.py @@ -1,10 +1,8 @@ from typing import Optional -from dbt.adapters.databricks.relation_configs.base import ( - DatabricksComponentConfig, - DatabricksRelationChangeSet, - DatabricksRelationConfigBase, -) +from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig +from dbt.adapters.databricks.relation_configs.base import DatabricksRelationChangeSet +from dbt.adapters.databricks.relation_configs.base import DatabricksRelationConfigBase from dbt.adapters.databricks.relation_configs.tags import TagsProcessor from dbt.adapters.databricks.relation_configs.tblproperties import TblPropertiesProcessor diff --git a/dbt/adapters/databricks/relation_configs/materialized_view.py b/dbt/adapters/databricks/relation_configs/materialized_view.py index e6a8a47e8..9a57edd9f 100644 --- a/dbt/adapters/databricks/relation_configs/materialized_view.py +++ b/dbt/adapters/databricks/relation_configs/materialized_view.py @@ -1,10 +1,8 @@ from typing import Optional -from dbt.adapters.databricks.relation_configs.base import ( - DatabricksComponentConfig, - DatabricksRelationChangeSet, - DatabricksRelationConfigBase, -) +from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig +from dbt.adapters.databricks.relation_configs.base import DatabricksRelationChangeSet +from dbt.adapters.databricks.relation_configs.base import DatabricksRelationConfigBase from dbt.adapters.databricks.relation_configs.comment import ( CommentProcessor, ) diff --git a/dbt/adapters/databricks/relation_configs/partitioning.py b/dbt/adapters/databricks/relation_configs/partitioning.py index bba4b5720..89120b8ac 100644 --- a/dbt/adapters/databricks/relation_configs/partitioning.py +++ b/dbt/adapters/databricks/relation_configs/partitioning.py @@ -1,12 +1,11 @@ import itertools -from typing import ClassVar, Union +from typing import ClassVar +from typing import Union from dbt.adapters.contracts.relation import RelationConfig from dbt.adapters.databricks.relation_configs import base -from dbt.adapters.databricks.relation_configs.base import ( - DatabricksComponentConfig, - DatabricksComponentProcessor, -) +from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig +from dbt.adapters.databricks.relation_configs.base import DatabricksComponentProcessor from dbt.adapters.relation_configs.config_base import RelationResults diff --git a/dbt/adapters/databricks/relation_configs/query.py b/dbt/adapters/databricks/relation_configs/query.py index e9b8f73b3..469a117c6 100644 --- a/dbt/adapters/databricks/relation_configs/query.py +++ b/dbt/adapters/databricks/relation_configs/query.py @@ -3,10 +3,8 @@ from dbt_common.exceptions import DbtRuntimeError from dbt.adapters.contracts.relation import RelationConfig -from dbt.adapters.databricks.relation_configs.base import ( - DatabricksComponentConfig, - DatabricksComponentProcessor, -) +from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig +from dbt.adapters.databricks.relation_configs.base import DatabricksComponentProcessor from dbt.adapters.relation_configs.config_base import RelationResults diff --git a/dbt/adapters/databricks/relation_configs/refresh.py b/dbt/adapters/databricks/relation_configs/refresh.py index bfea49e16..c3637297c 100644 --- a/dbt/adapters/databricks/relation_configs/refresh.py +++ b/dbt/adapters/databricks/relation_configs/refresh.py @@ -1,14 +1,13 @@ import re -from typing import ClassVar, Optional +from typing import ClassVar +from typing import Optional from dbt_common.exceptions import DbtRuntimeError from dbt.adapters.contracts.relation import RelationConfig from dbt.adapters.databricks.relation_configs import base -from dbt.adapters.databricks.relation_configs.base import ( - DatabricksComponentConfig, - DatabricksComponentProcessor, -) +from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig +from dbt.adapters.databricks.relation_configs.base import DatabricksComponentProcessor from dbt.adapters.relation_configs.config_base import RelationResults SCHEDULE_REGEX = re.compile(r"CRON '(.*)' AT TIME ZONE '(.*)'") diff --git a/dbt/adapters/databricks/relation_configs/streaming_table.py b/dbt/adapters/databricks/relation_configs/streaming_table.py index f07ed3669..d891da13d 100644 --- a/dbt/adapters/databricks/relation_configs/streaming_table.py +++ b/dbt/adapters/databricks/relation_configs/streaming_table.py @@ -1,17 +1,16 @@ from typing import Optional -from dbt.adapters.databricks.relation_configs.base import ( - DatabricksComponentConfig, - DatabricksRelationChangeSet, - DatabricksRelationConfigBase, -) +from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig +from dbt.adapters.databricks.relation_configs.base import DatabricksRelationChangeSet +from dbt.adapters.databricks.relation_configs.base import DatabricksRelationConfigBase from dbt.adapters.databricks.relation_configs.comment import ( CommentProcessor, ) from dbt.adapters.databricks.relation_configs.partitioning import ( PartitionedByProcessor, ) -from dbt.adapters.databricks.relation_configs.refresh import RefreshConfig, RefreshProcessor +from dbt.adapters.databricks.relation_configs.refresh import RefreshConfig +from dbt.adapters.databricks.relation_configs.refresh import RefreshProcessor from dbt.adapters.databricks.relation_configs.tblproperties import ( TblPropertiesProcessor, ) diff --git a/dbt/adapters/databricks/relation_configs/tags.py b/dbt/adapters/databricks/relation_configs/tags.py index 461fa77b2..f8c734ac8 100644 --- a/dbt/adapters/databricks/relation_configs/tags.py +++ b/dbt/adapters/databricks/relation_configs/tags.py @@ -1,13 +1,12 @@ -from typing import ClassVar, Optional +from typing import ClassVar +from typing import Optional from dbt_common.exceptions import DbtRuntimeError from dbt.adapters.contracts.relation import RelationConfig from dbt.adapters.databricks.relation_configs import base -from dbt.adapters.databricks.relation_configs.base import ( - DatabricksComponentConfig, - DatabricksComponentProcessor, -) +from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig +from dbt.adapters.databricks.relation_configs.base import DatabricksComponentProcessor from dbt.adapters.relation_configs.config_base import RelationResults diff --git a/dbt/adapters/databricks/relation_configs/tblproperties.py b/dbt/adapters/databricks/relation_configs/tblproperties.py index 41d18f3fd..84e9b8e50 100644 --- a/dbt/adapters/databricks/relation_configs/tblproperties.py +++ b/dbt/adapters/databricks/relation_configs/tblproperties.py @@ -1,13 +1,13 @@ -from typing import Any, ClassVar, Optional +from typing import Any +from typing import ClassVar +from typing import Optional from dbt_common.exceptions import DbtRuntimeError from dbt.adapters.contracts.relation import RelationConfig from dbt.adapters.databricks.relation_configs import base -from dbt.adapters.databricks.relation_configs.base import ( - DatabricksComponentConfig, - DatabricksComponentProcessor, -) +from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig +from dbt.adapters.databricks.relation_configs.base import DatabricksComponentProcessor from dbt.adapters.relation_configs.config_base import RelationResults diff --git a/dbt/adapters/databricks/session_connection.py b/dbt/adapters/databricks/session_connection.py index 64a6ff8f8..9fc998406 100644 --- a/dbt/adapters/databricks/session_connection.py +++ b/dbt/adapters/databricks/session_connection.py @@ -2,7 +2,8 @@ import sys from typing import Tuple -from dbt.adapters.spark.session import Connection, SessionConnectionWrapper +from dbt.adapters.spark.session import Connection +from dbt.adapters.spark.session import SessionConnectionWrapper DBR_VERSION_REGEX = re.compile(r"([1-9][0-9]*)\.(x|0|[1-9][0-9]*)") diff --git a/dbt/adapters/databricks/utils.py b/dbt/adapters/databricks/utils.py index 616e5b368..6ba06331b 100644 --- a/dbt/adapters/databricks/utils.py +++ b/dbt/adapters/databricks/utils.py @@ -1,6 +1,8 @@ import re from collections.abc import Callable -from typing import TYPE_CHECKING, Any, TypeVar +from typing import Any +from typing import TYPE_CHECKING +from typing import TypeVar from dbt_common.exceptions import DbtRuntimeError from jinja2 import Undefined diff --git a/dbt/include/databricks/macros/adapters/columns.sql b/dbt/include/databricks/macros/adapters/columns.sql index e1fc1d116..1fc64ebf5 100644 --- a/dbt/include/databricks/macros/adapters/columns.sql +++ b/dbt/include/databricks/macros/adapters/columns.sql @@ -19,7 +19,7 @@ from `system`.`information_schema`.`columns` where table_catalog = '{{ relation.database|lower }}' and - table_schema = '{{ relation.schema|lower }}' and + table_schema = '{{ relation.schema|lower }}' and table_name = '{{ relation.identifier|lower }}' {% endcall %} diff --git a/dbt/include/databricks/macros/adapters/metadata.sql b/dbt/include/databricks/macros/adapters/metadata.sql index 53f382e25..355638736 100644 --- a/dbt/include/databricks/macros/adapters/metadata.sql +++ b/dbt/include/databricks/macros/adapters/metadata.sql @@ -131,4 +131,4 @@ {% endif -%} {{ return(load_result('list_schemas').table) }} -{% endmacro %} \ No newline at end of file +{% endmacro %} diff --git a/dbt/include/databricks/macros/materializations/incremental/strategies.sql b/dbt/include/databricks/macros/materializations/incremental/strategies.sql index 57db5496a..be272d637 100644 --- a/dbt/include/databricks/macros/materializations/incremental/strategies.sql +++ b/dbt/include/databricks/macros/materializations/incremental/strategies.sql @@ -91,7 +91,7 @@ select {{source_cols_csv}} from {{ source_relation }} {%- set not_matched_by_source_action = config.get('not_matched_by_source_action') -%} {%- set not_matched_by_source_condition = config.get('not_matched_by_source_condition') -%} - + {% if unique_key %} {% if unique_key is sequence and unique_key is not mapping and unique_key is not string %} @@ -184,4 +184,4 @@ select {{source_cols_csv}} from {{ source_relation }} {%- endif -%} {%- do arg_dict.update({'incremental_predicates': incremental_predicates}) -%} {{ return(get_replace_where_sql(arg_dict)) }} -{% endmacro %} \ No newline at end of file +{% endmacro %} diff --git a/dbt/include/databricks/macros/materializations/seeds/helpers.sql b/dbt/include/databricks/macros/materializations/seeds/helpers.sql index df690f18b..a2d2af6d7 100644 --- a/dbt/include/databricks/macros/materializations/seeds/helpers.sql +++ b/dbt/include/databricks/macros/materializations/seeds/helpers.sql @@ -77,7 +77,7 @@ {%- set type = column_override.get(col_name, inferred_type) -%} {%- set column_name = (col_name | string) -%} {%- set column_comment_clause = "" -%} - {%- if column_comment and col_name in model.columns.keys() -%} + {%- if column_comment and col_name in model.columns.keys() -%} {%- set comment = model.columns[col_name]['description'] | replace("'", "\\'") -%} {%- if comment and comment != "" -%} {%- set column_comment_clause = "comment '" ~ comment ~ "'" -%} diff --git a/dbt/include/databricks/macros/relations/constraints.sql b/dbt/include/databricks/macros/relations/constraints.sql index bb77145ff..d0a46b151 100644 --- a/dbt/include/databricks/macros/relations/constraints.sql +++ b/dbt/include/databricks/macros/relations/constraints.sql @@ -179,7 +179,7 @@ {% endif %} {% set name = constraint.get("name") %} - + {% if constraint.get('expression') %} {% if not name %} @@ -188,7 +188,7 @@ {%- set name = local_md5("foreign_key;" ~ relation.identifier ~ ";" ~ constraint.get('expression') ~ ";") -%} {% else %} {{ exceptions.raise_compiler_error("Constraint of type " ~ type ~ " with no `name` provided, and no md5 utility.") }} - {% endif %} + {% endif %} {% endif %} {% set stmt = "alter table " ~ relation ~ " add constraint " ~ name ~ " foreign key" ~ constraint.get('expression') %} @@ -224,7 +224,7 @@ {%- set name = local_md5("foreign_key;" ~ relation.identifier ~ ";" ~ column_names ~ ";" ~ parent ~ ";") -%} {% else %} {{ exceptions.raise_compiler_error("Constraint of type " ~ type ~ " with no `name` provided, and no md5 utility.") }} - {% endif %} + {% endif %} {% endif %} {% set stmt = "alter table " ~ relation ~ " add constraint " ~ name ~ " foreign key(" ~ joined_names ~ ") references " ~ parent %} diff --git a/dbt/include/databricks/macros/relations/liquid_clustering.sql b/dbt/include/databricks/macros/relations/liquid_clustering.sql index 3cf810488..a034e6d12 100644 --- a/dbt/include/databricks/macros/relations/liquid_clustering.sql +++ b/dbt/include/databricks/macros/relations/liquid_clustering.sql @@ -18,4 +18,4 @@ ALTER {{ target_relation.type }} {{ target_relation }} CLUSTER BY ({{ cols | join(', ') }}) {%- endcall -%} {%- endif %} -{%- endmacro -%} \ No newline at end of file +{%- endmacro -%} diff --git a/dbt/include/databricks/macros/relations/tags.sql b/dbt/include/databricks/macros/relations/tags.sql index 3467631df..d26da9a20 100644 --- a/dbt/include/databricks/macros/relations/tags.sql +++ b/dbt/include/databricks/macros/relations/tags.sql @@ -11,7 +11,7 @@ {% macro fetch_tags_sql(relation) -%} SELECT tag_name, tag_value FROM `system`.`information_schema`.`table_tags` - WHERE catalog_name = '{{ relation.database|lower }}' + WHERE catalog_name = '{{ relation.database|lower }}' AND schema_name = '{{ relation.schema|lower }}' AND table_name = '{{ relation.identifier|lower }}' {%- endmacro -%} diff --git a/dbt/include/databricks/macros/utils/split_part.sql b/dbt/include/databricks/macros/utils/split_part.sql index 92b0cecd0..3f2773500 100644 --- a/dbt/include/databricks/macros/utils/split_part.sql +++ b/dbt/include/databricks/macros/utils/split_part.sql @@ -27,7 +27,7 @@ get(split( {{ string_text }}, {{ delimiter_expr }} - ), + ), length({{ string_text }}) - length( replace({{ string_text }}, {{ delimiter_text }}, '') diff --git a/pyproject.toml b/pyproject.toml index f1f680ead..287ab5327 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -91,15 +91,35 @@ uc-cluster-e2e = "pytest --color=yes -v --profile databricks_uc_cluster -n auto sqlw-e2e = "pytest --color=yes -v --profile databricks_uc_sql_endpoint -n auto --dist=loadscope tests/functional" [tool.hatch.envs.test.scripts] -unit = "pytest --color=yes -v --profile databricks_cluster -n auto --dist=loadscope tests/unit" +unit = """ +if [[ "$DBT_DATABRICKS_SESSION_CONNECTION" == "True" ]]; then + PROFILE="session_connection" +else + PROFILE="databricks_cluster" +fi + +pytest --color=yes -v --profile $PROFILE -n auto --dist=loadscope tests/unit +""" [[tool.hatch.envs.test.matrix]] python = ["3.9", "3.10", "3.11", "3.12"] +session_support = ["no_session", "session"] + + +[tool.hatch.envs.test.overrides] +matrix.session_support.env-vars = [ + { key = "DBT_DATABRICKS_SESSION_CONNECTION", value = "False", if = ["no_session"]}, + { key = "DBT_DATABRICKS_SESSION_CONNECTION", value = "True", if = ["session"] } +] + + + [tool.ruff] line-length = 100 target-version = 'py39' + [tool.ruff.lint] select = ["E", "W", "F", "I"] ignore = ["E203"] @@ -121,6 +141,8 @@ markers = [ strict_optional = true no_implicit_optional = true disallow_untyped_defs = true +allow_redefinition = true + [[tool.mypy.overrides]] module = "tests.*" diff --git a/tests/functional/adapter/aliases/test_aliases.py b/tests/functional/adapter/aliases/test_aliases.py index f0f90823d..a0f4495e9 100644 --- a/tests/functional/adapter/aliases/test_aliases.py +++ b/tests/functional/adapter/aliases/test_aliases.py @@ -1,12 +1,10 @@ import pytest from dbt.tests.adapter.aliases import fixtures as dbt_fixtures -from dbt.tests.adapter.aliases.test_aliases import ( - BaseAliasErrors, - BaseAliases, - BaseSameAliasDifferentDatabases, - BaseSameAliasDifferentSchemas, -) +from dbt.tests.adapter.aliases.test_aliases import BaseAliasErrors +from dbt.tests.adapter.aliases.test_aliases import BaseAliases +from dbt.tests.adapter.aliases.test_aliases import BaseSameAliasDifferentDatabases +from dbt.tests.adapter.aliases.test_aliases import BaseSameAliasDifferentSchemas from tests.functional.adapter.aliases import fixtures as databricks_fixtures macro_override = { diff --git a/tests/functional/adapter/basic/test_docs_generate.py b/tests/functional/adapter/basic/test_docs_generate.py index 83593eb5b..677879df4 100644 --- a/tests/functional/adapter/basic/test_docs_generate.py +++ b/tests/functional/adapter/basic/test_docs_generate.py @@ -1,9 +1,11 @@ import pytest from dbt.tests.adapter.basic import expected_catalog -from dbt.tests.adapter.basic.test_docs_generate import BaseDocsGenerate, BaseDocsGenReferences +from dbt.tests.adapter.basic.test_docs_generate import BaseDocsGenerate +from dbt.tests.adapter.basic.test_docs_generate import BaseDocsGenReferences from dbt.tests.util import AnyString -from tests.functional.adapter.basic.typing import AnyLongType, StatsLikeDict +from tests.functional.adapter.basic.typing import AnyLongType +from tests.functional.adapter.basic.typing import StatsLikeDict class TestDocsGenerate(BaseDocsGenerate): diff --git a/tests/functional/adapter/basic/test_incremental.py b/tests/functional/adapter/basic/test_incremental.py index a353d7802..3958b3304 100644 --- a/tests/functional/adapter/basic/test_incremental.py +++ b/tests/functional/adapter/basic/test_incremental.py @@ -2,7 +2,8 @@ import pytest -from dbt.tests.adapter.basic.test_incremental import BaseIncremental, BaseIncrementalNotSchemaChange +from dbt.tests.adapter.basic.test_incremental import BaseIncremental +from dbt.tests.adapter.basic.test_incremental import BaseIncrementalNotSchemaChange class TestIncrementalDelta(BaseIncremental): diff --git a/tests/functional/adapter/caching/test_caching.py b/tests/functional/adapter/caching/test_caching.py index d32b3d534..206f092f1 100644 --- a/tests/functional/adapter/caching/test_caching.py +++ b/tests/functional/adapter/caching/test_caching.py @@ -1,9 +1,7 @@ -from dbt.tests.adapter.caching.test_caching import ( - BaseCachingLowercaseModel, - BaseCachingSelectedSchemaOnly, - BaseCachingUppercaseModel, - BaseNoPopulateCache, -) +from dbt.tests.adapter.caching.test_caching import BaseCachingLowercaseModel +from dbt.tests.adapter.caching.test_caching import BaseCachingSelectedSchemaOnly +from dbt.tests.adapter.caching.test_caching import BaseCachingUppercaseModel +from dbt.tests.adapter.caching.test_caching import BaseNoPopulateCache class TestNoPopulateCache(BaseNoPopulateCache): diff --git a/tests/functional/adapter/concurrency/test_concurrency.py b/tests/functional/adapter/concurrency/test_concurrency.py index b1b8aded5..a4e410940 100644 --- a/tests/functional/adapter/concurrency/test_concurrency.py +++ b/tests/functional/adapter/concurrency/test_concurrency.py @@ -1,5 +1,6 @@ from dbt.tests import util -from dbt.tests.adapter.concurrency.test_concurrency import BaseConcurrency, seeds__update_csv +from dbt.tests.adapter.concurrency.test_concurrency import BaseConcurrency +from dbt.tests.adapter.concurrency.test_concurrency import seeds__update_csv # Copied from dbt-core diff --git a/tests/functional/adapter/constraints/test_constraints.py b/tests/functional/adapter/constraints/test_constraints.py index 350910d00..cfd4d1156 100644 --- a/tests/functional/adapter/constraints/test_constraints.py +++ b/tests/functional/adapter/constraints/test_constraints.py @@ -2,16 +2,14 @@ from dbt.tests import util from dbt.tests.adapter.constraints import fixtures -from dbt.tests.adapter.constraints.test_constraints import ( - BaseConstraintQuotedColumn, - BaseConstraintsRollback, - BaseConstraintsRuntimeDdlEnforcement, - BaseIncrementalConstraintsColumnsEqual, - BaseIncrementalConstraintsRollback, - BaseIncrementalConstraintsRuntimeDdlEnforcement, - BaseTableConstraintsColumnsEqual, - BaseViewConstraintsColumnsEqual, -) +from dbt.tests.adapter.constraints.test_constraints import BaseConstraintQuotedColumn +from dbt.tests.adapter.constraints.test_constraints import BaseConstraintsRollback +from dbt.tests.adapter.constraints.test_constraints import BaseConstraintsRuntimeDdlEnforcement +from dbt.tests.adapter.constraints.test_constraints import BaseIncrementalConstraintsColumnsEqual +from dbt.tests.adapter.constraints.test_constraints import BaseIncrementalConstraintsRollback +from dbt.tests.adapter.constraints.test_constraints import BaseIncrementalConstraintsRuntimeDdlEnforcement +from dbt.tests.adapter.constraints.test_constraints import BaseTableConstraintsColumnsEqual +from dbt.tests.adapter.constraints.test_constraints import BaseViewConstraintsColumnsEqual from tests.functional.adapter.constraints import fixtures as override_fixtures diff --git a/tests/functional/adapter/dbt_clone/test_dbt_clone.py b/tests/functional/adapter/dbt_clone/test_dbt_clone.py index 6d619abff..291554658 100644 --- a/tests/functional/adapter/dbt_clone/test_dbt_clone.py +++ b/tests/functional/adapter/dbt_clone/test_dbt_clone.py @@ -1,7 +1,8 @@ import pytest from dbt.tests import util -from dbt.tests.adapter.dbt_clone.test_dbt_clone import BaseClone, BaseClonePossible +from dbt.tests.adapter.dbt_clone.test_dbt_clone import BaseClone +from dbt.tests.adapter.dbt_clone.test_dbt_clone import BaseClonePossible from dbt.tests.util import run_dbt from tests.functional.adapter.dbt_clone import fixtures diff --git a/tests/functional/adapter/dbt_show/test_dbt_show.py b/tests/functional/adapter/dbt_show/test_dbt_show.py index 98d603154..340828c02 100644 --- a/tests/functional/adapter/dbt_show/test_dbt_show.py +++ b/tests/functional/adapter/dbt_show/test_dbt_show.py @@ -1,4 +1,5 @@ -from dbt.tests.adapter.dbt_show.test_dbt_show import BaseShowLimit, BaseShowSqlHeader +from dbt.tests.adapter.dbt_show.test_dbt_show import BaseShowLimit +from dbt.tests.adapter.dbt_show.test_dbt_show import BaseShowSqlHeader class TestShowLimit(BaseShowLimit): diff --git a/tests/functional/adapter/ephemeral/test_ephemeral.py b/tests/functional/adapter/ephemeral/test_ephemeral.py index c00585b81..81a320dbf 100644 --- a/tests/functional/adapter/ephemeral/test_ephemeral.py +++ b/tests/functional/adapter/ephemeral/test_ephemeral.py @@ -5,7 +5,8 @@ from dbt.tests import util from dbt.tests.adapter.ephemeral import test_ephemeral -from dbt.tests.adapter.ephemeral.test_ephemeral import BaseEphemeral, BaseEphemeralMulti +from dbt.tests.adapter.ephemeral.test_ephemeral import BaseEphemeral +from dbt.tests.adapter.ephemeral.test_ephemeral import BaseEphemeralMulti class TestEphemeralMulti(BaseEphemeralMulti): diff --git a/tests/functional/adapter/hooks/test_run_hooks.py b/tests/functional/adapter/hooks/test_run_hooks.py index 5c7dd5c2d..91c228437 100644 --- a/tests/functional/adapter/hooks/test_run_hooks.py +++ b/tests/functional/adapter/hooks/test_run_hooks.py @@ -3,7 +3,8 @@ import pytest from dbt.tests import util -from dbt.tests.adapter.hooks.test_run_hooks import BaseAfterRunHooks, BasePrePostRunHooks +from dbt.tests.adapter.hooks.test_run_hooks import BaseAfterRunHooks +from dbt.tests.adapter.hooks.test_run_hooks import BasePrePostRunHooks from tests.functional.adapter.hooks import fixtures as override_fixtures diff --git a/tests/functional/adapter/materialized_view_tests/test_changes.py b/tests/functional/adapter/materialized_view_tests/test_changes.py index 1c1fa4c71..49f715a1c 100644 --- a/tests/functional/adapter/materialized_view_tests/test_changes.py +++ b/tests/functional/adapter/materialized_view_tests/test_changes.py @@ -8,12 +8,10 @@ ) from dbt.adapters.databricks.relation_configs.tblproperties import TblPropertiesConfig from dbt.tests import util -from dbt.tests.adapter.materialized_view.changes import ( - MaterializedViewChanges, - MaterializedViewChangesApplyMixin, - MaterializedViewChangesContinueMixin, - MaterializedViewChangesFailMixin, -) +from dbt.tests.adapter.materialized_view.changes import MaterializedViewChanges +from dbt.tests.adapter.materialized_view.changes import MaterializedViewChangesApplyMixin +from dbt.tests.adapter.materialized_view.changes import MaterializedViewChangesContinueMixin +from dbt.tests.adapter.materialized_view.changes import MaterializedViewChangesFailMixin from tests.functional.adapter.materialized_view_tests import fixtures diff --git a/tests/functional/adapter/persist_constraints/test_persist_constraints.py b/tests/functional/adapter/persist_constraints/test_persist_constraints.py index ac7c50b0a..42c9862fa 100644 --- a/tests/functional/adapter/persist_constraints/test_persist_constraints.py +++ b/tests/functional/adapter/persist_constraints/test_persist_constraints.py @@ -1,6 +1,7 @@ import pytest -from dbt.contracts.results import RunResult, RunStatus +from dbt.contracts.results import RunResult +from dbt.contracts.results import RunStatus from dbt.tests import util from tests.functional.adapter.persist_constraints import fixtures diff --git a/tests/functional/adapter/python_model/test_python_model.py b/tests/functional/adapter/python_model/test_python_model.py index 858214b79..453807041 100644 --- a/tests/functional/adapter/python_model/test_python_model.py +++ b/tests/functional/adapter/python_model/test_python_model.py @@ -4,10 +4,8 @@ from dbt.tests import util from dbt.tests.adapter.python_model import test_python_model as fixtures -from dbt.tests.adapter.python_model.test_python_model import ( - BasePythonIncrementalTests, - BasePythonModelTests, -) +from dbt.tests.adapter.python_model.test_python_model import BasePythonIncrementalTests +from dbt.tests.adapter.python_model.test_python_model import BasePythonModelTests from tests.functional.adapter.python_model import fixtures as override_fixtures diff --git a/tests/functional/adapter/query_comments/test_query_comment.py b/tests/functional/adapter/query_comments/test_query_comment.py index 8d64835ba..3f05b41e2 100644 --- a/tests/functional/adapter/query_comments/test_query_comment.py +++ b/tests/functional/adapter/query_comments/test_query_comment.py @@ -1,13 +1,11 @@ import pytest -from dbt.tests.adapter.query_comment.test_query_comment import ( - BaseEmptyQueryComments, - BaseMacroArgsQueryComments, - BaseMacroInvalidQueryComments, - BaseMacroQueryComments, - BaseNullQueryComments, - BaseQueryComments, -) +from dbt.tests.adapter.query_comment.test_query_comment import BaseEmptyQueryComments +from dbt.tests.adapter.query_comment.test_query_comment import BaseMacroArgsQueryComments +from dbt.tests.adapter.query_comment.test_query_comment import BaseMacroInvalidQueryComments +from dbt.tests.adapter.query_comment.test_query_comment import BaseMacroQueryComments +from dbt.tests.adapter.query_comment.test_query_comment import BaseNullQueryComments +from dbt.tests.adapter.query_comment.test_query_comment import BaseQueryComments class TestQueryComments(BaseQueryComments): diff --git a/tests/functional/adapter/simple_copy/test_simple_copy.py b/tests/functional/adapter/simple_copy/test_simple_copy.py index ba65ad15e..feac39bd5 100644 --- a/tests/functional/adapter/simple_copy/test_simple_copy.py +++ b/tests/functional/adapter/simple_copy/test_simple_copy.py @@ -1,6 +1,7 @@ import pytest -from dbt.tests.adapter.simple_copy.test_simple_copy import EmptyModelsArentRunBase, SimpleCopyBase +from dbt.tests.adapter.simple_copy.test_simple_copy import EmptyModelsArentRunBase +from dbt.tests.adapter.simple_copy.test_simple_copy import SimpleCopyBase # Tests with materialized_views, which only works for SQL Warehouse diff --git a/tests/functional/adapter/simple_seed/test_seeds.py b/tests/functional/adapter/simple_seed/test_seeds.py index 242f79074..7e18b61d2 100644 --- a/tests/functional/adapter/simple_seed/test_seeds.py +++ b/tests/functional/adapter/simple_seed/test_seeds.py @@ -4,19 +4,17 @@ from dbt.tests import util from dbt.tests.adapter.simple_seed import seeds -from dbt.tests.adapter.simple_seed.test_seed import ( - BaseSeedConfigFullRefreshOff, - BaseSeedCustomSchema, - BaseSeedParsing, - BaseSeedSpecificFormats, - BaseSeedWithEmptyDelimiter, - BaseSeedWithUniqueDelimiter, - BaseSeedWithWrongDelimiter, - BaseSimpleSeedEnabledViaConfig, - BaseSimpleSeedWithBOM, - BaseTestEmptySeed, - SeedTestBase, -) +from dbt.tests.adapter.simple_seed.test_seed import BaseSeedConfigFullRefreshOff +from dbt.tests.adapter.simple_seed.test_seed import BaseSeedCustomSchema +from dbt.tests.adapter.simple_seed.test_seed import BaseSeedParsing +from dbt.tests.adapter.simple_seed.test_seed import BaseSeedSpecificFormats +from dbt.tests.adapter.simple_seed.test_seed import BaseSeedWithEmptyDelimiter +from dbt.tests.adapter.simple_seed.test_seed import BaseSeedWithUniqueDelimiter +from dbt.tests.adapter.simple_seed.test_seed import BaseSeedWithWrongDelimiter +from dbt.tests.adapter.simple_seed.test_seed import BaseSimpleSeedEnabledViaConfig +from dbt.tests.adapter.simple_seed.test_seed import BaseSimpleSeedWithBOM +from dbt.tests.adapter.simple_seed.test_seed import BaseTestEmptySeed +from dbt.tests.adapter.simple_seed.test_seed import SeedTestBase from tests.functional.adapter.simple_seed import fixtures diff --git a/tests/functional/adapter/simple_snapshot/test_snapshot.py b/tests/functional/adapter/simple_snapshot/test_snapshot.py index 56186e4ef..ef2b873e9 100644 --- a/tests/functional/adapter/simple_snapshot/test_snapshot.py +++ b/tests/functional/adapter/simple_snapshot/test_snapshot.py @@ -3,7 +3,8 @@ import pytest from dbt.tests import util -from dbt.tests.adapter.simple_snapshot.test_snapshot import BaseSimpleSnapshot, BaseSnapshotCheck +from dbt.tests.adapter.simple_snapshot.test_snapshot import BaseSimpleSnapshot +from dbt.tests.adapter.simple_snapshot.test_snapshot import BaseSnapshotCheck from dbt.tests.util import run_dbt from tests.functional.adapter.simple_snapshot import fixtures diff --git a/tests/functional/adapter/store_test_failures/test_store_failures_as.py b/tests/functional/adapter/store_test_failures/test_store_failures_as.py index 1d830dcb7..ccebc8115 100644 --- a/tests/functional/adapter/store_test_failures/test_store_failures_as.py +++ b/tests/functional/adapter/store_test_failures/test_store_failures_as.py @@ -1,11 +1,9 @@ -from dbt.tests.adapter.store_test_failures_tests.basic import ( - StoreTestFailuresAsExceptions, - StoreTestFailuresAsGeneric, - StoreTestFailuresAsInteractions, - StoreTestFailuresAsProjectLevelEphemeral, - StoreTestFailuresAsProjectLevelOff, - StoreTestFailuresAsProjectLevelView, -) +from dbt.tests.adapter.store_test_failures_tests.basic import StoreTestFailuresAsExceptions +from dbt.tests.adapter.store_test_failures_tests.basic import StoreTestFailuresAsGeneric +from dbt.tests.adapter.store_test_failures_tests.basic import StoreTestFailuresAsInteractions +from dbt.tests.adapter.store_test_failures_tests.basic import StoreTestFailuresAsProjectLevelEphemeral +from dbt.tests.adapter.store_test_failures_tests.basic import StoreTestFailuresAsProjectLevelOff +from dbt.tests.adapter.store_test_failures_tests.basic import StoreTestFailuresAsProjectLevelView from dbt.tests.adapter.store_test_failures_tests.test_store_test_failures import ( BaseStoreTestFailures, ) diff --git a/tests/functional/adapter/streaming_tables/test_st_basic.py b/tests/functional/adapter/streaming_tables/test_st_basic.py index 59f5ea5e7..ba95143b2 100644 --- a/tests/functional/adapter/streaming_tables/test_st_basic.py +++ b/tests/functional/adapter/streaming_tables/test_st_basic.py @@ -5,7 +5,9 @@ from dbt.adapters.base.relation import BaseRelation from dbt.adapters.databricks.relation import DatabricksRelationType from dbt.tests import util -from dbt.tests.adapter.materialized_view.files import MY_SEED, MY_TABLE, MY_VIEW +from dbt.tests.adapter.materialized_view.files import MY_SEED +from dbt.tests.adapter.materialized_view.files import MY_TABLE +from dbt.tests.adapter.materialized_view.files import MY_VIEW from tests.functional.adapter.streaming_tables import fixtures diff --git a/tests/functional/adapter/utils/test_source_freshness.py b/tests/functional/adapter/utils/test_source_freshness.py index 857690b75..5b2024c73 100644 --- a/tests/functional/adapter/utils/test_source_freshness.py +++ b/tests/functional/adapter/utils/test_source_freshness.py @@ -2,7 +2,8 @@ import pytest -from dbt.tests.util import get_artifact, run_dbt +from dbt.tests.util import get_artifact +from dbt.tests.util import run_dbt freshness_via_metadata_schema_yml = """ version: 2 diff --git a/tests/functional/adapter/utils/test_utils.py b/tests/functional/adapter/utils/test_utils.py index 06086f25d..a6b99634b 100644 --- a/tests/functional/adapter/utils/test_utils.py +++ b/tests/functional/adapter/utils/test_utils.py @@ -20,7 +20,8 @@ from dbt.tests.adapter.utils.test_intersect import BaseIntersect from dbt.tests.adapter.utils.test_last_day import BaseLastDay from dbt.tests.adapter.utils.test_length import BaseLength -from dbt.tests.adapter.utils.test_null_compare import BaseMixedNullCompare, BaseNullCompare +from dbt.tests.adapter.utils.test_null_compare import BaseMixedNullCompare +from dbt.tests.adapter.utils.test_null_compare import BaseNullCompare from dbt.tests.adapter.utils.test_position import BasePosition from dbt.tests.adapter.utils.test_replace import BaseReplace from dbt.tests.adapter.utils.test_right import BaseRight diff --git a/tests/profiles.py b/tests/profiles.py index e0c88447e..e34c5073f 100644 --- a/tests/profiles.py +++ b/tests/profiles.py @@ -1,5 +1,6 @@ import os -from typing import Any, Optional +from typing import Any +from typing import Optional def get_databricks_cluster_target(profile_type: str): diff --git a/tests/unit/api_client/api_test_base.py b/tests/unit/api_client/api_test_base.py index 3d60dc047..8377bd19d 100644 --- a/tests/unit/api_client/api_test_base.py +++ b/tests/unit/api_client/api_test_base.py @@ -1,4 +1,5 @@ -from typing import Any, Callable +from typing import Any +from typing import Callable from unittest.mock import Mock import pytest diff --git a/tests/unit/api_client/test_command_api.py b/tests/unit/api_client/test_command_api.py index 90efaf472..57a94ae16 100644 --- a/tests/unit/api_client/test_command_api.py +++ b/tests/unit/api_client/test_command_api.py @@ -1,10 +1,12 @@ -from unittest.mock import Mock, patch +from unittest.mock import Mock +from unittest.mock import patch import freezegun import pytest from dbt_common.exceptions import DbtRuntimeError -from dbt.adapters.databricks.api_client import CommandApi, CommandExecution +from dbt.adapters.databricks.api_client import CommandApi +from dbt.adapters.databricks.api_client import CommandExecution from tests.unit.api_client.api_test_base import ApiTestBase diff --git a/tests/unit/api_client/test_user_folder_api.py b/tests/unit/api_client/test_user_folder_api.py index c0ab4b081..ca540dce4 100644 --- a/tests/unit/api_client/test_user_folder_api.py +++ b/tests/unit/api_client/test_user_folder_api.py @@ -1,6 +1,7 @@ import pytest -from dbt.adapters.databricks.api_client import CurrUserApi, UserFolderApi +from dbt.adapters.databricks.api_client import CurrUserApi +from dbt.adapters.databricks.api_client import UserFolderApi from tests.unit.api_client.api_test_base import ApiTestBase diff --git a/tests/unit/events/test_base_events.py b/tests/unit/events/test_base_events.py index 7b480029d..990c51a42 100644 --- a/tests/unit/events/test_base_events.py +++ b/tests/unit/events/test_base_events.py @@ -1,6 +1,7 @@ from databricks.sql.exc import Error -from dbt.adapters.databricks.events.base import ErrorEvent, SQLErrorEvent +from dbt.adapters.databricks.events.base import ErrorEvent +from dbt.adapters.databricks.events.base import SQLErrorEvent class ErrorTestEvent(ErrorEvent): diff --git a/tests/unit/events/test_connection_events.py b/tests/unit/events/test_connection_events.py index 81a30c1dc..995110ca3 100644 --- a/tests/unit/events/test_connection_events.py +++ b/tests/unit/events/test_connection_events.py @@ -1,10 +1,8 @@ from unittest.mock import Mock -from dbt.adapters.databricks.events.connection_events import ( - ConnectionAcquire, - ConnectionCloseError, - ConnectionEvent, -) +from dbt.adapters.databricks.events.connection_events import ConnectionAcquire +from dbt.adapters.databricks.events.connection_events import ConnectionCloseError +from dbt.adapters.databricks.events.connection_events import ConnectionEvent class ConnectionTestEvent(ConnectionEvent): diff --git a/tests/unit/events/test_cursor_events.py b/tests/unit/events/test_cursor_events.py index e492415a6..e30caa849 100644 --- a/tests/unit/events/test_cursor_events.py +++ b/tests/unit/events/test_cursor_events.py @@ -1,6 +1,7 @@ from unittest.mock import Mock -from dbt.adapters.databricks.events.cursor_events import CursorCloseError, CursorEvent +from dbt.adapters.databricks.events.cursor_events import CursorCloseError +from dbt.adapters.databricks.events.cursor_events import CursorEvent class CursorTestEvent(CursorEvent): diff --git a/tests/unit/macros/base.py b/tests/unit/macros/base.py index cb9c7f282..02c8cf94a 100644 --- a/tests/unit/macros/base.py +++ b/tests/unit/macros/base.py @@ -3,7 +3,10 @@ from unittest.mock import Mock import pytest -from jinja2 import Environment, FileSystemLoader, PackageLoader, Template +from jinja2 import Environment +from jinja2 import FileSystemLoader +from jinja2 import PackageLoader +from jinja2 import Template from dbt.adapters.databricks.relation import DatabricksRelation diff --git a/tests/unit/python/test_python_config.py b/tests/unit/python/test_python_config.py index 3559e044b..6e00cd6f5 100644 --- a/tests/unit/python/test_python_config.py +++ b/tests/unit/python/test_python_config.py @@ -1,11 +1,9 @@ import pytest from pydantic import ValidationError -from dbt.adapters.databricks.python_models.python_config import ( - ParsedPythonModel, - PythonJobConfig, - PythonModelConfig, -) +from dbt.adapters.databricks.python_models.python_config import ParsedPythonModel +from dbt.adapters.databricks.python_models.python_config import PythonJobConfig +from dbt.adapters.databricks.python_models.python_config import PythonModelConfig class TestParsedPythonModel: diff --git a/tests/unit/python/test_python_helpers.py b/tests/unit/python/test_python_helpers.py index 8c93bf20e..22f6f5f0c 100644 --- a/tests/unit/python/test_python_helpers.py +++ b/tests/unit/python/test_python_helpers.py @@ -2,15 +2,13 @@ import pytest -from dbt.adapters.databricks.python_models.python_submissions import ( - AllPurposeClusterPythonJobHelper, - JobClusterPythonJobHelper, - PythonCommandSubmitter, - PythonNotebookSubmitter, - PythonNotebookWorkflowSubmitter, - ServerlessClusterPythonJobHelper, - WorkflowPythonJobHelper, -) +from dbt.adapters.databricks.python_models.python_submissions import AllPurposeClusterPythonJobHelper +from dbt.adapters.databricks.python_models.python_submissions import JobClusterPythonJobHelper +from dbt.adapters.databricks.python_models.python_submissions import PythonCommandSubmitter +from dbt.adapters.databricks.python_models.python_submissions import PythonNotebookSubmitter +from dbt.adapters.databricks.python_models.python_submissions import PythonNotebookWorkflowSubmitter +from dbt.adapters.databricks.python_models.python_submissions import ServerlessClusterPythonJobHelper +from dbt.adapters.databricks.python_models.python_submissions import WorkflowPythonJobHelper @pytest.fixture diff --git a/tests/unit/python/test_python_job_support.py b/tests/unit/python/test_python_job_support.py index 41f480413..93dfe3b46 100644 --- a/tests/unit/python/test_python_job_support.py +++ b/tests/unit/python/test_python_job_support.py @@ -3,11 +3,9 @@ import pytest from dbt.adapters.databricks.python_models import python_submissions -from dbt.adapters.databricks.python_models.python_submissions import ( - PythonJobConfigCompiler, - PythonNotebookUploader, - PythonPermissionBuilder, -) +from dbt.adapters.databricks.python_models.python_submissions import PythonJobConfigCompiler +from dbt.adapters.databricks.python_models.python_submissions import PythonNotebookUploader +from dbt.adapters.databricks.python_models.python_submissions import PythonPermissionBuilder @pytest.fixture diff --git a/tests/unit/python/test_python_submitters.py b/tests/unit/python/test_python_submitters.py index cf7077302..1abd50ad9 100644 --- a/tests/unit/python/test_python_submitters.py +++ b/tests/unit/python/test_python_submitters.py @@ -2,17 +2,15 @@ import pytest -from dbt.adapters.databricks.python_models.python_submissions import ( - PythonCommandSubmitter, - PythonJobConfigCompiler, - PythonJobDetails, - PythonNotebookSubmitter, - PythonNotebookUploader, - PythonNotebookWorkflowSubmitter, - PythonPermissionBuilder, - PythonWorkflowConfigCompiler, - PythonWorkflowCreator, -) +from dbt.adapters.databricks.python_models.python_submissions import PythonCommandSubmitter +from dbt.adapters.databricks.python_models.python_submissions import PythonJobConfigCompiler +from dbt.adapters.databricks.python_models.python_submissions import PythonJobDetails +from dbt.adapters.databricks.python_models.python_submissions import PythonNotebookSubmitter +from dbt.adapters.databricks.python_models.python_submissions import PythonNotebookUploader +from dbt.adapters.databricks.python_models.python_submissions import PythonNotebookWorkflowSubmitter +from dbt.adapters.databricks.python_models.python_submissions import PythonPermissionBuilder +from dbt.adapters.databricks.python_models.python_submissions import PythonWorkflowConfigCompiler +from dbt.adapters.databricks.python_models.python_submissions import PythonWorkflowCreator @pytest.fixture diff --git a/tests/unit/python/test_python_workflow_support.py b/tests/unit/python/test_python_workflow_support.py index 595bd90c9..9038b54ff 100644 --- a/tests/unit/python/test_python_workflow_support.py +++ b/tests/unit/python/test_python_workflow_support.py @@ -2,10 +2,8 @@ import pytest -from dbt.adapters.databricks.python_models.python_submissions import ( - PythonWorkflowConfigCompiler, - PythonWorkflowCreator, -) +from dbt.adapters.databricks.python_models.python_submissions import PythonWorkflowConfigCompiler +from dbt.adapters.databricks.python_models.python_submissions import PythonWorkflowCreator class TestPythonWorkflowConfigCompiler: diff --git a/tests/unit/relation_configs/test_comment.py b/tests/unit/relation_configs/test_comment.py index 71d802073..8e2fc1504 100644 --- a/tests/unit/relation_configs/test_comment.py +++ b/tests/unit/relation_configs/test_comment.py @@ -2,7 +2,8 @@ from agate import Table -from dbt.adapters.databricks.relation_configs.comment import CommentConfig, CommentProcessor +from dbt.adapters.databricks.relation_configs.comment import CommentConfig +from dbt.adapters.databricks.relation_configs.comment import CommentProcessor class TestCommentProcessor: diff --git a/tests/unit/relation_configs/test_materialized_view_config.py b/tests/unit/relation_configs/test_materialized_view_config.py index 09aa95717..80b6d2040 100644 --- a/tests/unit/relation_configs/test_materialized_view_config.py +++ b/tests/unit/relation_configs/test_materialized_view_config.py @@ -1,6 +1,7 @@ from unittest.mock import Mock -from agate import Row, Table +from agate import Row +from agate import Table from dbt.adapters.databricks.relation_configs.comment import CommentConfig from dbt.adapters.databricks.relation_configs.materialized_view import ( diff --git a/tests/unit/relation_configs/test_partitioning.py b/tests/unit/relation_configs/test_partitioning.py index fb345f066..731b09c5d 100644 --- a/tests/unit/relation_configs/test_partitioning.py +++ b/tests/unit/relation_configs/test_partitioning.py @@ -1,9 +1,7 @@ from unittest.mock import Mock -from dbt.adapters.databricks.relation_configs.partitioning import ( - PartitionedByConfig, - PartitionedByProcessor, -) +from dbt.adapters.databricks.relation_configs.partitioning import PartitionedByConfig +from dbt.adapters.databricks.relation_configs.partitioning import PartitionedByProcessor from tests.unit import fixtures diff --git a/tests/unit/relation_configs/test_query.py b/tests/unit/relation_configs/test_query.py index 3cd1f1b34..30748802e 100644 --- a/tests/unit/relation_configs/test_query.py +++ b/tests/unit/relation_configs/test_query.py @@ -3,7 +3,8 @@ import pytest from agate import Row -from dbt.adapters.databricks.relation_configs.query import QueryConfig, QueryProcessor +from dbt.adapters.databricks.relation_configs.query import QueryConfig +from dbt.adapters.databricks.relation_configs.query import QueryProcessor from dbt.exceptions import DbtRuntimeError sql = "select * from foo" diff --git a/tests/unit/relation_configs/test_refresh.py b/tests/unit/relation_configs/test_refresh.py index 9da56e4d8..33efc22d2 100644 --- a/tests/unit/relation_configs/test_refresh.py +++ b/tests/unit/relation_configs/test_refresh.py @@ -2,7 +2,8 @@ import pytest -from dbt.adapters.databricks.relation_configs.refresh import RefreshConfig, RefreshProcessor +from dbt.adapters.databricks.relation_configs.refresh import RefreshConfig +from dbt.adapters.databricks.relation_configs.refresh import RefreshProcessor from dbt.exceptions import DbtRuntimeError from tests.unit import fixtures diff --git a/tests/unit/relation_configs/test_tags.py b/tests/unit/relation_configs/test_tags.py index e893d3f59..d45a813e5 100644 --- a/tests/unit/relation_configs/test_tags.py +++ b/tests/unit/relation_configs/test_tags.py @@ -3,7 +3,8 @@ import pytest from agate import Table -from dbt.adapters.databricks.relation_configs.tags import TagsConfig, TagsProcessor +from dbt.adapters.databricks.relation_configs.tags import TagsConfig +from dbt.adapters.databricks.relation_configs.tags import TagsProcessor from dbt.exceptions import DbtRuntimeError diff --git a/tests/unit/relation_configs/test_tblproperties.py b/tests/unit/relation_configs/test_tblproperties.py index dab860b7d..23f563735 100644 --- a/tests/unit/relation_configs/test_tblproperties.py +++ b/tests/unit/relation_configs/test_tblproperties.py @@ -2,10 +2,8 @@ import pytest -from dbt.adapters.databricks.relation_configs.tblproperties import ( - TblPropertiesConfig, - TblPropertiesProcessor, -) +from dbt.adapters.databricks.relation_configs.tblproperties import TblPropertiesConfig +from dbt.adapters.databricks.relation_configs.tblproperties import TblPropertiesProcessor from dbt.exceptions import DbtRuntimeError from tests.unit import fixtures diff --git a/tests/unit/test_adapter.py b/tests/unit/test_adapter.py index e707dcbb3..394256870 100644 --- a/tests/unit/test_adapter.py +++ b/tests/unit/test_adapter.py @@ -1,21 +1,24 @@ from multiprocessing import get_context -from typing import Any, Optional -from unittest.mock import Mock, patch +from typing import Any +from typing import Optional +from unittest.mock import Mock +from unittest.mock import patch import pytest from agate import Row -from dbt_common.exceptions import DbtConfigError, DbtValidationError +from dbt_common.exceptions import DbtConfigError +from dbt_common.exceptions import DbtValidationError import dbt.flags as flags -from dbt.adapters.databricks import DatabricksAdapter, __version__ +from dbt.adapters.databricks import __version__ +from dbt.adapters.databricks import DatabricksAdapter from dbt.adapters.databricks.column import DatabricksColumn -from dbt.adapters.databricks.credentials import ( - CATALOG_KEY_IN_SESSION_PROPERTIES, - DBT_DATABRICKS_HTTP_SESSION_HEADERS, - DBT_DATABRICKS_INVOCATION_ENV, -) +from dbt.adapters.databricks.credentials import CATALOG_KEY_IN_SESSION_PROPERTIES +from dbt.adapters.databricks.credentials import DBT_DATABRICKS_HTTP_SESSION_HEADERS +from dbt.adapters.databricks.credentials import DBT_DATABRICKS_INVOCATION_ENV from dbt.adapters.databricks.impl import get_identifier_list_string -from dbt.adapters.databricks.relation import DatabricksRelation, DatabricksRelationType +from dbt.adapters.databricks.relation import DatabricksRelation +from dbt.adapters.databricks.relation import DatabricksRelationType from dbt.adapters.databricks.utils import check_not_found_error from dbt.config import RuntimeConfig from tests.unit.utils import config_from_parts_or_dicts diff --git a/tests/unit/test_idle_config.py b/tests/unit/test_idle_config.py index de3545680..5e79f7df9 100644 --- a/tests/unit/test_idle_config.py +++ b/tests/unit/test_idle_config.py @@ -3,7 +3,8 @@ from dbt.adapters.databricks import connections from dbt.adapters.databricks.credentials import DatabricksCredentials -from dbt.contracts.graph import model_config, nodes +from dbt.contracts.graph import model_config +from dbt.contracts.graph import nodes class TestDatabricksConnectionMaxIdleTime: diff --git a/tests/unit/test_relation.py b/tests/unit/test_relation.py index 5bde2fa5e..02f939ea2 100644 --- a/tests/unit/test_relation.py +++ b/tests/unit/test_relation.py @@ -1,7 +1,8 @@ import pytest from dbt.adapters.databricks import relation -from dbt.adapters.databricks.relation import DatabricksQuotePolicy, DatabricksRelation +from dbt.adapters.databricks.relation import DatabricksQuotePolicy +from dbt.adapters.databricks.relation import DatabricksRelation class TestDatabricksRelation: diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 81e772113..cbe9ed3db 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1,4 +1,5 @@ -from dbt.adapters.databricks.utils import redact_credentials, remove_ansi +from dbt.adapters.databricks.utils import redact_credentials +from dbt.adapters.databricks.utils import remove_ansi class TestDatabricksUtils: diff --git a/tests/unit/utils.py b/tests/unit/utils.py index 6fdf86cdf..d3e4fed25 100644 --- a/tests/unit/utils.py +++ b/tests/unit/utils.py @@ -1,9 +1,12 @@ import os from copy import deepcopy -from dbt.config import Profile, Project, RuntimeConfig +from dbt.config import Profile +from dbt.config import Project +from dbt.config import RuntimeConfig from dbt.config.project import PartialProject -from dbt.config.renderer import DbtProjectYamlRenderer, ProfileRenderer +from dbt.config.renderer import DbtProjectYamlRenderer +from dbt.config.renderer import ProfileRenderer from dbt.config.utils import parse_cli_vars From c2a3aac5e4bb40182a17be49c66314955ac75f2d Mon Sep 17 00:00:00 2001 From: dkruh1 <101973979+dkruh1@users.noreply.github.com> Date: Mon, 2 Dec 2024 09:01:59 +0200 Subject: [PATCH 6/9] updste precommit --- .pre-commit-config.yaml | 29 ++++++++++++----------------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b95733694..5c06a59e9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,20 +1,15 @@ -# See https://pre-commit.com for more information -# See https://pre-commit.com/hooks.html for more hooks repos: - - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v3.2.0 + - repo: https://github.com/astral-sh/ruff-pre-commit + # Ruff version. + rev: v0.8.0 hooks: - - id: trailing-whitespace - - id: end-of-file-fixer - - id: check-yaml - - id: check-added-large-files - - repo: https://github.com/psf/black-pre-commit-mirror - rev: 24.3.0 + # Run the linter. + - id: ruff + # Run the formatter. + - id: ruff-format + + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.13.0 hooks: - - id: black - language_version: python3.11 - args: [--config, black.ini] - - repo: https://github.com/asottile/reorder-python-imports - rev: v3.12.0 - hooks: - - id: reorder-python-imports + - id: mypy + additional_dependencies: [types-requests] \ No newline at end of file From f30c727a468e4422828b69842d34a9d8bcfa8844 Mon Sep 17 00:00:00 2001 From: dkruh1 <101973979+dkruh1@users.noreply.github.com> Date: Mon, 2 Dec 2024 09:31:02 +0200 Subject: [PATCH 7/9] revert wrong changes --- .pre-commit-config.yaml | 2 +- dbt/adapters/databricks/api_client.py | 9 +- dbt/adapters/databricks/auth.py | 15 +-- dbt/adapters/databricks/behaviors/columns.py | 3 +- dbt/adapters/databricks/column.py | 3 +- dbt/adapters/databricks/connections.py | 102 +++++++++--------- dbt/adapters/databricks/credentials.py | 24 ++--- .../databricks/events/connection_events.py | 3 +- dbt/adapters/databricks/impl.py | 101 ++++++++--------- dbt/adapters/databricks/logging.py | 4 +- .../databricks/python_models/python_config.py | 6 +- .../python_models/python_submissions.py | 10 +- .../databricks/python_models/run_tracking.py | 3 +- dbt/adapters/databricks/relation.py | 14 +-- .../databricks/relation_configs/base.py | 14 +-- .../databricks/relation_configs/comment.py | 9 +- .../relation_configs/incremental.py | 8 +- .../relation_configs/materialized_view.py | 8 +- .../relation_configs/partitioning.py | 9 +- .../databricks/relation_configs/query.py | 6 +- .../databricks/relation_configs/refresh.py | 9 +- .../relation_configs/streaming_table.py | 11 +- .../databricks/relation_configs/tags.py | 9 +- .../relation_configs/tblproperties.py | 10 +- dbt/adapters/databricks/utils.py | 4 +- .../databricks/macros/adapters/columns.sql | 2 +- .../incremental/strategies.sql | 4 +- .../macros/materializations/seeds/helpers.sql | 2 +- .../macros/relations/constraints.sql | 6 +- .../macros/relations/liquid_clustering.sql | 2 +- .../databricks/macros/relations/tags.sql | 2 +- .../databricks/macros/utils/split_part.sql | 2 +- pyproject.toml | 3 - .../adapter/aliases/test_aliases.py | 10 +- .../adapter/basic/test_docs_generate.py | 6 +- .../adapter/basic/test_incremental.py | 3 +- .../adapter/caching/test_caching.py | 10 +- .../adapter/concurrency/test_concurrency.py | 3 +- .../adapter/constraints/test_constraints.py | 18 ++-- .../adapter/dbt_clone/test_dbt_clone.py | 3 +- .../adapter/dbt_show/test_dbt_show.py | 3 +- .../adapter/ephemeral/test_ephemeral.py | 3 +- .../adapter/hooks/test_run_hooks.py | 3 +- .../materialized_view_tests/test_changes.py | 10 +- .../test_persist_constraints.py | 3 +- .../adapter/python_model/test_python_model.py | 6 +- .../query_comments/test_query_comment.py | 14 +-- .../adapter/simple_copy/test_simple_copy.py | 3 +- .../adapter/simple_seed/test_seeds.py | 24 +++-- .../adapter/simple_snapshot/test_snapshot.py | 3 +- .../test_store_failures_as.py | 14 +-- .../adapter/streaming_tables/test_st_basic.py | 4 +- .../adapter/utils/test_source_freshness.py | 3 +- tests/functional/adapter/utils/test_utils.py | 3 +- tests/profiles.py | 3 +- tests/unit/api_client/api_test_base.py | 3 +- tests/unit/api_client/test_command_api.py | 6 +- tests/unit/api_client/test_user_folder_api.py | 3 +- tests/unit/events/test_base_events.py | 3 +- tests/unit/events/test_connection_events.py | 8 +- tests/unit/events/test_cursor_events.py | 3 +- tests/unit/macros/base.py | 5 +- tests/unit/python/test_python_config.py | 8 +- tests/unit/python/test_python_helpers.py | 16 +-- tests/unit/python/test_python_job_support.py | 8 +- tests/unit/python/test_python_submitters.py | 20 ++-- .../python/test_python_workflow_support.py | 6 +- tests/unit/relation_configs/test_comment.py | 3 +- .../test_materialized_view_config.py | 3 +- .../relation_configs/test_partitioning.py | 6 +- tests/unit/relation_configs/test_query.py | 3 +- tests/unit/relation_configs/test_refresh.py | 3 +- tests/unit/relation_configs/test_tags.py | 3 +- .../relation_configs/test_tblproperties.py | 6 +- tests/unit/test_adapter.py | 23 ++-- tests/unit/test_idle_config.py | 3 +- tests/unit/test_relation.py | 3 +- tests/unit/test_utils.py | 3 +- tests/unit/utils.py | 7 +- 79 files changed, 340 insertions(+), 386 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5c06a59e9..81ca0ce98 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -12,4 +12,4 @@ repos: rev: v1.13.0 hooks: - id: mypy - additional_dependencies: [types-requests] \ No newline at end of file + additional_dependencies: [types-requests] diff --git a/dbt/adapters/databricks/api_client.py b/dbt/adapters/databricks/api_client.py index aa86d7402..c14af9f53 100644 --- a/dbt/adapters/databricks/api_client.py +++ b/dbt/adapters/databricks/api_client.py @@ -1,16 +1,13 @@ import base64 import re import time -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod from collections.abc import Callable from dataclasses import dataclass -from typing import Any -from typing import Optional +from typing import Any, Optional from dbt_common.exceptions import DbtRuntimeError -from requests import Response -from requests import Session +from requests import Response, Session from requests.adapters import HTTPAdapter from urllib3.util.retry import Retry diff --git a/dbt/adapters/databricks/auth.py b/dbt/adapters/databricks/auth.py index 8662f794d..439fa23fd 100644 --- a/dbt/adapters/databricks/auth.py +++ b/dbt/adapters/databricks/auth.py @@ -1,16 +1,11 @@ -from typing import Any -from typing import Optional - -from databricks.sdk.core import Config -from databricks.sdk.core import credentials_provider -from databricks.sdk.core import CredentialsProvider -from databricks.sdk.core import HeaderFactory -from databricks.sdk.oauth import ClientCredentials -from databricks.sdk.oauth import Token -from databricks.sdk.oauth import TokenSource +from typing import Any, Optional + from requests import PreparedRequest from requests.auth import AuthBase +from databricks.sdk.core import Config, CredentialsProvider, HeaderFactory, credentials_provider +from databricks.sdk.oauth import ClientCredentials, Token, TokenSource + class token_auth(CredentialsProvider): _token: str diff --git a/dbt/adapters/databricks/behaviors/columns.py b/dbt/adapters/databricks/behaviors/columns.py index fc72e66ec..8f3d1eca7 100644 --- a/dbt/adapters/databricks/behaviors/columns.py +++ b/dbt/adapters/databricks/behaviors/columns.py @@ -1,5 +1,4 @@ -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod from dbt_common.utils.dict import AttrDict diff --git a/dbt/adapters/databricks/column.py b/dbt/adapters/databricks/column.py index f9511201c..4d08ad4d9 100644 --- a/dbt/adapters/databricks/column.py +++ b/dbt/adapters/databricks/column.py @@ -1,6 +1,5 @@ from dataclasses import dataclass -from typing import ClassVar -from typing import Optional +from typing import ClassVar, Optional from dbt.adapters.spark.column import SparkColumn diff --git a/dbt/adapters/databricks/connections.py b/dbt/adapters/databricks/connections.py index 06636c1ec..5e2465c9f 100644 --- a/dbt/adapters/databricks/connections.py +++ b/dbt/adapters/databricks/connections.py @@ -5,77 +5,75 @@ import time import uuid import warnings -from collections.abc import Callable -from collections.abc import Iterator -from collections.abc import Sequence +from collections.abc import Callable, Iterator, Sequence from contextlib import contextmanager from dataclasses import dataclass from multiprocessing.context import SpawnContext from numbers import Number from threading import get_ident -from typing import Any -from typing import cast -from typing import Hashable -from typing import Optional -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Hashable, Optional, cast -import databricks.sql as dbsql -from databricks.sql.client import Connection as DatabricksSQLConnection -from databricks.sql.client import Cursor as DatabricksSQLCursor -from databricks.sql.exc import Error from dbt_common.events.contextvars import get_node_info from dbt_common.events.functions import fire_event -from dbt_common.exceptions import DbtDatabaseError -from dbt_common.exceptions import DbtInternalError -from dbt_common.exceptions import DbtRuntimeError +from dbt_common.exceptions import DbtDatabaseError, DbtInternalError, DbtRuntimeError from dbt_common.utils import cast_to_str from requests import Session +import databricks.sql as dbsql +from databricks.sql.client import Connection as DatabricksSQLConnection +from databricks.sql.client import Cursor as DatabricksSQLCursor +from databricks.sql.exc import Error from dbt.adapters.base.query_headers import MacroQueryStringSetter -from dbt.adapters.contracts.connection import AdapterRequiredConfig -from dbt.adapters.contracts.connection import AdapterResponse -from dbt.adapters.contracts.connection import Connection -from dbt.adapters.contracts.connection import ConnectionState -from dbt.adapters.contracts.connection import DEFAULT_QUERY_COMMENT -from dbt.adapters.contracts.connection import Identifier -from dbt.adapters.contracts.connection import LazyHandle +from dbt.adapters.contracts.connection import ( + DEFAULT_QUERY_COMMENT, + AdapterRequiredConfig, + AdapterResponse, + Connection, + ConnectionState, + Identifier, + LazyHandle, +) from dbt.adapters.databricks.__version__ import version as __version__ from dbt.adapters.databricks.api_client import DatabricksApiClient from dbt.adapters.databricks.auth import BearerAuth -from dbt.adapters.databricks.credentials import DatabricksCredentials -from dbt.adapters.databricks.credentials import TCredentialProvider -from dbt.adapters.databricks.events.connection_events import ConnectionAcquire -from dbt.adapters.databricks.events.connection_events import ConnectionCancel -from dbt.adapters.databricks.events.connection_events import ConnectionCancelError -from dbt.adapters.databricks.events.connection_events import ConnectionClose -from dbt.adapters.databricks.events.connection_events import ConnectionCloseError -from dbt.adapters.databricks.events.connection_events import ConnectionCreate -from dbt.adapters.databricks.events.connection_events import ConnectionCreated -from dbt.adapters.databricks.events.connection_events import ConnectionCreateError -from dbt.adapters.databricks.events.connection_events import ConnectionIdleCheck -from dbt.adapters.databricks.events.connection_events import ConnectionIdleClose -from dbt.adapters.databricks.events.connection_events import ConnectionRelease -from dbt.adapters.databricks.events.connection_events import ConnectionReset -from dbt.adapters.databricks.events.connection_events import ConnectionRetrieve -from dbt.adapters.databricks.events.connection_events import ConnectionReuse -from dbt.adapters.databricks.events.cursor_events import CursorCancel -from dbt.adapters.databricks.events.cursor_events import CursorCancelError -from dbt.adapters.databricks.events.cursor_events import CursorClose -from dbt.adapters.databricks.events.cursor_events import CursorCloseError -from dbt.adapters.databricks.events.cursor_events import CursorCreate +from dbt.adapters.databricks.credentials import DatabricksCredentials, TCredentialProvider +from dbt.adapters.databricks.events.connection_events import ( + ConnectionAcquire, + ConnectionCancel, + ConnectionCancelError, + ConnectionClose, + ConnectionCloseError, + ConnectionCreate, + ConnectionCreated, + ConnectionCreateError, + ConnectionIdleCheck, + ConnectionIdleClose, + ConnectionRelease, + ConnectionReset, + ConnectionRetrieve, + ConnectionReuse, +) +from dbt.adapters.databricks.events.cursor_events import ( + CursorCancel, + CursorCancelError, + CursorClose, + CursorCloseError, + CursorCreate, +) from dbt.adapters.databricks.events.other_events import QueryError -from dbt.adapters.databricks.events.pipeline_events import PipelineRefresh -from dbt.adapters.databricks.events.pipeline_events import PipelineRefreshError +from dbt.adapters.databricks.events.pipeline_events import PipelineRefresh, PipelineRefreshError from dbt.adapters.databricks.logging import logger from dbt.adapters.databricks.python_models.run_tracking import PythonRunTracker from dbt.adapters.databricks.utils import redact_credentials -from dbt.adapters.events.types import ConnectionClosedInCleanup -from dbt.adapters.events.types import ConnectionLeftOpenInCleanup -from dbt.adapters.events.types import ConnectionReused -from dbt.adapters.events.types import ConnectionUsed -from dbt.adapters.events.types import NewConnection -from dbt.adapters.events.types import SQLQuery -from dbt.adapters.events.types import SQLQueryStatus +from dbt.adapters.events.types import ( + ConnectionClosedInCleanup, + ConnectionLeftOpenInCleanup, + ConnectionReused, + ConnectionUsed, + NewConnection, + SQLQuery, + SQLQueryStatus, +) from dbt.adapters.spark.connections import SparkConnectionManager if TYPE_CHECKING: diff --git a/dbt/adapters/databricks/credentials.py b/dbt/adapters/databricks/credentials.py index 30f058091..0f729a5b1 100644 --- a/dbt/adapters/databricks/credentials.py +++ b/dbt/adapters/databricks/credentials.py @@ -5,24 +5,20 @@ import threading from collections.abc import Iterable from dataclasses import dataclass -from typing import Any -from typing import cast -from typing import Optional -from typing import Union +from typing import Any, Optional, Union, cast import keyring -from databricks.sdk.core import CredentialsProvider -from databricks.sdk.oauth import OAuthClient -from databricks.sdk.oauth import SessionCredentials -from dbt_common.exceptions import DbtConfigError -from dbt_common.exceptions import DbtValidationError +from dbt_common.exceptions import DbtConfigError, DbtValidationError +from databricks.sdk.core import CredentialsProvider +from databricks.sdk.oauth import OAuthClient, SessionCredentials from dbt.adapters.contracts.connection import Credentials -from dbt.adapters.databricks.auth import m2m_auth -from dbt.adapters.databricks.auth import token_auth -from dbt.adapters.databricks.events.credential_events import CredentialLoadError -from dbt.adapters.databricks.events.credential_events import CredentialSaveError -from dbt.adapters.databricks.events.credential_events import CredentialShardEvent +from dbt.adapters.databricks.auth import m2m_auth, token_auth +from dbt.adapters.databricks.events.credential_events import ( + CredentialLoadError, + CredentialSaveError, + CredentialShardEvent, +) from dbt.adapters.databricks.logging import logger CATALOG_KEY_IN_SESSION_PROPERTIES = "databricks.catalog" diff --git a/dbt/adapters/databricks/events/connection_events.py b/dbt/adapters/databricks/events/connection_events.py index 0b77c2a16..9f8ec8c1b 100644 --- a/dbt/adapters/databricks/events/connection_events.py +++ b/dbt/adapters/databricks/events/connection_events.py @@ -1,6 +1,5 @@ from abc import ABC -from typing import Any -from typing import Optional +from typing import Any, Optional from databricks.sql.client import Connection diff --git a/dbt/adapters/databricks/impl.py b/dbt/adapters/databricks/impl.py index 8ce1de93c..03df68920 100644 --- a/dbt/adapters/databricks/impl.py +++ b/dbt/adapters/databricks/impl.py @@ -1,67 +1,60 @@ import os import re -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod from collections import defaultdict -from collections.abc import Iterable -from collections.abc import Iterator +from collections.abc import Iterable, Iterator from concurrent.futures import Future from contextlib import contextmanager from dataclasses import dataclass from importlib import metadata from multiprocessing.context import SpawnContext -from typing import Any -from typing import cast -from typing import ClassVar -from typing import Generic -from typing import Optional -from typing import TYPE_CHECKING -from typing import Union +from typing import TYPE_CHECKING, Any, ClassVar, Generic, Optional, Union, cast from uuid import uuid4 from dbt_common.behavior_flags import BehaviorFlag from dbt_common.contracts.config.base import BaseConfig -from dbt_common.exceptions import CompilationError -from dbt_common.exceptions import DbtConfigError -from dbt_common.exceptions import DbtInternalError +from dbt_common.exceptions import CompilationError, DbtConfigError, DbtInternalError from dbt_common.utils import executor from dbt_common.utils.dict import AttrDict from packaging import version -from dbt.adapters.base import AdapterConfig -from dbt.adapters.base import PythonJobHelper -from dbt.adapters.base.impl import catch_as_completed -from dbt.adapters.base.impl import log_code_execution +from dbt.adapters.base import AdapterConfig, PythonJobHelper +from dbt.adapters.base.impl import catch_as_completed, log_code_execution from dbt.adapters.base.meta import available from dbt.adapters.base.relation import BaseRelation -from dbt.adapters.capability import Capability -from dbt.adapters.capability import CapabilityDict -from dbt.adapters.capability import CapabilitySupport -from dbt.adapters.capability import Support -from dbt.adapters.contracts.connection import AdapterResponse -from dbt.adapters.contracts.connection import Connection -from dbt.adapters.contracts.relation import RelationConfig -from dbt.adapters.contracts.relation import RelationType -from dbt.adapters.databricks.behaviors.columns import GetColumnsBehavior -from dbt.adapters.databricks.behaviors.columns import GetColumnsByDescribe -from dbt.adapters.databricks.behaviors.columns import GetColumnsByInformationSchema +from dbt.adapters.capability import Capability, CapabilityDict, CapabilitySupport, Support +from dbt.adapters.contracts.connection import AdapterResponse, Connection +from dbt.adapters.contracts.relation import RelationConfig, RelationType +from dbt.adapters.databricks.behaviors.columns import ( + GetColumnsBehavior, + GetColumnsByDescribe, + GetColumnsByInformationSchema, +) from dbt.adapters.databricks.column import DatabricksColumn -from dbt.adapters.databricks.connections import DatabricksConnectionManager -from dbt.adapters.databricks.connections import DatabricksDBTConnection -from dbt.adapters.databricks.connections import DatabricksSessionConnectionManager -from dbt.adapters.databricks.connections import DatabricksSQLConnectionWrapper -from dbt.adapters.databricks.connections import ExtendedSessionConnectionManager -from dbt.adapters.databricks.connections import USE_LONG_SESSIONS -from dbt.adapters.databricks.connections import USE_SESSION_CONNECTION -from dbt.adapters.databricks.python_models.python_submissions import AllPurposeClusterPythonJobHelper -from dbt.adapters.databricks.python_models.python_submissions import JobClusterPythonJobHelper -from dbt.adapters.databricks.python_models.python_submissions import ServerlessClusterPythonJobHelper -from dbt.adapters.databricks.python_models.python_submissions import WorkflowPythonJobHelper -from dbt.adapters.databricks.relation import DatabricksRelation -from dbt.adapters.databricks.relation import DatabricksRelationType -from dbt.adapters.databricks.relation import KEY_TABLE_PROVIDER -from dbt.adapters.databricks.relation_configs.base import DatabricksRelationConfig -from dbt.adapters.databricks.relation_configs.base import DatabricksRelationConfigBase +from dbt.adapters.databricks.connections import ( + USE_LONG_SESSIONS, + USE_SESSION_CONNECTION, + DatabricksConnectionManager, + DatabricksDBTConnection, + DatabricksSQLConnectionWrapper, + DatabricksSessionConnectionManager, + ExtendedSessionConnectionManager, +) +from dbt.adapters.databricks.python_models.python_submissions import ( + AllPurposeClusterPythonJobHelper, + JobClusterPythonJobHelper, + ServerlessClusterPythonJobHelper, + WorkflowPythonJobHelper, +) +from dbt.adapters.databricks.relation import ( + KEY_TABLE_PROVIDER, + DatabricksRelation, + DatabricksRelationType, +) +from dbt.adapters.databricks.relation_configs.base import ( + DatabricksRelationConfig, + DatabricksRelationConfigBase, +) from dbt.adapters.databricks.relation_configs.incremental import IncrementalTableConfig from dbt.adapters.databricks.relation_configs.materialized_view import ( MaterializedViewConfig, @@ -71,16 +64,16 @@ ) from dbt.adapters.databricks.relation_configs.table_format import TableFormat from dbt.adapters.databricks.relation_configs.tblproperties import TblPropertiesConfig -from dbt.adapters.databricks.utils import get_first_row -from dbt.adapters.databricks.utils import handle_missing_objects -from dbt.adapters.databricks.utils import redact_credentials +from dbt.adapters.databricks.utils import get_first_row, handle_missing_objects, redact_credentials from dbt.adapters.relation_configs import RelationResults -from dbt.adapters.spark.impl import DESCRIBE_TABLE_EXTENDED_MACRO_NAME -from dbt.adapters.spark.impl import GET_COLUMNS_IN_RELATION_RAW_MACRO_NAME -from dbt.adapters.spark.impl import KEY_TABLE_OWNER -from dbt.adapters.spark.impl import KEY_TABLE_STATISTICS -from dbt.adapters.spark.impl import LIST_SCHEMAS_MACRO_NAME -from dbt.adapters.spark.impl import SparkAdapter +from dbt.adapters.spark.impl import ( + DESCRIBE_TABLE_EXTENDED_MACRO_NAME, + GET_COLUMNS_IN_RELATION_RAW_MACRO_NAME, + KEY_TABLE_OWNER, + KEY_TABLE_STATISTICS, + LIST_SCHEMAS_MACRO_NAME, + SparkAdapter, +) if TYPE_CHECKING: from agate import Row, Table diff --git a/dbt/adapters/databricks/logging.py b/dbt/adapters/databricks/logging.py index c309e12c7..d0f1d42ba 100644 --- a/dbt/adapters/databricks/logging.py +++ b/dbt/adapters/databricks/logging.py @@ -1,7 +1,5 @@ import os -from logging import getLogger -from logging import Handler -from logging import LogRecord +from logging import Handler, LogRecord, getLogger from typing import Union from dbt.adapters.events.logging import AdapterLogger diff --git a/dbt/adapters/databricks/python_models/python_config.py b/dbt/adapters/databricks/python_models/python_config.py index 761b0ef65..29aa44efa 100644 --- a/dbt/adapters/databricks/python_models/python_config.py +++ b/dbt/adapters/databricks/python_models/python_config.py @@ -1,9 +1,7 @@ import uuid -from typing import Any -from typing import Optional +from typing import Any, Optional -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field DEFAULT_TIMEOUT = 60 * 60 * 24 diff --git a/dbt/adapters/databricks/python_models/python_submissions.py b/dbt/adapters/databricks/python_models/python_submissions.py index b241fb358..afcb383c2 100644 --- a/dbt/adapters/databricks/python_models/python_submissions.py +++ b/dbt/adapters/databricks/python_models/python_submissions.py @@ -1,16 +1,12 @@ -from abc import ABC -from abc import abstractmethod -from typing import Any -from typing import Optional +from abc import ABC, abstractmethod +from typing import Any, Optional from dbt_common.exceptions import DbtRuntimeError from pydantic import BaseModel from typing_extensions import override from dbt.adapters.base import PythonJobHelper -from dbt.adapters.databricks.api_client import CommandExecution -from dbt.adapters.databricks.api_client import DatabricksApiClient -from dbt.adapters.databricks.api_client import WorkflowJobApi +from dbt.adapters.databricks.api_client import CommandExecution, DatabricksApiClient, WorkflowJobApi from dbt.adapters.databricks.credentials import DatabricksCredentials from dbt.adapters.databricks.logging import logger from dbt.adapters.databricks.python_models.python_config import ParsedPythonModel diff --git a/dbt/adapters/databricks/python_models/run_tracking.py b/dbt/adapters/databricks/python_models/run_tracking.py index ae703cdd1..e48dae7d4 100644 --- a/dbt/adapters/databricks/python_models/run_tracking.py +++ b/dbt/adapters/databricks/python_models/run_tracking.py @@ -2,8 +2,7 @@ from dbt_common.exceptions import DbtRuntimeError -from dbt.adapters.databricks.api_client import CommandExecution -from dbt.adapters.databricks.api_client import DatabricksApiClient +from dbt.adapters.databricks.api_client import CommandExecution, DatabricksApiClient from dbt.adapters.databricks.logging import logger diff --git a/dbt/adapters/databricks/relation.py b/dbt/adapters/databricks/relation.py index ac6189384..bc1dc991f 100644 --- a/dbt/adapters/databricks/relation.py +++ b/dbt/adapters/databricks/relation.py @@ -1,23 +1,17 @@ from collections.abc import Iterable -from dataclasses import dataclass -from dataclasses import field -from typing import Any -from typing import Optional -from typing import Type +from dataclasses import dataclass, field +from typing import Any, Optional, Type from dbt_common.dataclass_schema import StrEnum from dbt_common.exceptions import DbtRuntimeError from dbt_common.utils import filter_null_values -from dbt.adapters.base.relation import BaseRelation -from dbt.adapters.base.relation import InformationSchema -from dbt.adapters.base.relation import Policy +from dbt.adapters.base.relation import BaseRelation, InformationSchema, Policy from dbt.adapters.contracts.relation import ( ComponentName, ) from dbt.adapters.databricks.utils import remove_undefined -from dbt.adapters.spark.impl import KEY_TABLE_OWNER -from dbt.adapters.spark.impl import KEY_TABLE_STATISTICS +from dbt.adapters.spark.impl import KEY_TABLE_OWNER, KEY_TABLE_STATISTICS from dbt.adapters.utils import classproperty KEY_TABLE_PROVIDER = "Provider" diff --git a/dbt/adapters/databricks/relation_configs/base.py b/dbt/adapters/databricks/relation_configs/base.py index 9366226c3..53e5bdb8d 100644 --- a/dbt/adapters/databricks/relation_configs/base.py +++ b/dbt/adapters/databricks/relation_configs/base.py @@ -1,13 +1,7 @@ -from abc import ABC -from abc import abstractmethod -from typing import Any -from typing import ClassVar -from typing import Generic -from typing import Optional -from typing import TypeVar - -from pydantic import BaseModel -from pydantic import ConfigDict +from abc import ABC, abstractmethod +from typing import Any, ClassVar, Generic, Optional, TypeVar + +from pydantic import BaseModel, ConfigDict from typing_extensions import Self from dbt.adapters.contracts.relation import RelationConfig diff --git a/dbt/adapters/databricks/relation_configs/comment.py b/dbt/adapters/databricks/relation_configs/comment.py index 0af362d1a..3dfa07bc1 100644 --- a/dbt/adapters/databricks/relation_configs/comment.py +++ b/dbt/adapters/databricks/relation_configs/comment.py @@ -1,9 +1,10 @@ -from typing import ClassVar -from typing import Optional +from typing import ClassVar, Optional from dbt.adapters.contracts.relation import RelationConfig -from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig -from dbt.adapters.databricks.relation_configs.base import DatabricksComponentProcessor +from dbt.adapters.databricks.relation_configs.base import ( + DatabricksComponentConfig, + DatabricksComponentProcessor, +) from dbt.adapters.relation_configs.config_base import RelationResults diff --git a/dbt/adapters/databricks/relation_configs/incremental.py b/dbt/adapters/databricks/relation_configs/incremental.py index f31baf581..156f5bdf7 100644 --- a/dbt/adapters/databricks/relation_configs/incremental.py +++ b/dbt/adapters/databricks/relation_configs/incremental.py @@ -1,8 +1,10 @@ from typing import Optional -from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig -from dbt.adapters.databricks.relation_configs.base import DatabricksRelationChangeSet -from dbt.adapters.databricks.relation_configs.base import DatabricksRelationConfigBase +from dbt.adapters.databricks.relation_configs.base import ( + DatabricksComponentConfig, + DatabricksRelationChangeSet, + DatabricksRelationConfigBase, +) from dbt.adapters.databricks.relation_configs.tags import TagsProcessor from dbt.adapters.databricks.relation_configs.tblproperties import TblPropertiesProcessor diff --git a/dbt/adapters/databricks/relation_configs/materialized_view.py b/dbt/adapters/databricks/relation_configs/materialized_view.py index 9a57edd9f..e6a8a47e8 100644 --- a/dbt/adapters/databricks/relation_configs/materialized_view.py +++ b/dbt/adapters/databricks/relation_configs/materialized_view.py @@ -1,8 +1,10 @@ from typing import Optional -from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig -from dbt.adapters.databricks.relation_configs.base import DatabricksRelationChangeSet -from dbt.adapters.databricks.relation_configs.base import DatabricksRelationConfigBase +from dbt.adapters.databricks.relation_configs.base import ( + DatabricksComponentConfig, + DatabricksRelationChangeSet, + DatabricksRelationConfigBase, +) from dbt.adapters.databricks.relation_configs.comment import ( CommentProcessor, ) diff --git a/dbt/adapters/databricks/relation_configs/partitioning.py b/dbt/adapters/databricks/relation_configs/partitioning.py index 89120b8ac..bba4b5720 100644 --- a/dbt/adapters/databricks/relation_configs/partitioning.py +++ b/dbt/adapters/databricks/relation_configs/partitioning.py @@ -1,11 +1,12 @@ import itertools -from typing import ClassVar -from typing import Union +from typing import ClassVar, Union from dbt.adapters.contracts.relation import RelationConfig from dbt.adapters.databricks.relation_configs import base -from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig -from dbt.adapters.databricks.relation_configs.base import DatabricksComponentProcessor +from dbt.adapters.databricks.relation_configs.base import ( + DatabricksComponentConfig, + DatabricksComponentProcessor, +) from dbt.adapters.relation_configs.config_base import RelationResults diff --git a/dbt/adapters/databricks/relation_configs/query.py b/dbt/adapters/databricks/relation_configs/query.py index 469a117c6..e9b8f73b3 100644 --- a/dbt/adapters/databricks/relation_configs/query.py +++ b/dbt/adapters/databricks/relation_configs/query.py @@ -3,8 +3,10 @@ from dbt_common.exceptions import DbtRuntimeError from dbt.adapters.contracts.relation import RelationConfig -from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig -from dbt.adapters.databricks.relation_configs.base import DatabricksComponentProcessor +from dbt.adapters.databricks.relation_configs.base import ( + DatabricksComponentConfig, + DatabricksComponentProcessor, +) from dbt.adapters.relation_configs.config_base import RelationResults diff --git a/dbt/adapters/databricks/relation_configs/refresh.py b/dbt/adapters/databricks/relation_configs/refresh.py index c3637297c..bfea49e16 100644 --- a/dbt/adapters/databricks/relation_configs/refresh.py +++ b/dbt/adapters/databricks/relation_configs/refresh.py @@ -1,13 +1,14 @@ import re -from typing import ClassVar -from typing import Optional +from typing import ClassVar, Optional from dbt_common.exceptions import DbtRuntimeError from dbt.adapters.contracts.relation import RelationConfig from dbt.adapters.databricks.relation_configs import base -from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig -from dbt.adapters.databricks.relation_configs.base import DatabricksComponentProcessor +from dbt.adapters.databricks.relation_configs.base import ( + DatabricksComponentConfig, + DatabricksComponentProcessor, +) from dbt.adapters.relation_configs.config_base import RelationResults SCHEDULE_REGEX = re.compile(r"CRON '(.*)' AT TIME ZONE '(.*)'") diff --git a/dbt/adapters/databricks/relation_configs/streaming_table.py b/dbt/adapters/databricks/relation_configs/streaming_table.py index d891da13d..f07ed3669 100644 --- a/dbt/adapters/databricks/relation_configs/streaming_table.py +++ b/dbt/adapters/databricks/relation_configs/streaming_table.py @@ -1,16 +1,17 @@ from typing import Optional -from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig -from dbt.adapters.databricks.relation_configs.base import DatabricksRelationChangeSet -from dbt.adapters.databricks.relation_configs.base import DatabricksRelationConfigBase +from dbt.adapters.databricks.relation_configs.base import ( + DatabricksComponentConfig, + DatabricksRelationChangeSet, + DatabricksRelationConfigBase, +) from dbt.adapters.databricks.relation_configs.comment import ( CommentProcessor, ) from dbt.adapters.databricks.relation_configs.partitioning import ( PartitionedByProcessor, ) -from dbt.adapters.databricks.relation_configs.refresh import RefreshConfig -from dbt.adapters.databricks.relation_configs.refresh import RefreshProcessor +from dbt.adapters.databricks.relation_configs.refresh import RefreshConfig, RefreshProcessor from dbt.adapters.databricks.relation_configs.tblproperties import ( TblPropertiesProcessor, ) diff --git a/dbt/adapters/databricks/relation_configs/tags.py b/dbt/adapters/databricks/relation_configs/tags.py index f8c734ac8..461fa77b2 100644 --- a/dbt/adapters/databricks/relation_configs/tags.py +++ b/dbt/adapters/databricks/relation_configs/tags.py @@ -1,12 +1,13 @@ -from typing import ClassVar -from typing import Optional +from typing import ClassVar, Optional from dbt_common.exceptions import DbtRuntimeError from dbt.adapters.contracts.relation import RelationConfig from dbt.adapters.databricks.relation_configs import base -from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig -from dbt.adapters.databricks.relation_configs.base import DatabricksComponentProcessor +from dbt.adapters.databricks.relation_configs.base import ( + DatabricksComponentConfig, + DatabricksComponentProcessor, +) from dbt.adapters.relation_configs.config_base import RelationResults diff --git a/dbt/adapters/databricks/relation_configs/tblproperties.py b/dbt/adapters/databricks/relation_configs/tblproperties.py index 84e9b8e50..41d18f3fd 100644 --- a/dbt/adapters/databricks/relation_configs/tblproperties.py +++ b/dbt/adapters/databricks/relation_configs/tblproperties.py @@ -1,13 +1,13 @@ -from typing import Any -from typing import ClassVar -from typing import Optional +from typing import Any, ClassVar, Optional from dbt_common.exceptions import DbtRuntimeError from dbt.adapters.contracts.relation import RelationConfig from dbt.adapters.databricks.relation_configs import base -from dbt.adapters.databricks.relation_configs.base import DatabricksComponentConfig -from dbt.adapters.databricks.relation_configs.base import DatabricksComponentProcessor +from dbt.adapters.databricks.relation_configs.base import ( + DatabricksComponentConfig, + DatabricksComponentProcessor, +) from dbt.adapters.relation_configs.config_base import RelationResults diff --git a/dbt/adapters/databricks/utils.py b/dbt/adapters/databricks/utils.py index 6ba06331b..616e5b368 100644 --- a/dbt/adapters/databricks/utils.py +++ b/dbt/adapters/databricks/utils.py @@ -1,8 +1,6 @@ import re from collections.abc import Callable -from typing import Any -from typing import TYPE_CHECKING -from typing import TypeVar +from typing import TYPE_CHECKING, Any, TypeVar from dbt_common.exceptions import DbtRuntimeError from jinja2 import Undefined diff --git a/dbt/include/databricks/macros/adapters/columns.sql b/dbt/include/databricks/macros/adapters/columns.sql index 1fc64ebf5..e1fc1d116 100644 --- a/dbt/include/databricks/macros/adapters/columns.sql +++ b/dbt/include/databricks/macros/adapters/columns.sql @@ -19,7 +19,7 @@ from `system`.`information_schema`.`columns` where table_catalog = '{{ relation.database|lower }}' and - table_schema = '{{ relation.schema|lower }}' and + table_schema = '{{ relation.schema|lower }}' and table_name = '{{ relation.identifier|lower }}' {% endcall %} diff --git a/dbt/include/databricks/macros/materializations/incremental/strategies.sql b/dbt/include/databricks/macros/materializations/incremental/strategies.sql index be272d637..57db5496a 100644 --- a/dbt/include/databricks/macros/materializations/incremental/strategies.sql +++ b/dbt/include/databricks/macros/materializations/incremental/strategies.sql @@ -91,7 +91,7 @@ select {{source_cols_csv}} from {{ source_relation }} {%- set not_matched_by_source_action = config.get('not_matched_by_source_action') -%} {%- set not_matched_by_source_condition = config.get('not_matched_by_source_condition') -%} - + {% if unique_key %} {% if unique_key is sequence and unique_key is not mapping and unique_key is not string %} @@ -184,4 +184,4 @@ select {{source_cols_csv}} from {{ source_relation }} {%- endif -%} {%- do arg_dict.update({'incremental_predicates': incremental_predicates}) -%} {{ return(get_replace_where_sql(arg_dict)) }} -{% endmacro %} +{% endmacro %} \ No newline at end of file diff --git a/dbt/include/databricks/macros/materializations/seeds/helpers.sql b/dbt/include/databricks/macros/materializations/seeds/helpers.sql index a2d2af6d7..df690f18b 100644 --- a/dbt/include/databricks/macros/materializations/seeds/helpers.sql +++ b/dbt/include/databricks/macros/materializations/seeds/helpers.sql @@ -77,7 +77,7 @@ {%- set type = column_override.get(col_name, inferred_type) -%} {%- set column_name = (col_name | string) -%} {%- set column_comment_clause = "" -%} - {%- if column_comment and col_name in model.columns.keys() -%} + {%- if column_comment and col_name in model.columns.keys() -%} {%- set comment = model.columns[col_name]['description'] | replace("'", "\\'") -%} {%- if comment and comment != "" -%} {%- set column_comment_clause = "comment '" ~ comment ~ "'" -%} diff --git a/dbt/include/databricks/macros/relations/constraints.sql b/dbt/include/databricks/macros/relations/constraints.sql index d0a46b151..bb77145ff 100644 --- a/dbt/include/databricks/macros/relations/constraints.sql +++ b/dbt/include/databricks/macros/relations/constraints.sql @@ -179,7 +179,7 @@ {% endif %} {% set name = constraint.get("name") %} - + {% if constraint.get('expression') %} {% if not name %} @@ -188,7 +188,7 @@ {%- set name = local_md5("foreign_key;" ~ relation.identifier ~ ";" ~ constraint.get('expression') ~ ";") -%} {% else %} {{ exceptions.raise_compiler_error("Constraint of type " ~ type ~ " with no `name` provided, and no md5 utility.") }} - {% endif %} + {% endif %} {% endif %} {% set stmt = "alter table " ~ relation ~ " add constraint " ~ name ~ " foreign key" ~ constraint.get('expression') %} @@ -224,7 +224,7 @@ {%- set name = local_md5("foreign_key;" ~ relation.identifier ~ ";" ~ column_names ~ ";" ~ parent ~ ";") -%} {% else %} {{ exceptions.raise_compiler_error("Constraint of type " ~ type ~ " with no `name` provided, and no md5 utility.") }} - {% endif %} + {% endif %} {% endif %} {% set stmt = "alter table " ~ relation ~ " add constraint " ~ name ~ " foreign key(" ~ joined_names ~ ") references " ~ parent %} diff --git a/dbt/include/databricks/macros/relations/liquid_clustering.sql b/dbt/include/databricks/macros/relations/liquid_clustering.sql index a034e6d12..3cf810488 100644 --- a/dbt/include/databricks/macros/relations/liquid_clustering.sql +++ b/dbt/include/databricks/macros/relations/liquid_clustering.sql @@ -18,4 +18,4 @@ ALTER {{ target_relation.type }} {{ target_relation }} CLUSTER BY ({{ cols | join(', ') }}) {%- endcall -%} {%- endif %} -{%- endmacro -%} +{%- endmacro -%} \ No newline at end of file diff --git a/dbt/include/databricks/macros/relations/tags.sql b/dbt/include/databricks/macros/relations/tags.sql index d26da9a20..3467631df 100644 --- a/dbt/include/databricks/macros/relations/tags.sql +++ b/dbt/include/databricks/macros/relations/tags.sql @@ -11,7 +11,7 @@ {% macro fetch_tags_sql(relation) -%} SELECT tag_name, tag_value FROM `system`.`information_schema`.`table_tags` - WHERE catalog_name = '{{ relation.database|lower }}' + WHERE catalog_name = '{{ relation.database|lower }}' AND schema_name = '{{ relation.schema|lower }}' AND table_name = '{{ relation.identifier|lower }}' {%- endmacro -%} diff --git a/dbt/include/databricks/macros/utils/split_part.sql b/dbt/include/databricks/macros/utils/split_part.sql index 3f2773500..92b0cecd0 100644 --- a/dbt/include/databricks/macros/utils/split_part.sql +++ b/dbt/include/databricks/macros/utils/split_part.sql @@ -27,7 +27,7 @@ get(split( {{ string_text }}, {{ delimiter_expr }} - ), + ), length({{ string_text }}) - length( replace({{ string_text }}, {{ delimiter_text }}, '') diff --git a/pyproject.toml b/pyproject.toml index 287ab5327..15efc9e4f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -119,7 +119,6 @@ matrix.session_support.env-vars = [ line-length = 100 target-version = 'py39' - [tool.ruff.lint] select = ["E", "W", "F", "I"] ignore = ["E203"] @@ -141,8 +140,6 @@ markers = [ strict_optional = true no_implicit_optional = true disallow_untyped_defs = true -allow_redefinition = true - [[tool.mypy.overrides]] module = "tests.*" diff --git a/tests/functional/adapter/aliases/test_aliases.py b/tests/functional/adapter/aliases/test_aliases.py index a0f4495e9..f0f90823d 100644 --- a/tests/functional/adapter/aliases/test_aliases.py +++ b/tests/functional/adapter/aliases/test_aliases.py @@ -1,10 +1,12 @@ import pytest from dbt.tests.adapter.aliases import fixtures as dbt_fixtures -from dbt.tests.adapter.aliases.test_aliases import BaseAliasErrors -from dbt.tests.adapter.aliases.test_aliases import BaseAliases -from dbt.tests.adapter.aliases.test_aliases import BaseSameAliasDifferentDatabases -from dbt.tests.adapter.aliases.test_aliases import BaseSameAliasDifferentSchemas +from dbt.tests.adapter.aliases.test_aliases import ( + BaseAliasErrors, + BaseAliases, + BaseSameAliasDifferentDatabases, + BaseSameAliasDifferentSchemas, +) from tests.functional.adapter.aliases import fixtures as databricks_fixtures macro_override = { diff --git a/tests/functional/adapter/basic/test_docs_generate.py b/tests/functional/adapter/basic/test_docs_generate.py index 677879df4..83593eb5b 100644 --- a/tests/functional/adapter/basic/test_docs_generate.py +++ b/tests/functional/adapter/basic/test_docs_generate.py @@ -1,11 +1,9 @@ import pytest from dbt.tests.adapter.basic import expected_catalog -from dbt.tests.adapter.basic.test_docs_generate import BaseDocsGenerate -from dbt.tests.adapter.basic.test_docs_generate import BaseDocsGenReferences +from dbt.tests.adapter.basic.test_docs_generate import BaseDocsGenerate, BaseDocsGenReferences from dbt.tests.util import AnyString -from tests.functional.adapter.basic.typing import AnyLongType -from tests.functional.adapter.basic.typing import StatsLikeDict +from tests.functional.adapter.basic.typing import AnyLongType, StatsLikeDict class TestDocsGenerate(BaseDocsGenerate): diff --git a/tests/functional/adapter/basic/test_incremental.py b/tests/functional/adapter/basic/test_incremental.py index 3958b3304..a353d7802 100644 --- a/tests/functional/adapter/basic/test_incremental.py +++ b/tests/functional/adapter/basic/test_incremental.py @@ -2,8 +2,7 @@ import pytest -from dbt.tests.adapter.basic.test_incremental import BaseIncremental -from dbt.tests.adapter.basic.test_incremental import BaseIncrementalNotSchemaChange +from dbt.tests.adapter.basic.test_incremental import BaseIncremental, BaseIncrementalNotSchemaChange class TestIncrementalDelta(BaseIncremental): diff --git a/tests/functional/adapter/caching/test_caching.py b/tests/functional/adapter/caching/test_caching.py index 206f092f1..d32b3d534 100644 --- a/tests/functional/adapter/caching/test_caching.py +++ b/tests/functional/adapter/caching/test_caching.py @@ -1,7 +1,9 @@ -from dbt.tests.adapter.caching.test_caching import BaseCachingLowercaseModel -from dbt.tests.adapter.caching.test_caching import BaseCachingSelectedSchemaOnly -from dbt.tests.adapter.caching.test_caching import BaseCachingUppercaseModel -from dbt.tests.adapter.caching.test_caching import BaseNoPopulateCache +from dbt.tests.adapter.caching.test_caching import ( + BaseCachingLowercaseModel, + BaseCachingSelectedSchemaOnly, + BaseCachingUppercaseModel, + BaseNoPopulateCache, +) class TestNoPopulateCache(BaseNoPopulateCache): diff --git a/tests/functional/adapter/concurrency/test_concurrency.py b/tests/functional/adapter/concurrency/test_concurrency.py index a4e410940..b1b8aded5 100644 --- a/tests/functional/adapter/concurrency/test_concurrency.py +++ b/tests/functional/adapter/concurrency/test_concurrency.py @@ -1,6 +1,5 @@ from dbt.tests import util -from dbt.tests.adapter.concurrency.test_concurrency import BaseConcurrency -from dbt.tests.adapter.concurrency.test_concurrency import seeds__update_csv +from dbt.tests.adapter.concurrency.test_concurrency import BaseConcurrency, seeds__update_csv # Copied from dbt-core diff --git a/tests/functional/adapter/constraints/test_constraints.py b/tests/functional/adapter/constraints/test_constraints.py index cfd4d1156..350910d00 100644 --- a/tests/functional/adapter/constraints/test_constraints.py +++ b/tests/functional/adapter/constraints/test_constraints.py @@ -2,14 +2,16 @@ from dbt.tests import util from dbt.tests.adapter.constraints import fixtures -from dbt.tests.adapter.constraints.test_constraints import BaseConstraintQuotedColumn -from dbt.tests.adapter.constraints.test_constraints import BaseConstraintsRollback -from dbt.tests.adapter.constraints.test_constraints import BaseConstraintsRuntimeDdlEnforcement -from dbt.tests.adapter.constraints.test_constraints import BaseIncrementalConstraintsColumnsEqual -from dbt.tests.adapter.constraints.test_constraints import BaseIncrementalConstraintsRollback -from dbt.tests.adapter.constraints.test_constraints import BaseIncrementalConstraintsRuntimeDdlEnforcement -from dbt.tests.adapter.constraints.test_constraints import BaseTableConstraintsColumnsEqual -from dbt.tests.adapter.constraints.test_constraints import BaseViewConstraintsColumnsEqual +from dbt.tests.adapter.constraints.test_constraints import ( + BaseConstraintQuotedColumn, + BaseConstraintsRollback, + BaseConstraintsRuntimeDdlEnforcement, + BaseIncrementalConstraintsColumnsEqual, + BaseIncrementalConstraintsRollback, + BaseIncrementalConstraintsRuntimeDdlEnforcement, + BaseTableConstraintsColumnsEqual, + BaseViewConstraintsColumnsEqual, +) from tests.functional.adapter.constraints import fixtures as override_fixtures diff --git a/tests/functional/adapter/dbt_clone/test_dbt_clone.py b/tests/functional/adapter/dbt_clone/test_dbt_clone.py index 291554658..6d619abff 100644 --- a/tests/functional/adapter/dbt_clone/test_dbt_clone.py +++ b/tests/functional/adapter/dbt_clone/test_dbt_clone.py @@ -1,8 +1,7 @@ import pytest from dbt.tests import util -from dbt.tests.adapter.dbt_clone.test_dbt_clone import BaseClone -from dbt.tests.adapter.dbt_clone.test_dbt_clone import BaseClonePossible +from dbt.tests.adapter.dbt_clone.test_dbt_clone import BaseClone, BaseClonePossible from dbt.tests.util import run_dbt from tests.functional.adapter.dbt_clone import fixtures diff --git a/tests/functional/adapter/dbt_show/test_dbt_show.py b/tests/functional/adapter/dbt_show/test_dbt_show.py index 340828c02..98d603154 100644 --- a/tests/functional/adapter/dbt_show/test_dbt_show.py +++ b/tests/functional/adapter/dbt_show/test_dbt_show.py @@ -1,5 +1,4 @@ -from dbt.tests.adapter.dbt_show.test_dbt_show import BaseShowLimit -from dbt.tests.adapter.dbt_show.test_dbt_show import BaseShowSqlHeader +from dbt.tests.adapter.dbt_show.test_dbt_show import BaseShowLimit, BaseShowSqlHeader class TestShowLimit(BaseShowLimit): diff --git a/tests/functional/adapter/ephemeral/test_ephemeral.py b/tests/functional/adapter/ephemeral/test_ephemeral.py index 81a320dbf..c00585b81 100644 --- a/tests/functional/adapter/ephemeral/test_ephemeral.py +++ b/tests/functional/adapter/ephemeral/test_ephemeral.py @@ -5,8 +5,7 @@ from dbt.tests import util from dbt.tests.adapter.ephemeral import test_ephemeral -from dbt.tests.adapter.ephemeral.test_ephemeral import BaseEphemeral -from dbt.tests.adapter.ephemeral.test_ephemeral import BaseEphemeralMulti +from dbt.tests.adapter.ephemeral.test_ephemeral import BaseEphemeral, BaseEphemeralMulti class TestEphemeralMulti(BaseEphemeralMulti): diff --git a/tests/functional/adapter/hooks/test_run_hooks.py b/tests/functional/adapter/hooks/test_run_hooks.py index 91c228437..5c7dd5c2d 100644 --- a/tests/functional/adapter/hooks/test_run_hooks.py +++ b/tests/functional/adapter/hooks/test_run_hooks.py @@ -3,8 +3,7 @@ import pytest from dbt.tests import util -from dbt.tests.adapter.hooks.test_run_hooks import BaseAfterRunHooks -from dbt.tests.adapter.hooks.test_run_hooks import BasePrePostRunHooks +from dbt.tests.adapter.hooks.test_run_hooks import BaseAfterRunHooks, BasePrePostRunHooks from tests.functional.adapter.hooks import fixtures as override_fixtures diff --git a/tests/functional/adapter/materialized_view_tests/test_changes.py b/tests/functional/adapter/materialized_view_tests/test_changes.py index 49f715a1c..1c1fa4c71 100644 --- a/tests/functional/adapter/materialized_view_tests/test_changes.py +++ b/tests/functional/adapter/materialized_view_tests/test_changes.py @@ -8,10 +8,12 @@ ) from dbt.adapters.databricks.relation_configs.tblproperties import TblPropertiesConfig from dbt.tests import util -from dbt.tests.adapter.materialized_view.changes import MaterializedViewChanges -from dbt.tests.adapter.materialized_view.changes import MaterializedViewChangesApplyMixin -from dbt.tests.adapter.materialized_view.changes import MaterializedViewChangesContinueMixin -from dbt.tests.adapter.materialized_view.changes import MaterializedViewChangesFailMixin +from dbt.tests.adapter.materialized_view.changes import ( + MaterializedViewChanges, + MaterializedViewChangesApplyMixin, + MaterializedViewChangesContinueMixin, + MaterializedViewChangesFailMixin, +) from tests.functional.adapter.materialized_view_tests import fixtures diff --git a/tests/functional/adapter/persist_constraints/test_persist_constraints.py b/tests/functional/adapter/persist_constraints/test_persist_constraints.py index 42c9862fa..ac7c50b0a 100644 --- a/tests/functional/adapter/persist_constraints/test_persist_constraints.py +++ b/tests/functional/adapter/persist_constraints/test_persist_constraints.py @@ -1,7 +1,6 @@ import pytest -from dbt.contracts.results import RunResult -from dbt.contracts.results import RunStatus +from dbt.contracts.results import RunResult, RunStatus from dbt.tests import util from tests.functional.adapter.persist_constraints import fixtures diff --git a/tests/functional/adapter/python_model/test_python_model.py b/tests/functional/adapter/python_model/test_python_model.py index 453807041..858214b79 100644 --- a/tests/functional/adapter/python_model/test_python_model.py +++ b/tests/functional/adapter/python_model/test_python_model.py @@ -4,8 +4,10 @@ from dbt.tests import util from dbt.tests.adapter.python_model import test_python_model as fixtures -from dbt.tests.adapter.python_model.test_python_model import BasePythonIncrementalTests -from dbt.tests.adapter.python_model.test_python_model import BasePythonModelTests +from dbt.tests.adapter.python_model.test_python_model import ( + BasePythonIncrementalTests, + BasePythonModelTests, +) from tests.functional.adapter.python_model import fixtures as override_fixtures diff --git a/tests/functional/adapter/query_comments/test_query_comment.py b/tests/functional/adapter/query_comments/test_query_comment.py index 3f05b41e2..8d64835ba 100644 --- a/tests/functional/adapter/query_comments/test_query_comment.py +++ b/tests/functional/adapter/query_comments/test_query_comment.py @@ -1,11 +1,13 @@ import pytest -from dbt.tests.adapter.query_comment.test_query_comment import BaseEmptyQueryComments -from dbt.tests.adapter.query_comment.test_query_comment import BaseMacroArgsQueryComments -from dbt.tests.adapter.query_comment.test_query_comment import BaseMacroInvalidQueryComments -from dbt.tests.adapter.query_comment.test_query_comment import BaseMacroQueryComments -from dbt.tests.adapter.query_comment.test_query_comment import BaseNullQueryComments -from dbt.tests.adapter.query_comment.test_query_comment import BaseQueryComments +from dbt.tests.adapter.query_comment.test_query_comment import ( + BaseEmptyQueryComments, + BaseMacroArgsQueryComments, + BaseMacroInvalidQueryComments, + BaseMacroQueryComments, + BaseNullQueryComments, + BaseQueryComments, +) class TestQueryComments(BaseQueryComments): diff --git a/tests/functional/adapter/simple_copy/test_simple_copy.py b/tests/functional/adapter/simple_copy/test_simple_copy.py index feac39bd5..ba65ad15e 100644 --- a/tests/functional/adapter/simple_copy/test_simple_copy.py +++ b/tests/functional/adapter/simple_copy/test_simple_copy.py @@ -1,7 +1,6 @@ import pytest -from dbt.tests.adapter.simple_copy.test_simple_copy import EmptyModelsArentRunBase -from dbt.tests.adapter.simple_copy.test_simple_copy import SimpleCopyBase +from dbt.tests.adapter.simple_copy.test_simple_copy import EmptyModelsArentRunBase, SimpleCopyBase # Tests with materialized_views, which only works for SQL Warehouse diff --git a/tests/functional/adapter/simple_seed/test_seeds.py b/tests/functional/adapter/simple_seed/test_seeds.py index 7e18b61d2..242f79074 100644 --- a/tests/functional/adapter/simple_seed/test_seeds.py +++ b/tests/functional/adapter/simple_seed/test_seeds.py @@ -4,17 +4,19 @@ from dbt.tests import util from dbt.tests.adapter.simple_seed import seeds -from dbt.tests.adapter.simple_seed.test_seed import BaseSeedConfigFullRefreshOff -from dbt.tests.adapter.simple_seed.test_seed import BaseSeedCustomSchema -from dbt.tests.adapter.simple_seed.test_seed import BaseSeedParsing -from dbt.tests.adapter.simple_seed.test_seed import BaseSeedSpecificFormats -from dbt.tests.adapter.simple_seed.test_seed import BaseSeedWithEmptyDelimiter -from dbt.tests.adapter.simple_seed.test_seed import BaseSeedWithUniqueDelimiter -from dbt.tests.adapter.simple_seed.test_seed import BaseSeedWithWrongDelimiter -from dbt.tests.adapter.simple_seed.test_seed import BaseSimpleSeedEnabledViaConfig -from dbt.tests.adapter.simple_seed.test_seed import BaseSimpleSeedWithBOM -from dbt.tests.adapter.simple_seed.test_seed import BaseTestEmptySeed -from dbt.tests.adapter.simple_seed.test_seed import SeedTestBase +from dbt.tests.adapter.simple_seed.test_seed import ( + BaseSeedConfigFullRefreshOff, + BaseSeedCustomSchema, + BaseSeedParsing, + BaseSeedSpecificFormats, + BaseSeedWithEmptyDelimiter, + BaseSeedWithUniqueDelimiter, + BaseSeedWithWrongDelimiter, + BaseSimpleSeedEnabledViaConfig, + BaseSimpleSeedWithBOM, + BaseTestEmptySeed, + SeedTestBase, +) from tests.functional.adapter.simple_seed import fixtures diff --git a/tests/functional/adapter/simple_snapshot/test_snapshot.py b/tests/functional/adapter/simple_snapshot/test_snapshot.py index ef2b873e9..56186e4ef 100644 --- a/tests/functional/adapter/simple_snapshot/test_snapshot.py +++ b/tests/functional/adapter/simple_snapshot/test_snapshot.py @@ -3,8 +3,7 @@ import pytest from dbt.tests import util -from dbt.tests.adapter.simple_snapshot.test_snapshot import BaseSimpleSnapshot -from dbt.tests.adapter.simple_snapshot.test_snapshot import BaseSnapshotCheck +from dbt.tests.adapter.simple_snapshot.test_snapshot import BaseSimpleSnapshot, BaseSnapshotCheck from dbt.tests.util import run_dbt from tests.functional.adapter.simple_snapshot import fixtures diff --git a/tests/functional/adapter/store_test_failures/test_store_failures_as.py b/tests/functional/adapter/store_test_failures/test_store_failures_as.py index ccebc8115..1d830dcb7 100644 --- a/tests/functional/adapter/store_test_failures/test_store_failures_as.py +++ b/tests/functional/adapter/store_test_failures/test_store_failures_as.py @@ -1,9 +1,11 @@ -from dbt.tests.adapter.store_test_failures_tests.basic import StoreTestFailuresAsExceptions -from dbt.tests.adapter.store_test_failures_tests.basic import StoreTestFailuresAsGeneric -from dbt.tests.adapter.store_test_failures_tests.basic import StoreTestFailuresAsInteractions -from dbt.tests.adapter.store_test_failures_tests.basic import StoreTestFailuresAsProjectLevelEphemeral -from dbt.tests.adapter.store_test_failures_tests.basic import StoreTestFailuresAsProjectLevelOff -from dbt.tests.adapter.store_test_failures_tests.basic import StoreTestFailuresAsProjectLevelView +from dbt.tests.adapter.store_test_failures_tests.basic import ( + StoreTestFailuresAsExceptions, + StoreTestFailuresAsGeneric, + StoreTestFailuresAsInteractions, + StoreTestFailuresAsProjectLevelEphemeral, + StoreTestFailuresAsProjectLevelOff, + StoreTestFailuresAsProjectLevelView, +) from dbt.tests.adapter.store_test_failures_tests.test_store_test_failures import ( BaseStoreTestFailures, ) diff --git a/tests/functional/adapter/streaming_tables/test_st_basic.py b/tests/functional/adapter/streaming_tables/test_st_basic.py index ba95143b2..59f5ea5e7 100644 --- a/tests/functional/adapter/streaming_tables/test_st_basic.py +++ b/tests/functional/adapter/streaming_tables/test_st_basic.py @@ -5,9 +5,7 @@ from dbt.adapters.base.relation import BaseRelation from dbt.adapters.databricks.relation import DatabricksRelationType from dbt.tests import util -from dbt.tests.adapter.materialized_view.files import MY_SEED -from dbt.tests.adapter.materialized_view.files import MY_TABLE -from dbt.tests.adapter.materialized_view.files import MY_VIEW +from dbt.tests.adapter.materialized_view.files import MY_SEED, MY_TABLE, MY_VIEW from tests.functional.adapter.streaming_tables import fixtures diff --git a/tests/functional/adapter/utils/test_source_freshness.py b/tests/functional/adapter/utils/test_source_freshness.py index 5b2024c73..857690b75 100644 --- a/tests/functional/adapter/utils/test_source_freshness.py +++ b/tests/functional/adapter/utils/test_source_freshness.py @@ -2,8 +2,7 @@ import pytest -from dbt.tests.util import get_artifact -from dbt.tests.util import run_dbt +from dbt.tests.util import get_artifact, run_dbt freshness_via_metadata_schema_yml = """ version: 2 diff --git a/tests/functional/adapter/utils/test_utils.py b/tests/functional/adapter/utils/test_utils.py index a6b99634b..06086f25d 100644 --- a/tests/functional/adapter/utils/test_utils.py +++ b/tests/functional/adapter/utils/test_utils.py @@ -20,8 +20,7 @@ from dbt.tests.adapter.utils.test_intersect import BaseIntersect from dbt.tests.adapter.utils.test_last_day import BaseLastDay from dbt.tests.adapter.utils.test_length import BaseLength -from dbt.tests.adapter.utils.test_null_compare import BaseMixedNullCompare -from dbt.tests.adapter.utils.test_null_compare import BaseNullCompare +from dbt.tests.adapter.utils.test_null_compare import BaseMixedNullCompare, BaseNullCompare from dbt.tests.adapter.utils.test_position import BasePosition from dbt.tests.adapter.utils.test_replace import BaseReplace from dbt.tests.adapter.utils.test_right import BaseRight diff --git a/tests/profiles.py b/tests/profiles.py index e34c5073f..e0c88447e 100644 --- a/tests/profiles.py +++ b/tests/profiles.py @@ -1,6 +1,5 @@ import os -from typing import Any -from typing import Optional +from typing import Any, Optional def get_databricks_cluster_target(profile_type: str): diff --git a/tests/unit/api_client/api_test_base.py b/tests/unit/api_client/api_test_base.py index 8377bd19d..3d60dc047 100644 --- a/tests/unit/api_client/api_test_base.py +++ b/tests/unit/api_client/api_test_base.py @@ -1,5 +1,4 @@ -from typing import Any -from typing import Callable +from typing import Any, Callable from unittest.mock import Mock import pytest diff --git a/tests/unit/api_client/test_command_api.py b/tests/unit/api_client/test_command_api.py index 57a94ae16..90efaf472 100644 --- a/tests/unit/api_client/test_command_api.py +++ b/tests/unit/api_client/test_command_api.py @@ -1,12 +1,10 @@ -from unittest.mock import Mock -from unittest.mock import patch +from unittest.mock import Mock, patch import freezegun import pytest from dbt_common.exceptions import DbtRuntimeError -from dbt.adapters.databricks.api_client import CommandApi -from dbt.adapters.databricks.api_client import CommandExecution +from dbt.adapters.databricks.api_client import CommandApi, CommandExecution from tests.unit.api_client.api_test_base import ApiTestBase diff --git a/tests/unit/api_client/test_user_folder_api.py b/tests/unit/api_client/test_user_folder_api.py index ca540dce4..c0ab4b081 100644 --- a/tests/unit/api_client/test_user_folder_api.py +++ b/tests/unit/api_client/test_user_folder_api.py @@ -1,7 +1,6 @@ import pytest -from dbt.adapters.databricks.api_client import CurrUserApi -from dbt.adapters.databricks.api_client import UserFolderApi +from dbt.adapters.databricks.api_client import CurrUserApi, UserFolderApi from tests.unit.api_client.api_test_base import ApiTestBase diff --git a/tests/unit/events/test_base_events.py b/tests/unit/events/test_base_events.py index 990c51a42..7b480029d 100644 --- a/tests/unit/events/test_base_events.py +++ b/tests/unit/events/test_base_events.py @@ -1,7 +1,6 @@ from databricks.sql.exc import Error -from dbt.adapters.databricks.events.base import ErrorEvent -from dbt.adapters.databricks.events.base import SQLErrorEvent +from dbt.adapters.databricks.events.base import ErrorEvent, SQLErrorEvent class ErrorTestEvent(ErrorEvent): diff --git a/tests/unit/events/test_connection_events.py b/tests/unit/events/test_connection_events.py index 995110ca3..81a30c1dc 100644 --- a/tests/unit/events/test_connection_events.py +++ b/tests/unit/events/test_connection_events.py @@ -1,8 +1,10 @@ from unittest.mock import Mock -from dbt.adapters.databricks.events.connection_events import ConnectionAcquire -from dbt.adapters.databricks.events.connection_events import ConnectionCloseError -from dbt.adapters.databricks.events.connection_events import ConnectionEvent +from dbt.adapters.databricks.events.connection_events import ( + ConnectionAcquire, + ConnectionCloseError, + ConnectionEvent, +) class ConnectionTestEvent(ConnectionEvent): diff --git a/tests/unit/events/test_cursor_events.py b/tests/unit/events/test_cursor_events.py index e30caa849..e492415a6 100644 --- a/tests/unit/events/test_cursor_events.py +++ b/tests/unit/events/test_cursor_events.py @@ -1,7 +1,6 @@ from unittest.mock import Mock -from dbt.adapters.databricks.events.cursor_events import CursorCloseError -from dbt.adapters.databricks.events.cursor_events import CursorEvent +from dbt.adapters.databricks.events.cursor_events import CursorCloseError, CursorEvent class CursorTestEvent(CursorEvent): diff --git a/tests/unit/macros/base.py b/tests/unit/macros/base.py index 02c8cf94a..cb9c7f282 100644 --- a/tests/unit/macros/base.py +++ b/tests/unit/macros/base.py @@ -3,10 +3,7 @@ from unittest.mock import Mock import pytest -from jinja2 import Environment -from jinja2 import FileSystemLoader -from jinja2 import PackageLoader -from jinja2 import Template +from jinja2 import Environment, FileSystemLoader, PackageLoader, Template from dbt.adapters.databricks.relation import DatabricksRelation diff --git a/tests/unit/python/test_python_config.py b/tests/unit/python/test_python_config.py index 6e00cd6f5..3559e044b 100644 --- a/tests/unit/python/test_python_config.py +++ b/tests/unit/python/test_python_config.py @@ -1,9 +1,11 @@ import pytest from pydantic import ValidationError -from dbt.adapters.databricks.python_models.python_config import ParsedPythonModel -from dbt.adapters.databricks.python_models.python_config import PythonJobConfig -from dbt.adapters.databricks.python_models.python_config import PythonModelConfig +from dbt.adapters.databricks.python_models.python_config import ( + ParsedPythonModel, + PythonJobConfig, + PythonModelConfig, +) class TestParsedPythonModel: diff --git a/tests/unit/python/test_python_helpers.py b/tests/unit/python/test_python_helpers.py index 22f6f5f0c..8c93bf20e 100644 --- a/tests/unit/python/test_python_helpers.py +++ b/tests/unit/python/test_python_helpers.py @@ -2,13 +2,15 @@ import pytest -from dbt.adapters.databricks.python_models.python_submissions import AllPurposeClusterPythonJobHelper -from dbt.adapters.databricks.python_models.python_submissions import JobClusterPythonJobHelper -from dbt.adapters.databricks.python_models.python_submissions import PythonCommandSubmitter -from dbt.adapters.databricks.python_models.python_submissions import PythonNotebookSubmitter -from dbt.adapters.databricks.python_models.python_submissions import PythonNotebookWorkflowSubmitter -from dbt.adapters.databricks.python_models.python_submissions import ServerlessClusterPythonJobHelper -from dbt.adapters.databricks.python_models.python_submissions import WorkflowPythonJobHelper +from dbt.adapters.databricks.python_models.python_submissions import ( + AllPurposeClusterPythonJobHelper, + JobClusterPythonJobHelper, + PythonCommandSubmitter, + PythonNotebookSubmitter, + PythonNotebookWorkflowSubmitter, + ServerlessClusterPythonJobHelper, + WorkflowPythonJobHelper, +) @pytest.fixture diff --git a/tests/unit/python/test_python_job_support.py b/tests/unit/python/test_python_job_support.py index 93dfe3b46..41f480413 100644 --- a/tests/unit/python/test_python_job_support.py +++ b/tests/unit/python/test_python_job_support.py @@ -3,9 +3,11 @@ import pytest from dbt.adapters.databricks.python_models import python_submissions -from dbt.adapters.databricks.python_models.python_submissions import PythonJobConfigCompiler -from dbt.adapters.databricks.python_models.python_submissions import PythonNotebookUploader -from dbt.adapters.databricks.python_models.python_submissions import PythonPermissionBuilder +from dbt.adapters.databricks.python_models.python_submissions import ( + PythonJobConfigCompiler, + PythonNotebookUploader, + PythonPermissionBuilder, +) @pytest.fixture diff --git a/tests/unit/python/test_python_submitters.py b/tests/unit/python/test_python_submitters.py index 1abd50ad9..cf7077302 100644 --- a/tests/unit/python/test_python_submitters.py +++ b/tests/unit/python/test_python_submitters.py @@ -2,15 +2,17 @@ import pytest -from dbt.adapters.databricks.python_models.python_submissions import PythonCommandSubmitter -from dbt.adapters.databricks.python_models.python_submissions import PythonJobConfigCompiler -from dbt.adapters.databricks.python_models.python_submissions import PythonJobDetails -from dbt.adapters.databricks.python_models.python_submissions import PythonNotebookSubmitter -from dbt.adapters.databricks.python_models.python_submissions import PythonNotebookUploader -from dbt.adapters.databricks.python_models.python_submissions import PythonNotebookWorkflowSubmitter -from dbt.adapters.databricks.python_models.python_submissions import PythonPermissionBuilder -from dbt.adapters.databricks.python_models.python_submissions import PythonWorkflowConfigCompiler -from dbt.adapters.databricks.python_models.python_submissions import PythonWorkflowCreator +from dbt.adapters.databricks.python_models.python_submissions import ( + PythonCommandSubmitter, + PythonJobConfigCompiler, + PythonJobDetails, + PythonNotebookSubmitter, + PythonNotebookUploader, + PythonNotebookWorkflowSubmitter, + PythonPermissionBuilder, + PythonWorkflowConfigCompiler, + PythonWorkflowCreator, +) @pytest.fixture diff --git a/tests/unit/python/test_python_workflow_support.py b/tests/unit/python/test_python_workflow_support.py index 9038b54ff..595bd90c9 100644 --- a/tests/unit/python/test_python_workflow_support.py +++ b/tests/unit/python/test_python_workflow_support.py @@ -2,8 +2,10 @@ import pytest -from dbt.adapters.databricks.python_models.python_submissions import PythonWorkflowConfigCompiler -from dbt.adapters.databricks.python_models.python_submissions import PythonWorkflowCreator +from dbt.adapters.databricks.python_models.python_submissions import ( + PythonWorkflowConfigCompiler, + PythonWorkflowCreator, +) class TestPythonWorkflowConfigCompiler: diff --git a/tests/unit/relation_configs/test_comment.py b/tests/unit/relation_configs/test_comment.py index 8e2fc1504..71d802073 100644 --- a/tests/unit/relation_configs/test_comment.py +++ b/tests/unit/relation_configs/test_comment.py @@ -2,8 +2,7 @@ from agate import Table -from dbt.adapters.databricks.relation_configs.comment import CommentConfig -from dbt.adapters.databricks.relation_configs.comment import CommentProcessor +from dbt.adapters.databricks.relation_configs.comment import CommentConfig, CommentProcessor class TestCommentProcessor: diff --git a/tests/unit/relation_configs/test_materialized_view_config.py b/tests/unit/relation_configs/test_materialized_view_config.py index 80b6d2040..09aa95717 100644 --- a/tests/unit/relation_configs/test_materialized_view_config.py +++ b/tests/unit/relation_configs/test_materialized_view_config.py @@ -1,7 +1,6 @@ from unittest.mock import Mock -from agate import Row -from agate import Table +from agate import Row, Table from dbt.adapters.databricks.relation_configs.comment import CommentConfig from dbt.adapters.databricks.relation_configs.materialized_view import ( diff --git a/tests/unit/relation_configs/test_partitioning.py b/tests/unit/relation_configs/test_partitioning.py index 731b09c5d..fb345f066 100644 --- a/tests/unit/relation_configs/test_partitioning.py +++ b/tests/unit/relation_configs/test_partitioning.py @@ -1,7 +1,9 @@ from unittest.mock import Mock -from dbt.adapters.databricks.relation_configs.partitioning import PartitionedByConfig -from dbt.adapters.databricks.relation_configs.partitioning import PartitionedByProcessor +from dbt.adapters.databricks.relation_configs.partitioning import ( + PartitionedByConfig, + PartitionedByProcessor, +) from tests.unit import fixtures diff --git a/tests/unit/relation_configs/test_query.py b/tests/unit/relation_configs/test_query.py index 30748802e..3cd1f1b34 100644 --- a/tests/unit/relation_configs/test_query.py +++ b/tests/unit/relation_configs/test_query.py @@ -3,8 +3,7 @@ import pytest from agate import Row -from dbt.adapters.databricks.relation_configs.query import QueryConfig -from dbt.adapters.databricks.relation_configs.query import QueryProcessor +from dbt.adapters.databricks.relation_configs.query import QueryConfig, QueryProcessor from dbt.exceptions import DbtRuntimeError sql = "select * from foo" diff --git a/tests/unit/relation_configs/test_refresh.py b/tests/unit/relation_configs/test_refresh.py index 33efc22d2..9da56e4d8 100644 --- a/tests/unit/relation_configs/test_refresh.py +++ b/tests/unit/relation_configs/test_refresh.py @@ -2,8 +2,7 @@ import pytest -from dbt.adapters.databricks.relation_configs.refresh import RefreshConfig -from dbt.adapters.databricks.relation_configs.refresh import RefreshProcessor +from dbt.adapters.databricks.relation_configs.refresh import RefreshConfig, RefreshProcessor from dbt.exceptions import DbtRuntimeError from tests.unit import fixtures diff --git a/tests/unit/relation_configs/test_tags.py b/tests/unit/relation_configs/test_tags.py index d45a813e5..e893d3f59 100644 --- a/tests/unit/relation_configs/test_tags.py +++ b/tests/unit/relation_configs/test_tags.py @@ -3,8 +3,7 @@ import pytest from agate import Table -from dbt.adapters.databricks.relation_configs.tags import TagsConfig -from dbt.adapters.databricks.relation_configs.tags import TagsProcessor +from dbt.adapters.databricks.relation_configs.tags import TagsConfig, TagsProcessor from dbt.exceptions import DbtRuntimeError diff --git a/tests/unit/relation_configs/test_tblproperties.py b/tests/unit/relation_configs/test_tblproperties.py index 23f563735..dab860b7d 100644 --- a/tests/unit/relation_configs/test_tblproperties.py +++ b/tests/unit/relation_configs/test_tblproperties.py @@ -2,8 +2,10 @@ import pytest -from dbt.adapters.databricks.relation_configs.tblproperties import TblPropertiesConfig -from dbt.adapters.databricks.relation_configs.tblproperties import TblPropertiesProcessor +from dbt.adapters.databricks.relation_configs.tblproperties import ( + TblPropertiesConfig, + TblPropertiesProcessor, +) from dbt.exceptions import DbtRuntimeError from tests.unit import fixtures diff --git a/tests/unit/test_adapter.py b/tests/unit/test_adapter.py index 394256870..e707dcbb3 100644 --- a/tests/unit/test_adapter.py +++ b/tests/unit/test_adapter.py @@ -1,24 +1,21 @@ from multiprocessing import get_context -from typing import Any -from typing import Optional -from unittest.mock import Mock -from unittest.mock import patch +from typing import Any, Optional +from unittest.mock import Mock, patch import pytest from agate import Row -from dbt_common.exceptions import DbtConfigError -from dbt_common.exceptions import DbtValidationError +from dbt_common.exceptions import DbtConfigError, DbtValidationError import dbt.flags as flags -from dbt.adapters.databricks import __version__ -from dbt.adapters.databricks import DatabricksAdapter +from dbt.adapters.databricks import DatabricksAdapter, __version__ from dbt.adapters.databricks.column import DatabricksColumn -from dbt.adapters.databricks.credentials import CATALOG_KEY_IN_SESSION_PROPERTIES -from dbt.adapters.databricks.credentials import DBT_DATABRICKS_HTTP_SESSION_HEADERS -from dbt.adapters.databricks.credentials import DBT_DATABRICKS_INVOCATION_ENV +from dbt.adapters.databricks.credentials import ( + CATALOG_KEY_IN_SESSION_PROPERTIES, + DBT_DATABRICKS_HTTP_SESSION_HEADERS, + DBT_DATABRICKS_INVOCATION_ENV, +) from dbt.adapters.databricks.impl import get_identifier_list_string -from dbt.adapters.databricks.relation import DatabricksRelation -from dbt.adapters.databricks.relation import DatabricksRelationType +from dbt.adapters.databricks.relation import DatabricksRelation, DatabricksRelationType from dbt.adapters.databricks.utils import check_not_found_error from dbt.config import RuntimeConfig from tests.unit.utils import config_from_parts_or_dicts diff --git a/tests/unit/test_idle_config.py b/tests/unit/test_idle_config.py index 5e79f7df9..de3545680 100644 --- a/tests/unit/test_idle_config.py +++ b/tests/unit/test_idle_config.py @@ -3,8 +3,7 @@ from dbt.adapters.databricks import connections from dbt.adapters.databricks.credentials import DatabricksCredentials -from dbt.contracts.graph import model_config -from dbt.contracts.graph import nodes +from dbt.contracts.graph import model_config, nodes class TestDatabricksConnectionMaxIdleTime: diff --git a/tests/unit/test_relation.py b/tests/unit/test_relation.py index 02f939ea2..5bde2fa5e 100644 --- a/tests/unit/test_relation.py +++ b/tests/unit/test_relation.py @@ -1,8 +1,7 @@ import pytest from dbt.adapters.databricks import relation -from dbt.adapters.databricks.relation import DatabricksQuotePolicy -from dbt.adapters.databricks.relation import DatabricksRelation +from dbt.adapters.databricks.relation import DatabricksQuotePolicy, DatabricksRelation class TestDatabricksRelation: diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index cbe9ed3db..81e772113 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1,5 +1,4 @@ -from dbt.adapters.databricks.utils import redact_credentials -from dbt.adapters.databricks.utils import remove_ansi +from dbt.adapters.databricks.utils import redact_credentials, remove_ansi class TestDatabricksUtils: diff --git a/tests/unit/utils.py b/tests/unit/utils.py index d3e4fed25..6fdf86cdf 100644 --- a/tests/unit/utils.py +++ b/tests/unit/utils.py @@ -1,12 +1,9 @@ import os from copy import deepcopy -from dbt.config import Profile -from dbt.config import Project -from dbt.config import RuntimeConfig +from dbt.config import Profile, Project, RuntimeConfig from dbt.config.project import PartialProject -from dbt.config.renderer import DbtProjectYamlRenderer -from dbt.config.renderer import ProfileRenderer +from dbt.config.renderer import DbtProjectYamlRenderer, ProfileRenderer from dbt.config.utils import parse_cli_vars From 190364dd19ff18b6663ccca75dab38ec703b1741 Mon Sep 17 00:00:00 2001 From: dkruh1 <101973979+dkruh1@users.noreply.github.com> Date: Tue, 3 Dec 2024 16:01:55 +0200 Subject: [PATCH 8/9] test From 38b7e0dbe83a4e0e3ceb68afa21d06c6643c67d6 Mon Sep 17 00:00:00 2001 From: dkruh1 <101973979+dkruh1@users.noreply.github.com> Date: Tue, 3 Dec 2024 17:00:47 +0200 Subject: [PATCH 9/9] update changelog --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index dfce1d4e6..a746b439f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,6 @@ ## dbt-databricks 1.9.0 (TBD) ### Features - - Add config for generating unique tmp table names for enabling parralel merge (thanks @huangxingyi-git!) ([854](https://github.com/databricks/dbt-databricks/pull/854)) - Add support for serverless job clusters on python models ([706](https://github.com/databricks/dbt-databricks/pull/706)) - Add 'user_folder_for_python' behavior to switch writing python model notebooks to the user's folder ([835](https://github.com/databricks/dbt-databricks/pull/835)) @@ -18,6 +17,7 @@ - Add a new `workflow_job` submission method for python, which creates a long-lived Databricks Workflow instead of a one-time run (thanks @kdazzle!) ([762](https://github.com/databricks/dbt-databricks/pull/762)) - Allow for additional options to be passed to the Databricks Job API when using other python submission methods. For example, enable email_notifications (thanks @kdazzle!) ([762](https://github.com/databricks/dbt-databricks/pull/762)) - Support microbatch incremental strategy using replace_where ([825](https://github.com/databricks/dbt-databricks/pull/825)) +- Support pyspark session connection ([862]https://github.com/databricks/dbt-databricks/pull/862) ### Fixes