diff --git a/.github/workflows/integration-kernel.yml b/.github/workflows/integration-kernel.yml new file mode 100644 index 000000000..0d6fc9012 --- /dev/null +++ b/.github/workflows/integration-kernel.yml @@ -0,0 +1,197 @@ +# Kernel Integration Tests for dbt-databricks. +# +# Runs the functional suite against the SQL-warehouse profile with the +# databricks-sql-connector Rust kernel backend (use_kernel=True) instead of +# the default Thrift backend, on a weekly Saturday schedule. Like +# integration-spog.yml this is a scheduled smoke, not a PR gate — PR-targeting +# and status-reporting are omitted. +# +# Why warehouse-only: the kernel routes over the SEA (Statement Execution API) +# HTTP transport, which is warehouse-only — the connector rejects an +# all-purpose-cluster http_path ("SEA only works for warehouses"). So the two +# cluster profiles from the main matrix cannot run on the kernel and are not +# included here. +# +# Auth: the kernel supports PAT and Databricks OAuth (M2M/U2M) and has no +# Azure-AD flow. The peco CI credential is an Azure service principal, which the +# kernel cannot authenticate, so tests/profiles.py resolves it to a workspace +# token and hands the kernel a PAT when DBT_DATABRICKS_USE_KERNEL=1 (see +# _resolve_azure_sp_to_token_for_kernel). The adapter's production path forwards +# only raw PAT/Databricks-OAuth creds — see SqlUtils.prepare_connection_arguments. +# +# Concurrency & sharding: unlike the other lanes (-n 10, 3 shards), this one runs +# -n 4 across 5 shards. The kernel evicts SEA sessions under sustained -n 10 load +# (late verification queries fail with "session ... not open"), and the Azure test +# token has a ~1h lifetime — finer shards keep each job well under that. Both are +# kernel-lane workarounds; revisit as the kernel matures. See the run step below. +name: Kernel Integration Tests +on: + workflow_dispatch: + schedule: + - cron: "17 9 * * 6" # Weekly: Saturday 09:17 UTC (14:47 IST). Clear of the + # daily 21:30 UTC nightly and the Sunday SPOG slot, + # which share the peco warehouse. + +permissions: + id-token: write + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + prepare-shards: + runs-on: + group: databricks-protected-runner-group + labels: linux-ubuntu-latest + env: + UV_FROZEN: "1" + steps: + - name: Check out the repository + uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4 + + - name: Setup JFrog PyPI Proxy + uses: ./.github/actions/setup-jfrog-pypi + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5 + with: + python-version: "3.10" + + - name: Install uv + uses: astral-sh/setup-uv@38f3f104447c67c051c4a08e39b64a148898af3a # v4 + with: + cache-local-path: ~/.cache/uv + + - name: Install Hatch + uses: pypa/hatch@257e27e51a6a5616ed08a39a408a21c35c9931bc # install + + - name: Collect tests and assign to shards + run: | + set -euo pipefail + mkdir -p shard-assignments + hatch run pytest --collect-only -q --profile databricks_uc_sql_endpoint tests/functional 2>&1 \ + | grep "::" \ + > "shard-assignments/databricks_uc_sql_endpoint-collected.txt" + python3 scripts/shard_assign.py \ + --profile databricks_uc_sql_endpoint \ + --num-shards 5 \ + --input "shard-assignments/databricks_uc_sql_endpoint-collected.txt" \ + --output-dir shard-assignments \ + --algo lpt_historical_time \ + --timings .github/test_timings.json + + - name: Upload shard assignments + uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4 + with: + name: shard-assignments-kernel + path: shard-assignments/ + retention-days: 5 + + uc-sql-endpoint: + needs: prepare-shards + strategy: + fail-fast: false + matrix: + shard: [0, 1, 2, 3, 4] + runs-on: + group: databricks-protected-runner-group + labels: linux-ubuntu-latest + environment: azure-prod + env: + DBT_DATABRICKS_HOST_NAME: ${{ secrets.DATABRICKS_HOST }} + DBT_DATABRICKS_CLIENT_ID: ${{ secrets.TEST_PECO_SP_ID }} + DBT_DATABRICKS_CLIENT_SECRET: ${{ secrets.TEST_PECO_SP_SECRET }} + DBT_DATABRICKS_HTTP_PATH: ${{ secrets.TEST_PECO_WAREHOUSE_HTTP_PATH }} + DBT_DATABRICKS_UC_INITIAL_CATALOG: peco + DBT_DATABRICKS_LOCATION_ROOT: ${{ secrets.TEST_PECO_EXTERNAL_LOCATION }}test + # Wire the UC cluster id so build_cluster_http_path.py can populate + # DBT_DATABRICKS_UC_CLUSTER_HTTP_PATH. The kernel-incompatible tests that + # pin a model to that cluster path are @pytest.mark.skip_kernel and skip + # here; this keeps the env var resolvable so their fixtures still parse. + TEST_PECO_UC_CLUSTER_ID: ${{ secrets.TEST_PECO_UC_CLUSTER_ID }} + TEST_PECO_SPOG_HOST: ${{ secrets.TEST_PECO_SPOG_HOST }} + TEST_PECO_SPOG_WORKSPACE_ID: ${{ secrets.TEST_PECO_SPOG_WORKSPACE_ID }} + # Route the whole suite through the connector's Rust kernel (SEA) backend. + DBT_DATABRICKS_USE_KERNEL: "1" + UV_FROZEN: "1" + steps: + - name: Check out repository + uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4 + + - name: Setup JFrog PyPI Proxy + uses: ./.github/actions/setup-jfrog-pypi + + - name: Set up python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5 + with: + python-version: "3.10" + + - name: Get http path from environment + run: python .github/workflows/build_cluster_http_path.py + shell: sh + + - name: Install uv + uses: astral-sh/setup-uv@38f3f104447c67c051c4a08e39b64a148898af3a # v4 + with: + cache-local-path: ~/.cache/uv + + - name: Install Hatch + uses: pypa/hatch@257e27e51a6a5616ed08a39a408a21c35c9931bc # install + + - name: Install the connector kernel extra + # The kernel backend ships in the optional databricks-sql-connector + # [kernel] extra (a PyO3 wheel), which the frozen lock does not include. + # Add it into the hatch env so use_kernel=True can import the kernel. + run: | + set -euo pipefail + HATCH_PY="$(hatch run python -c 'import sys; print(sys.executable)')" + hatch run python -c "import databricks.sql as s; print('connector version:', s.__version__)" + uv pip install --python "$HATCH_PY" "databricks-sql-connector[kernel]" + hatch run python -c "import databricks_sql_kernel; print('kernel extra installed')" + + - name: Download shard assignments + uses: actions/download-artifact@65a9edc5881444af0b9093a5e628f2fe47ea3b2e # v4.1.7 + with: + name: shard-assignments-kernel + path: shard-assignments/ + + - name: Resolve test list for this shard + run: | + set -euo pipefail + SHARD_FILE="shard-assignments/databricks_uc_sql_endpoint-shard-${{ matrix.shard }}.txt" + if [ ! -s "$SHARD_FILE" ]; then + echo "::error::Shard file missing or empty: $SHARD_FILE" + exit 1 + fi + echo "SHARD_TESTS_FILE=$SHARD_FILE" >> "$GITHUB_ENV" + echo "Files in shard ${{ matrix.shard }}: $(wc -l < "$SHARD_FILE")" + + - name: Run Sql Endpoint Functional Tests on kernel (shard ${{ matrix.shard }}) + # Runs at -n 4 rather than the -n 10 the other lanes use: under sustained + # -n 10 the kernel's SEA sessions get evicted server-side, so a test's + # post-build verification query intermittently fails with "Cannot + # execute_command without an open session". -n 4 keeps the kernel session + # count low enough to avoid that. This lane is also sharded more finely + # (5 shards, above) so no shard runs long enough to cross an auth token's + # ~1h lifetime. Revisit once the kernel handles session lifecycle under + # higher concurrency. + run: | + mkdir -p logs + DBT_TEST_USER=notnecessaryformosttests@example.com \ + xargs -r hatch -v run pytest \ + --color=yes -v \ + --profile databricks_uc_sql_endpoint \ + -n 4 --dist=loadfile \ + --reruns 2 --reruns-delay 120 \ + --junitxml=logs/junit-kernel-sql-endpoint-shard-${{ matrix.shard }}.xml \ + < "$SHARD_TESTS_FILE" + + - name: Upload Kernel SQL Endpoint Test Logs + if: always() + uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4 + with: + name: kernel-sql-endpoint-logs-shard-${{ matrix.shard }} + path: logs/ + retention-days: 14 diff --git a/CHANGELOG.md b/CHANGELOG.md index 7121cf0e7..664318653 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ - Add catalogs.yml v2 support (requires `use_catalogs_v2: true` in dbt-core) ([1440](https://github.com/databricks/dbt-databricks/pull/1440)) - Add `skip_optimize` model config to opt out of the post-materialization `OPTIMIZE` call without dropping `zorder` / `liquid_clustered_by` / `auto_liquid_cluster` from the table definition. Useful when `OPTIMIZE` is delegated to Predictive Optimization or scheduled out of band. Complements the existing run-wide `DATABRICKS_SKIP_OPTIMIZE` var by allowing project-, folder-, or model-level opt-out via standard dbt config inheritance ([#703](https://github.com/databricks/dbt-databricks/issues/703)). +- Support the connector's Rust kernel backend via `connection_parameters: {use_kernel: true}` for SQL warehouses, with personal access token or Databricks OAuth (M2M/U2M) auth (requires `databricks-sql-connector[kernel]` on Python 3.10+; Azure service principals are not supported by the kernel) ([#1576](https://github.com/databricks/dbt-databricks/pull/1576)) ### Fixes - Stop dropping existing constraints on incremental runs when `contract.enforced` is `false` ([#1557](https://github.com/databricks/dbt-databricks/pull/1557)) @@ -22,6 +23,7 @@ ### Under the Hood +- Add a weekly `Kernel Integration Tests` workflow that runs the functional suite against the SQL-warehouse profile through the connector's Rust kernel backend (`DBT_DATABRICKS_USE_KERNEL=1`) (test-only, no runtime impact) ([#1576](https://github.com/databricks/dbt-databricks/pull/1576)). - Add functional tests for the `query` relation-config component's change handling: a streaming table's defining-query change is applied in place via `CREATE OR REFRESH`, and re-running a materialized view with an unchanged query leaves the existing relation in place instead of rebuilding it (test-only, no runtime impact). - Raise the `databricks-sql-connector` upper bound to `<4.3.1` to support `4.3.0` ([#1518](https://github.com/databricks/dbt-databricks/pull/1518)) - Add a functional test for incremental column-mask removal: dropping a `column_mask` from a model with an existing incremental relation issues `ALTER COLUMN ... DROP MASK` and leaves the column unmasked (test-only, no runtime impact). ([#1514](https://github.com/databricks/dbt-databricks/pull/1514)) diff --git a/dbt/adapters/databricks/handle.py b/dbt/adapters/databricks/handle.py index 218e4f1fa..3da799c9c 100644 --- a/dbt/adapters/databricks/handle.py +++ b/dbt/adapters/databricks/handle.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any, Optional, TypeVar from dbt.adapters.contracts.connection import AdapterResponse -from dbt_common.exceptions import DbtRuntimeError +from dbt_common.exceptions import DbtConfigError, DbtRuntimeError import databricks.sql as dbsql from databricks.sql.client import Connection, Cursor @@ -407,10 +407,9 @@ def prepare_connection_arguments( query_tags ) - return { + args: dict[str, Any] = { "server_hostname": creds.host, "http_path": http_path, - "credentials_provider": creds_manager.credentials_provider, "http_headers": http_headers if http_headers else None, "session_configuration": session_config, "catalog": creds.database, @@ -418,5 +417,63 @@ def prepare_connection_arguments( "schema": creds.schema, "_user_agent_entry": user_agent_entry, "user_agent_entry": user_agent_entry, - **connection_parameters, } + + if connection_parameters.get("use_kernel"): + SqlUtils._add_kernel_auth_arguments(args, creds, creds_manager) + else: + args["credentials_provider"] = creds_manager.credentials_provider + + return {**args, **connection_parameters} + + @staticmethod + def _add_kernel_auth_arguments( + args: dict[str, Any], + creds: DatabricksCredentials, + creds_manager: DatabricksCredentialManager, + ) -> None: + """Populate connection kwargs for the connector's Rust kernel backend. + + The kernel (use_kernel=True) routes over SEA and its auth bridge rejects the + connector credentials_provider dbt passes on the Thrift path. It supports + PAT, Databricks OAuth M2M, and OAuth U2M — we forward the raw credentials it + understands so the connector owns the token lifecycle and refresh. + + The kernel has no Azure-AD flow, so an Azure service principal cannot connect + through it; any auth other than the three above is rejected here with a clear + error directing the user to the default backend. + """ + # PAT: forward the token directly. + if creds_manager.token: + args["access_token"] = creds_manager.token + return + + # OAuth U2M (browser): auth_type "oauth" with no secret. Translate to the + # kernel's "databricks-oauth" (client_id is optional). An Azure SP keeps its + # secret in azure_client_secret, so this does not match one. + # Keep this auth_type check on raw creds to avoid initializing SDK auth here. + if ( + creds.auth_type == "oauth" + and not creds_manager.client_secret + and not creds_manager.azure_client_secret + ): + args["auth_type"] = "databricks-oauth" + args["oauth_client_id"] = creds_manager.client_id + args["oauth_scopes"] = creds_manager.oauth_scopes + return + + # Databricks OAuth M2M: forward OAuth creds so the kernel owns refresh. The + # resolved auth_type distinguishes genuine Databricks OAuth from an Azure SP + # supplied through the client_id/client_secret fields. + if not creds_manager.azure_client_secret and creds_manager.config.auth_type == "oauth-m2m": + args["oauth_client_id"] = creds_manager.client_id + args["oauth_client_secret"] = creds_manager.client_secret + args["oauth_scopes"] = creds_manager.oauth_scopes + return + + raise DbtConfigError( + "use_kernel=True supports only personal access tokens and Databricks " + "OAuth (M2M/U2M); the configured authentication (e.g. an Azure service " + "principal) is not supported by the kernel backend. Remove use_kernel to " + "use the default backend." + ) diff --git a/pyproject.toml b/pyproject.toml index 806108afa..bb325ca10 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -165,6 +165,7 @@ markers = [ "external: mark test as requiring an external location", "python: mark test as running a python model", "dlt: mark test as running a DLT model", + "skip_kernel: skip test when running against the connector kernel backend (DBT_DATABRICKS_USE_KERNEL=1)", ] [tool.mypy] diff --git a/tests/conftest.py b/tests/conftest.py index bf45a1493..dfcac90c2 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -13,13 +13,17 @@ def pytest_addoption(parser): parser.addoption("--profile", action="store", default=default_profile, type=str) -# Using @pytest.mark.skip_profile('databricks_cluster') uses the 'skip_by_adapter_type' -# autouse fixture below +# @pytest.mark.skip_profile(...) and @pytest.mark.skip_kernel are enforced by the +# autouse fixtures below (skip_by_profile_type, skip_on_kernel_backend). def pytest_configure(config): config.addinivalue_line( "markers", "skip_profile(profile): skip test for the given profile", ) + config.addinivalue_line( + "markers", + "skip_kernel: skip test when running against the connector kernel backend", + ) @pytest.fixture(scope="session") @@ -37,6 +41,18 @@ def skip_by_profile_type(request): pytest.skip(f"skipped on '{profile_type}' profile") +@pytest.fixture(autouse=True) +def skip_on_kernel_backend(request): + # The connector kernel backend (DBT_DATABRICKS_USE_KERNEL=1) routes over SEA, + # which is warehouse-only. Tests that pin a model to an all-purpose-cluster + # http_path or compute cannot connect through it, so they opt out with + # @pytest.mark.skip_kernel. + if os.environ.get("DBT_DATABRICKS_USE_KERNEL") == "1" and request.node.get_closest_marker( + "skip_kernel" + ): + pytest.skip("skipped on the connector kernel backend (SEA is warehouse-only)") + + # The profile dictionary, used to write out profiles.yml. It will pull in updates # from two separate sources, the 'profile_target' and 'profiles_config_update'. # The second one is useful when using alternative targets, etc. diff --git a/tests/functional/adapter/incremental/test_incremental_column_tags.py b/tests/functional/adapter/incremental/test_incremental_column_tags.py index cc2e5e44e..a484b08be 100644 --- a/tests/functional/adapter/incremental/test_incremental_column_tags.py +++ b/tests/functional/adapter/incremental/test_incremental_column_tags.py @@ -7,6 +7,7 @@ from tests.functional.adapter.incremental import fixtures +@pytest.mark.skip_kernel # pins model http_path to a UC cluster; SEA/kernel is warehouse-only @pytest.mark.skip_profile("databricks_cluster") class TestIncrementalColumnTags(RerunSafeMixin, MaterializationV2Mixin): @pytest.fixture(scope="class") @@ -57,6 +58,7 @@ def test_changing_column_tags(self, project): @pytest.mark.python +@pytest.mark.skip_kernel # Python model runs on a cluster; kernel/SEA is warehouse-only @pytest.mark.skip_profile("databricks_cluster") class TestIncrementalPythonColumnTags(TestIncrementalColumnTags): @pytest.fixture(scope="class") diff --git a/tests/functional/adapter/incremental/test_incremental_strategies.py b/tests/functional/adapter/incremental/test_incremental_strategies.py index 8c19219c1..238bd3166 100644 --- a/tests/functional/adapter/incremental/test_incremental_strategies.py +++ b/tests/functional/adapter/incremental/test_incremental_strategies.py @@ -166,6 +166,7 @@ def test_incremental(self, project): # Only runs under SQL warehouse profile, but overrides compute at model level +@pytest.mark.skip_kernel # overrides compute to a UC cluster; SEA/kernel is warehouse-only @pytest.mark.skip_profile("databricks_uc_cluster", "databricks_cluster") class TestInsertOverwriteWithModelComputeOverride(IncrementalBase): @pytest.fixture(scope="class") diff --git a/tests/functional/adapter/python_model/test_python_model.py b/tests/functional/adapter/python_model/test_python_model.py index b9f7a2f68..5ab868e1c 100644 --- a/tests/functional/adapter/python_model/test_python_model.py +++ b/tests/functional/adapter/python_model/test_python_model.py @@ -182,6 +182,7 @@ def test_changing_schema_via_incremental(self, project): @pytest.mark.python +@pytest.mark.skip_kernel # pins model http_path to a cluster; SEA/kernel is warehouse-only @pytest.mark.skip_profile("databricks_cluster", "databricks_uc_cluster") @pytest.mark.flaky(reruns=2, reruns_delay=120) class TestSpecifyingHttpPath(PythonModelDataMixin, BasePythonModelTests): diff --git a/tests/functional/adapter/spog/test_spog_connection_open.py b/tests/functional/adapter/spog/test_spog_connection_open.py index ba538cae3..a79833236 100644 --- a/tests/functional/adapter/spog/test_spog_connection_open.py +++ b/tests/functional/adapter/spog/test_spog_connection_open.py @@ -32,6 +32,7 @@ def test_probe_failure_run_succeeds(self, project): run_dbt(["run", "--select", "spog_smoke_model"], expect_pass=True) +@pytest.mark.skip_kernel # pins model to a UC cluster; SEA/kernel is warehouse-only @pytest.mark.skip_profile("databricks_cluster", "databricks_uc_cluster") class TestSpogNamedCompute: @pytest.fixture(scope="class") diff --git a/tests/profiles.py b/tests/profiles.py index 0909cd539..4ceb5f81f 100644 --- a/tests/profiles.py +++ b/tests/profiles.py @@ -50,10 +50,51 @@ def _build_databricks_cluster_target( # If you are specifying a port for running tests, assume Docker # is being used and disable TLS verification connection_parameters["_tls_no_verify"] = True + # Route the whole suite through the connector's Rust kernel backend (SEA + # transport) when DBT_DATABRICKS_USE_KERNEL=1. Used by the weekend + # integration-kernel workflow. SEA is warehouse-only, so this is only + # meaningful on the databricks_uc_sql_endpoint profile. + if os.getenv("DBT_DATABRICKS_USE_KERNEL") == "1": + connection_parameters["use_kernel"] = True + # TEMPORARY: the CI credential is an Azure service principal, which the + # kernel cannot authenticate (no Azure-AD flow). Resolve it to a workspace + # token here so the suite can run. The token does not auto-refresh, so this + # is limited to short-lived test connections; remove once the kernel adds + # native Azure support. + if profile.get("azure_client_secret") or profile.get("client_secret"): + _resolve_azure_sp_to_token_for_kernel(profile) profile["connection_parameters"] = connection_parameters return profile +def _resolve_azure_sp_to_token_for_kernel(profile: dict[str, Any]) -> None: + """Rewrite the Azure-service-principal CI profile into a PAT for the kernel + backend: authenticate the SP through the credential manager, set the resolved + workspace token as the profile's token, and drop the SP fields. + """ + from dbt.adapters.databricks.credentials import ( + DatabricksCredentialManager, + DatabricksCredentials, + ) + + creds = DatabricksCredentials( + host=profile["host"], + http_path=profile["http_path"], + client_id=profile.get("client_id"), + client_secret=profile.get("client_secret"), + azure_client_id=profile.get("azure_client_id"), + azure_client_secret=profile.get("azure_client_secret"), + auth_type=profile.get("auth_type", "oauth"), + schema=profile.get("schema", "default"), + ) + manager = DatabricksCredentialManager.create_from(creds) + authorization = manager.header_factory().get("Authorization", "") + profile["token"] = authorization[len("Bearer ") :] + profile["auth_type"] = "pat" + for sp_field in ("client_id", "client_secret", "azure_client_id", "azure_client_secret"): + profile.pop(sp_field, None) + + def databricks_cluster_target(): return _build_databricks_cluster_target( http_path=os.getenv( diff --git a/tests/unit/test_handle.py b/tests/unit/test_handle.py index a0c36ae2b..8f58c25a6 100644 --- a/tests/unit/test_handle.py +++ b/tests/unit/test_handle.py @@ -3,12 +3,17 @@ import os import sys from decimal import Decimal +from typing import Any from unittest.mock import Mock, patch import pytest from databricks.sql.client import Cursor -from dbt_common.exceptions import DbtRuntimeError +from dbt_common.exceptions import DbtConfigError, DbtRuntimeError +from dbt.adapters.databricks.credentials import ( + DatabricksCredentialManager, + DatabricksCredentials, +) from dbt.adapters.databricks.handle import ( CursorWrapper, DatabricksAdapterResponse, @@ -54,6 +59,42 @@ def _build_session_headers( return json.dumps(payload) +_KERNEL_HTTP_PATH = "/sql/1.0/warehouses/abc" + + +def _prepare_connection_args(**creds_kwargs: Any) -> dict[str, Any]: + """Run prepare_connection_arguments for creds built from fixed host/path/schema + plus the given auth kwargs, and return the resulting connect() kwargs.""" + creds = DatabricksCredentials( + host="yourorg.databricks.com", + http_path=_KERNEL_HTTP_PATH, + schema="dbt", + **creds_kwargs, + ) + manager = DatabricksCredentialManager.create_from(creds) + return SqlUtils.prepare_connection_arguments(creds, manager, _KERNEL_HTTP_PATH, {}) + + +def _fail_if_config_accessed(self: DatabricksCredentialManager) -> Any: + raise AssertionError("creds_manager.config should not be accessed") + + +def _prepare_connection_args_without_config_auth(**creds_kwargs: Any) -> dict[str, Any]: + creds = DatabricksCredentials( + host="yourorg.databricks.com", + http_path=_KERNEL_HTTP_PATH, + schema="dbt", + **creds_kwargs, + ) + manager = DatabricksCredentialManager.create_from(creds) + with patch.object( + type(manager), + "config", + new_callable=lambda: property(_fail_if_config_accessed), + ): + return SqlUtils.prepare_connection_arguments(creds, manager, _KERNEL_HTTP_PATH, {}) + + class TestSqlUtils: @pytest.mark.parametrize( "bindings, expected", [(None, None), ([1], [1]), ([1, Decimal(0.73)], [1, 0.73])] @@ -190,6 +231,129 @@ def test_extract_dbr_version__invalid(self): with pytest.raises(DbtRuntimeError): SqlUtils.extract_dbr_version("foo") + def test_prepare_connection_arguments__default_uses_credentials_provider(self): + """Without use_kernel, the connector's credentials_provider is passed and + no raw credentials leak into the connect() kwargs.""" + args = _prepare_connection_args(client_id="cid", client_secret="dose-secret") + assert callable(args["credentials_provider"]) + assert "use_kernel" not in args + assert "oauth_client_id" not in args + assert "access_token" not in args + + def test_prepare_connection_arguments__kernel_pat_forwards_access_token(self): + """use_kernel with a PAT forwards the token directly as access_token + (PATs are long-lived, so no refresh concern), dropping credentials_provider.""" + args = _prepare_connection_args_without_config_auth( + token="dapiabc123", + connection_parameters={"use_kernel": True}, + ) + assert args["use_kernel"] is True + assert args["access_token"] == "dapiabc123" + assert args.get("credentials_provider") is None + assert "oauth_client_id" not in args + + def test_prepare_connection_arguments__kernel_m2m_forwards_oauth_creds(self): + """OAuth M2M forwards resolved client id/scopes plus the client secret so the + kernel owns the token lifecycle (acquire + refresh via workspace OIDC), not a + pre-resolved token that would not refresh.""" + creds = DatabricksCredentials( + host="yourorg.databricks.com", + http_path=_KERNEL_HTTP_PATH, + schema="dbt", + client_id="cid", + client_secret="dose-secret", + connection_parameters={"use_kernel": True}, + ) + manager = DatabricksCredentialManager.create_from(creds) + with patch.object( + type(manager), + "config", + new_callable=lambda: property(lambda self: Mock(auth_type="oauth-m2m")), + ): + args = SqlUtils.prepare_connection_arguments(creds, manager, _KERNEL_HTTP_PATH, {}) + assert args["oauth_client_id"] == "cid" + assert args["oauth_client_secret"] == "dose-secret" + # Scopes come from the manager's resolved defaults, not raw creds. + assert args["oauth_scopes"] == ["all-apis", "offline_access"] + assert args.get("credentials_provider") is None + assert "access_token" not in args + + def test_prepare_connection_arguments__kernel_azure_sp_raises(self): + """The kernel has no Azure-AD flow, so an Azure SP cannot connect through it. + The adapter rejects it with a clear error rather than silently converting the + credentials (that conversion, needed only to run the CI suite on a peco Azure + SP, lives in test setup — see tests/profiles.py).""" + creds = DatabricksCredentials( + host="yourorg.databricks.com", + http_path=_KERNEL_HTTP_PATH, + schema="dbt", + azure_client_id="acid", + azure_client_secret="asecret", + connection_parameters={"use_kernel": True}, + ) + manager = DatabricksCredentialManager.create_from(creds) + with patch.object( + type(manager), + "config", + new_callable=lambda: property(_fail_if_config_accessed), + ): + with pytest.raises(DbtConfigError, match="use_kernel"): + SqlUtils.prepare_connection_arguments(creds, manager, _KERNEL_HTTP_PATH, {}) + + def test_prepare_connection_arguments__kernel_legacy_azure_sp_raises(self): + """An Azure SP supplied through the client_id/client_secret fields resolves + to azure-client-secret; the kernel cannot serve it, so it is rejected too.""" + creds = DatabricksCredentials( + host="yourorg.databricks.com", + http_path=_KERNEL_HTTP_PATH, + schema="dbt", + client_id="cid", + client_secret="azure-secret", + connection_parameters={"use_kernel": True}, + ) + manager = DatabricksCredentialManager.create_from(creds) + with patch.object( + type(manager), + "config", + new_callable=lambda: property(lambda self: Mock(auth_type="azure-client-secret")), + ): + with pytest.raises(DbtConfigError, match="use_kernel"): + SqlUtils.prepare_connection_arguments(creds, manager, _KERNEL_HTTP_PATH, {}) + + def test_prepare_connection_arguments__kernel_u2m_explicit_client_id(self): + """use_kernel with OAuth U2M and an explicit client_id forwards that + client_id and translates dbt's auth_type='oauth' to the kernel's + 'databricks-oauth' so the kernel bridge routes to its U2M flow. U2M is a + browser flow the kernel runs itself, so no token is pre-resolved.""" + args = _prepare_connection_args_without_config_auth( + client_id="cid", + auth_type="oauth", + connection_parameters={"use_kernel": True}, + ) + assert args["use_kernel"] is True + assert args["oauth_client_id"] == "cid" + assert args["oauth_scopes"] == ["all-apis", "offline_access"] + assert args["auth_type"] == "databricks-oauth" + assert args.get("credentials_provider") is None + assert "oauth_client_secret" not in args + assert "access_token" not in args + + def test_prepare_connection_arguments__kernel_u2m_default_app_uses_manager_defaults(self): + """The common external-browser config is `auth_type: oauth` with NO + client_id/scopes set. We forward the credential manager's resolved defaults + (client_id=dbt-databricks, scopes=[all-apis, offline_access]) so the kernel + uses dbt's defaults rather than its own.""" + args = _prepare_connection_args_without_config_auth( + auth_type="oauth", + connection_parameters={"use_kernel": True}, + ) + assert args["auth_type"] == "databricks-oauth" + assert args["oauth_client_id"] == "dbt-databricks" + assert args["oauth_scopes"] == ["all-apis", "offline_access"] + assert args.get("credentials_provider") is None + assert "oauth_client_secret" not in args + assert "access_token" not in args + class TestCursorWrapper: @pytest.fixture