From dce9bc98052deb76589db059f6a1f36228bf83fd Mon Sep 17 00:00:00 2001 From: Shubham Dhal Date: Mon, 6 Jul 2026 18:34:15 +0530 Subject: [PATCH 1/8] Add kernel integration test lane + use_kernel auth support Adds a weekly SQL-warehouse integration workflow that runs the functional suite through the databricks-sql-connector Rust kernel backend (use_kernel=True), and the adapter change needed for connections to authenticate on that path. The kernel routes over SEA and owns the OAuth token lifecycle itself, so its auth bridge rejects the connector credentials_provider dbt passes on the Thrift path. prepare_connection_arguments now forwards the raw credential the bridge understands (PAT, Databricks OAuth M2M, or OAuth U2M) when use_kernel is set, and raises a clear error for native Azure service-principal auth (which the kernel cannot serve) rather than misrouting it. SEA is warehouse-only, so the workflow runs only the databricks_uc_sql_endpoint profile and tests that pin a model to a cluster are marked @pytest.mark.skip_kernel. --- .github/workflows/integration-kernel.yml | 183 ++++++++++++++++++ CHANGELOG.md | 2 + dbt/adapters/databricks/handle.py | 66 ++++++- pyproject.toml | 1 + tests/conftest.py | 20 +- .../test_incremental_column_tags.py | 2 + .../test_incremental_strategies.py | 1 + .../adapter/python_model/test_python_model.py | 1 + .../adapter/spog/test_spog_connection_open.py | 1 + tests/profiles.py | 6 + tests/unit/test_handle.py | 125 +++++++++++- 11 files changed, 401 insertions(+), 7 deletions(-) create mode 100644 .github/workflows/integration-kernel.yml diff --git a/.github/workflows/integration-kernel.yml b/.github/workflows/integration-kernel.yml new file mode 100644 index 000000000..7dd5ae5d4 --- /dev/null +++ b/.github/workflows/integration-kernel.yml @@ -0,0 +1,183 @@ +# 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 owns the token lifecycle itself and its auth bridge rejects +# the connector credentials_provider dbt passes on the Thrift path. When +# DBT_DATABRICKS_USE_KERNEL=1, tests/profiles.py sets use_kernel=True and the +# adapter resolves the configured credentials (PAT, OAuth M2M, or the peco Azure +# service principal) to a Databricks workspace bearer token and forwards it to +# the kernel's PAT path — see SqlUtils.prepare_connection_arguments. +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 3 \ + --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] + 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 }}) + 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 bb8f00a62..93576e5aa 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 (requires `databricks-sql-connector[kernel]` on Python 3.10+). The kernel's auth bridge rejects the connector `credentials_provider` dbt uses on the default path, so on this path the adapter resolves the configured credentials — personal access token, OAuth (M2M/U2M), or Azure service principal — to a bearer token and forwards that to the kernel; the forwarded token is point-in-time rather than auto-refreshing ([#1576](https://github.com/databricks/dbt-databricks/pull/1576)) ### Fixes - Honor the `expression` field on `primary_key` constraints on the V1 materialization path. A primary key declared with `expression: RELY` (or any trailing clause) previously had its expression silently dropped. ([#1551](https://github.com/databricks/dbt-databricks/pull/1551)) @@ -20,6 +21,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..dacee74aa 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,64 @@ 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 owns the token lifecycle + itself, so its auth bridge rejects the connector credentials_provider that + dbt passes on the Thrift path, and it accepts only PAT, Databricks OAuth + M2M, or OAuth U2M — it has no Azure-AD service-principal flow. + + Rather than reclassify the credential shape and forward raw creds (which + misroutes an Azure SP into the kernel's Databricks-OAuth M2M flow and fails + with invalid_client), we let the credential manager perform whatever auth + dbt already supports — PAT, OAuth M2M, or Azure service principal — and + forward the resolved Databricks workspace bearer token to the kernel's PAT + path. One path serves every machine-auth method the adapter supports. + + U2M (browser) is the exception: there is no token to pre-resolve because + the kernel runs the interactive flow itself, so we translate dbt's + auth_type to the kernel's "databricks-oauth" and let it drive the browser. + + Note: the forwarded token is a point-in-time value, not the refreshing + credentials_provider the Thrift path uses. That is fine for dbt's + short-lived per-thread connections; a connection outliving the token's + lifetime would need a reconnect. + """ + if not creds.token and creds.auth_type == "oauth" and not creds.client_secret: + # OAuth U2M (browser). Translate dbt's auth_type "oauth" to the kernel's + # "databricks-oauth"; client_id is optional when using the default dbt app. + args["auth_type"] = "databricks-oauth" + if creds.client_id: + args["oauth_client_id"] = creds.client_id + if creds.oauth_scopes: + args["oauth_scopes"] = creds.oauth_scopes + return + + # Machine auth (PAT / OAuth M2M / Azure SP): resolve to a bearer token. + headers: dict[str, str] = {} + headers.update(creds_manager.header_factory()) + authorization = headers.get("Authorization", "") + prefix = "Bearer " + if not authorization[: len(prefix)].lower() == prefix.lower(): + raise DbtConfigError( + "use_kernel=True could not obtain a bearer token from the configured " + "credentials. The kernel backend supports personal access tokens, " + "Databricks OAuth (M2M/U2M), and Azure service principals; other auth " + "methods must use the default backend (drop use_kernel)." + ) + args["access_token"] = authorization[len(prefix) :] 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..c04894e7c 100644 --- a/tests/profiles.py +++ b/tests/profiles.py @@ -50,6 +50,12 @@ 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 profile["connection_parameters"] = connection_parameters return profile diff --git a/tests/unit/test_handle.py b/tests/unit/test_handle.py index a0c36ae2b..591a1b796 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,22 @@ 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, {}) + + class TestSqlUtils: @pytest.mark.parametrize( "bindings, expected", [(None, None), ([1], [1]), ([1, Decimal(0.73)], [1, 0.73])] @@ -190,6 +211,108 @@ 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 _kernel_args_with_token(self, bearer, **creds_kwargs): + """Build kernel connect() kwargs with the credential manager's + header_factory stubbed to yield the given Authorization header, so the + machine-auth path can be exercised without live authentication.""" + creds = DatabricksCredentials( + host="yourorg.databricks.com", + http_path=_KERNEL_HTTP_PATH, + schema="dbt", + connection_parameters={"use_kernel": True}, + **creds_kwargs, + ) + manager = DatabricksCredentialManager.create_from(creds) + header = {"Authorization": bearer} if bearer is not None else {} + with patch.object( + type(manager), + "header_factory", + new_callable=lambda: property(lambda self: (lambda: header)), + ): + return SqlUtils.prepare_connection_arguments(creds, manager, _KERNEL_HTTP_PATH, {}) + + def test_prepare_connection_arguments__kernel_pat_forwards_resolved_token(self): + """use_kernel with a PAT resolves the bearer token from the credential + manager and forwards it as access_token, dropping credentials_provider.""" + args = self._kernel_args_with_token("Bearer dapiabc123", token="dapiabc123") + 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_resolved_token(self): + """OAuth M2M is resolved to a Databricks workspace token by the credential + manager, and that token is forwarded to the kernel's PAT path — not the raw + client id/secret.""" + args = self._kernel_args_with_token( + "Bearer exchanged-m2m-token", client_id="cid", client_secret="dose-secret" + ) + assert args["access_token"] == "exchanged-m2m-token" + assert args.get("credentials_provider") is None + assert "oauth_client_secret" not in args + + def test_prepare_connection_arguments__kernel_azure_sp_forwards_resolved_token(self): + """The Azure-AD service principal (the peco CI credential) is resolved to a + Databricks workspace token by the credential manager's azure-client-secret + auth, and that token drives the kernel's PAT path. This is what makes the + kernel work for Azure SP — forwarding the raw creds to the kernel's + Databricks-OAuth M2M would fail with invalid_client.""" + args = self._kernel_args_with_token( + "Bearer azure-exchanged-token", + azure_client_id="acid", + azure_client_secret="asecret", + ) + assert args["access_token"] == "azure-exchanged-token" + assert args.get("credentials_provider") is None + assert "oauth_client_id" not in args + + def test_prepare_connection_arguments__kernel_non_bearer_header_raises(self): + """If the credential manager yields a non-Bearer (or missing) Authorization + header, there is no token to forward — fail loudly rather than hand the + kernel unusable auth.""" + with pytest.raises(DbtConfigError): + self._kernel_args_with_token("Basic abc123", token="dapiabc123") + + def test_prepare_connection_arguments__kernel_u2m_explicit_client_id(self): + """use_kernel with OAuth U2M and an explicit client_id forwards + oauth_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( + 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["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_no_client_id(self): + """The common external-browser config is `auth_type: oauth` with NO + client_id (SDK uses the default dbt app). The kernel bridge routes U2M off + auth_type alone, so we must translate auth_type even without a client_id.""" + args = _prepare_connection_args( + auth_type="oauth", + connection_parameters={"use_kernel": True}, + ) + assert args["auth_type"] == "databricks-oauth" + assert args.get("credentials_provider") is None + assert "oauth_client_id" not in args + assert "oauth_client_secret" not in args + assert "access_token" not in args + class TestCursorWrapper: @pytest.fixture From ffa6ae3a8ec114e03c17ed1eb7e33bba501a7061 Mon Sep 17 00:00:00 2001 From: Shubham Dhal Date: Tue, 7 Jul 2026 10:45:47 +0530 Subject: [PATCH 2/8] TEMP: add pull_request trigger to verify kernel workflow on PR (revert before merge) --- .github/workflows/integration-kernel.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/integration-kernel.yml b/.github/workflows/integration-kernel.yml index 7dd5ae5d4..514e6c2f4 100644 --- a/.github/workflows/integration-kernel.yml +++ b/.github/workflows/integration-kernel.yml @@ -21,6 +21,7 @@ name: Kernel Integration Tests on: workflow_dispatch: + pull_request: # TEMP: verify the workflow runs on PR #1576; reverted before merge. 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, From ac09009e35f167037048fd81435f81e79b987eec Mon Sep 17 00:00:00 2001 From: Shubham Dhal Date: Tue, 7 Jul 2026 12:50:55 +0530 Subject: [PATCH 3/8] Prefer raw-cred forwarding on kernel path; bearer only as Azure fallback The kernel's oauth-m2m flow owns the full token lifecycle including refresh (re-mints via workspace OIDC), so forwarding raw client_id/client_secret keeps long-running connections authenticated. Reverts to raw-cred forwarding as the primary path for PAT and Databricks OAuth (M2M/U2M). Azure service principals have no kernel-native flow, so they remain a fallback: dbt resolves the Azure exchange and forwards the resulting bearer token to the kernel's PAT path, with a warning that this token does not auto-refresh. Routing uses dbt's resolved config.auth_type, since an Azure SP may arrive either via azure_client_id/secret or (legacily) the client_id/client_secret fields. Fixes the CI invalid_client failure (peco is an Azure SP) without giving up token refresh for the common PAT/Databricks-OAuth cases. --- .github/workflows/integration-kernel.yml | 7 +- CHANGELOG.md | 2 +- dbt/adapters/databricks/handle.py | 98 ++++++++++++++++-------- tests/unit/test_handle.py | 88 +++++++++++++-------- 4 files changed, 127 insertions(+), 68 deletions(-) diff --git a/.github/workflows/integration-kernel.yml b/.github/workflows/integration-kernel.yml index 514e6c2f4..dc5b208fb 100644 --- a/.github/workflows/integration-kernel.yml +++ b/.github/workflows/integration-kernel.yml @@ -15,9 +15,10 @@ # Auth: the kernel owns the token lifecycle itself and its auth bridge rejects # the connector credentials_provider dbt passes on the Thrift path. When # DBT_DATABRICKS_USE_KERNEL=1, tests/profiles.py sets use_kernel=True and the -# adapter resolves the configured credentials (PAT, OAuth M2M, or the peco Azure -# service principal) to a Databricks workspace bearer token and forwards it to -# the kernel's PAT path — see SqlUtils.prepare_connection_arguments. +# adapter forwards the raw credentials the kernel understands (PAT or Databricks +# OAuth) so the kernel owns refresh. The peco CI credential is an Azure service +# principal, which has no kernel-native flow, so it falls back to a pre-resolved +# bearer token (no auto-refresh) — see SqlUtils.prepare_connection_arguments. name: Kernel Integration Tests on: workflow_dispatch: diff --git a/CHANGELOG.md b/CHANGELOG.md index 93576e5aa..883db0889 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +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 (requires `databricks-sql-connector[kernel]` on Python 3.10+). The kernel's auth bridge rejects the connector `credentials_provider` dbt uses on the default path, so on this path the adapter resolves the configured credentials — personal access token, OAuth (M2M/U2M), or Azure service principal — to a bearer token and forwards that to the kernel; the forwarded token is point-in-time rather than auto-refreshing ([#1576](https://github.com/databricks/dbt-databricks/pull/1576)) +- Support the connector's Rust kernel backend via `connection_parameters: {use_kernel: true}` for SQL warehouses (requires `databricks-sql-connector[kernel]` on Python 3.10+). The kernel's auth bridge rejects the connector `credentials_provider` dbt uses on the default path, so the adapter forwards the raw credentials the kernel understands — personal access token or Databricks OAuth (M2M/U2M) — letting the kernel own the token lifecycle, including OAuth refresh. Azure service principals have no kernel-native flow, so they fall back to a pre-resolved bearer token (logged with a warning) that does not auto-refresh; use a PAT or Databricks OAuth, or the default backend, for long-lived Azure connections ([#1576](https://github.com/databricks/dbt-databricks/pull/1576)) ### Fixes - Honor the `expression` field on `primary_key` constraints on the V1 materialization path. A primary key declared with `expression: RELY` (or any trailing clause) previously had its expression silently dropped. ([#1551](https://github.com/databricks/dbt-databricks/pull/1551)) diff --git a/dbt/adapters/databricks/handle.py b/dbt/adapters/databricks/handle.py index dacee74aa..2a66e9530 100644 --- a/dbt/adapters/databricks/handle.py +++ b/dbt/adapters/databricks/handle.py @@ -436,28 +436,38 @@ def _add_kernel_auth_arguments( The kernel (use_kernel=True) routes over SEA and owns the token lifecycle itself, so its auth bridge rejects the connector credentials_provider that - dbt passes on the Thrift path, and it accepts only PAT, Databricks OAuth - M2M, or OAuth U2M — it has no Azure-AD service-principal flow. - - Rather than reclassify the credential shape and forward raw creds (which - misroutes an Azure SP into the kernel's Databricks-OAuth M2M flow and fails - with invalid_client), we let the credential manager perform whatever auth - dbt already supports — PAT, OAuth M2M, or Azure service principal — and - forward the resolved Databricks workspace bearer token to the kernel's PAT - path. One path serves every machine-auth method the adapter supports. - - U2M (browser) is the exception: there is no token to pre-resolve because - the kernel runs the interactive flow itself, so we translate dbt's - auth_type to the kernel's "databricks-oauth" and let it drive the browser. - - Note: the forwarded token is a point-in-time value, not the refreshing - credentials_provider the Thrift path uses. That is fine for dbt's - short-lived per-thread connections; a connection outliving the token's - lifetime would need a reconnect. + dbt passes on the Thrift path. It accepts only three flows — PAT, Databricks + OAuth M2M, and OAuth U2M — and has no Azure-AD service-principal flow. + + We forward the raw credentials the kernel understands so that IT owns the + token lifecycle, including refresh: for OAuth M2M the kernel re-mints tokens + via workspace OIDC, so a long-running connection stays authenticated. This + is the important property for real dbt runs, which are not bounded to a + token's lifetime. + + The Azure service principal is the exception: the kernel cannot mint Azure + tokens, so there is no raw-cred path for it. As a fallback we let dbt's + credential manager perform the Azure exchange and forward the resulting + Databricks workspace bearer token to the kernel's PAT path. That token is a + point-in-time value the kernel does NOT refresh, so a connection outliving + it (~1h) would need to be reopened — we log a warning to make that explicit. + Routing uses dbt's resolved ``config.auth_type`` because an Azure SP may be + supplied either natively (azure_client_id/secret) or, legacily, through the + M2M client_id/client_secret fields — indistinguishable by input shape alone. """ - if not creds.token and creds.auth_type == "oauth" and not creds.client_secret: - # OAuth U2M (browser). Translate dbt's auth_type "oauth" to the kernel's - # "databricks-oauth"; client_id is optional when using the default dbt app. + # PAT: forward the token directly (long-lived; no refresh concern). + if creds.token: + args["access_token"] = creds.token + return + + # OAuth U2M (browser): the kernel runs the interactive flow itself, so there + # is no token to pre-resolve. Translate dbt's auth_type "oauth" to the + # kernel's "databricks-oauth"; client_id is optional (default dbt app). + if ( + creds.auth_type == "oauth" + and not creds.client_secret + and not (creds.azure_client_id and creds.azure_client_secret) + ): args["auth_type"] = "databricks-oauth" if creds.client_id: args["oauth_client_id"] = creds.client_id @@ -465,16 +475,38 @@ def _add_kernel_auth_arguments( args["oauth_scopes"] = creds.oauth_scopes return - # Machine auth (PAT / OAuth M2M / Azure SP): resolve to a bearer token. - headers: dict[str, str] = {} - headers.update(creds_manager.header_factory()) - authorization = headers.get("Authorization", "") - prefix = "Bearer " - if not authorization[: len(prefix)].lower() == prefix.lower(): - raise DbtConfigError( - "use_kernel=True could not obtain a bearer token from the configured " - "credentials. The kernel backend supports personal access tokens, " - "Databricks OAuth (M2M/U2M), and Azure service principals; other auth " - "methods must use the default backend (drop use_kernel)." + # Resolve what dbt actually authenticates as. For OAuth M2M we forward raw + # creds so the kernel owns refresh; for Azure (native or legacy-via-M2M- + # fields) the kernel has no flow, so we fall back to a resolved bearer token. + resolved_auth_type = creds_manager.config.auth_type + + if resolved_auth_type == "oauth-m2m": + args["oauth_client_id"] = creds.client_id + args["oauth_client_secret"] = creds.client_secret + if creds.oauth_scopes: + args["oauth_scopes"] = creds.oauth_scopes + return + + if resolved_auth_type == "azure-client-secret": + logger.warning( + "use_kernel=True with an Azure service principal: the kernel has no " + "Azure-AD auth flow, so dbt forwards a pre-resolved bearer token that " + "the kernel will NOT refresh. A connection open for longer than the " + "token's lifetime (~1h) will fail to re-authenticate; use a personal " + "access token or Databricks OAuth, or the default backend, for " + "long-lived connections." ) - args["access_token"] = authorization[len(prefix) :] + headers: dict[str, str] = {} + headers.update(creds_manager.header_factory()) + authorization = headers.get("Authorization", "") + prefix = "Bearer " + if authorization[: len(prefix)].lower() == prefix.lower(): + args["access_token"] = authorization[len(prefix) :] + return + + raise DbtConfigError( + "use_kernel=True could not resolve usable credentials. The kernel backend " + "supports personal access tokens, Databricks OAuth (M2M/U2M), and Azure " + "service principals; other auth methods must use the default backend " + "(drop use_kernel)." + ) diff --git a/tests/unit/test_handle.py b/tests/unit/test_handle.py index 591a1b796..89b0d99aa 100644 --- a/tests/unit/test_handle.py +++ b/tests/unit/test_handle.py @@ -220,10 +220,11 @@ def test_prepare_connection_arguments__default_uses_credentials_provider(self): assert "oauth_client_id" not in args assert "access_token" not in args - def _kernel_args_with_token(self, bearer, **creds_kwargs): - """Build kernel connect() kwargs with the credential manager's - header_factory stubbed to yield the given Authorization header, so the - machine-auth path can be exercised without live authentication.""" + def _kernel_azure_args(self, bearer, **creds_kwargs): + """Build kernel connect() kwargs for an Azure SP: the credential manager's + resolved config.auth_type is stubbed to 'azure-client-secret' and its + header_factory yields the given Authorization header, so the Azure bearer + fallback is exercised without live authentication.""" creds = DatabricksCredentials( host="yourorg.databricks.com", http_path=_KERNEL_HTTP_PATH, @@ -233,40 +234,64 @@ def _kernel_args_with_token(self, bearer, **creds_kwargs): ) manager = DatabricksCredentialManager.create_from(creds) header = {"Authorization": bearer} if bearer is not None else {} - with patch.object( - type(manager), - "header_factory", - new_callable=lambda: property(lambda self: (lambda: header)), + fake_config = Mock(auth_type="azure-client-secret") + with ( + patch.object( + type(manager), + "config", + new_callable=lambda: property(lambda self: fake_config), + ), + patch.object( + type(manager), + "header_factory", + new_callable=lambda: property(lambda self: (lambda: header)), + ), ): return SqlUtils.prepare_connection_arguments(creds, manager, _KERNEL_HTTP_PATH, {}) - def test_prepare_connection_arguments__kernel_pat_forwards_resolved_token(self): - """use_kernel with a PAT resolves the bearer token from the credential - manager and forwards it as access_token, dropping credentials_provider.""" - args = self._kernel_args_with_token("Bearer dapiabc123", token="dapiabc123") + 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( + 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_resolved_token(self): - """OAuth M2M is resolved to a Databricks workspace token by the credential - manager, and that token is forwarded to the kernel's PAT path — not the raw - client id/secret.""" - args = self._kernel_args_with_token( - "Bearer exchanged-m2m-token", client_id="cid", client_secret="dose-secret" + def test_prepare_connection_arguments__kernel_m2m_forwards_raw_creds(self): + """OAuth M2M forwards the raw client id/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}, ) - assert args["access_token"] == "exchanged-m2m-token" + 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" 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_azure_sp_forwards_resolved_token(self): - """The Azure-AD service principal (the peco CI credential) is resolved to a - Databricks workspace token by the credential manager's azure-client-secret - auth, and that token drives the kernel's PAT path. This is what makes the - kernel work for Azure SP — forwarding the raw creds to the kernel's - Databricks-OAuth M2M would fail with invalid_client.""" - args = self._kernel_args_with_token( + """The Azure SP (the peco CI credential) has no kernel-native flow, so as a + fallback dbt resolves it to a Databricks workspace token via + azure-client-secret auth and forwards that token to the kernel's PAT path. + Forwarding the raw creds to the kernel's Databricks-OAuth M2M would fail with + invalid_client.""" + args = self._kernel_azure_args( "Bearer azure-exchanged-token", azure_client_id="acid", azure_client_secret="asecret", @@ -275,12 +300,13 @@ def test_prepare_connection_arguments__kernel_azure_sp_forwards_resolved_token(s assert args.get("credentials_provider") is None assert "oauth_client_id" not in args - def test_prepare_connection_arguments__kernel_non_bearer_header_raises(self): - """If the credential manager yields a non-Bearer (or missing) Authorization - header, there is no token to forward — fail loudly rather than hand the - kernel unusable auth.""" + def test_prepare_connection_arguments__kernel_azure_non_bearer_raises(self): + """If the Azure exchange yields no usable Bearer header, fail loudly rather + than hand the kernel unusable auth.""" with pytest.raises(DbtConfigError): - self._kernel_args_with_token("Basic abc123", token="dapiabc123") + self._kernel_azure_args( + "Basic abc123", azure_client_id="acid", azure_client_secret="asecret" + ) def test_prepare_connection_arguments__kernel_u2m_explicit_client_id(self): """use_kernel with OAuth U2M and an explicit client_id forwards From 273138dd509c6fbdf37a7e80f5964ed86a262319 Mon Sep 17 00:00:00 2001 From: Shubham Dhal Date: Tue, 7 Jul 2026 14:22:40 +0530 Subject: [PATCH 4/8] Address review: restructure kernel auth into PAT/U2M/M2M paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Split _add_kernel_auth_arguments into three clear identification paths per review — PAT, U2M, and M2M — and reject/route the Azure SP up front: - Native Azure SP (azure_client_id/secret) routes to the bearer fallback before the M2M resolution, since it has no kernel-native flow. - U2M is identified simply as auth_type == "oauth" with no client secret. - M2M forwards raw creds when the resolved config.auth_type is oauth-m2m (kernel owns refresh); a legacy Azure SP passed via client_id/client_secret resolves to azure-client-secret and takes the bearer fallback. Extract the bearer-fallback into a helper. Terser docstring and CHANGELOG. --- CHANGELOG.md | 2 +- dbt/adapters/databricks/handle.py | 112 +++++++++++++++--------------- tests/unit/test_handle.py | 40 +++++++---- 3 files changed, 84 insertions(+), 70 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 883db0889..22a13bd9f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +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 (requires `databricks-sql-connector[kernel]` on Python 3.10+). The kernel's auth bridge rejects the connector `credentials_provider` dbt uses on the default path, so the adapter forwards the raw credentials the kernel understands — personal access token or Databricks OAuth (M2M/U2M) — letting the kernel own the token lifecycle, including OAuth refresh. Azure service principals have no kernel-native flow, so they fall back to a pre-resolved bearer token (logged with a warning) that does not auto-refresh; use a PAT or Databricks OAuth, or the default backend, for long-lived Azure connections ([#1576](https://github.com/databricks/dbt-databricks/pull/1576)) +- Support the connector's Rust kernel backend via `connection_parameters: {use_kernel: true}` for SQL warehouses (requires `databricks-sql-connector[kernel]` on Python 3.10+). PAT and Databricks OAuth (M2M/U2M) are supported natively; Azure service principals fall back to a pre-resolved bearer token that does not auto-refresh ([#1576](https://github.com/databricks/dbt-databricks/pull/1576)) ### Fixes - Honor the `expression` field on `primary_key` constraints on the V1 materialization path. A primary key declared with `expression: RELY` (or any trailing clause) previously had its expression silently dropped. ([#1551](https://github.com/databricks/dbt-databricks/pull/1551)) diff --git a/dbt/adapters/databricks/handle.py b/dbt/adapters/databricks/handle.py index 2a66e9530..90456338c 100644 --- a/dbt/adapters/databricks/handle.py +++ b/dbt/adapters/databricks/handle.py @@ -434,40 +434,33 @@ def _add_kernel_auth_arguments( ) -> None: """Populate connection kwargs for the connector's Rust kernel backend. - The kernel (use_kernel=True) routes over SEA and owns the token lifecycle - itself, so its auth bridge rejects the connector credentials_provider that - dbt passes on the Thrift path. It accepts only three flows — PAT, Databricks - OAuth M2M, and OAuth U2M — and has no Azure-AD service-principal flow. - - We forward the raw credentials the kernel understands so that IT owns the - token lifecycle, including refresh: for OAuth M2M the kernel re-mints tokens - via workspace OIDC, so a long-running connection stays authenticated. This - is the important property for real dbt runs, which are not bounded to a - token's lifetime. - - The Azure service principal is the exception: the kernel cannot mint Azure - tokens, so there is no raw-cred path for it. As a fallback we let dbt's - credential manager perform the Azure exchange and forward the resulting - Databricks workspace bearer token to the kernel's PAT path. That token is a - point-in-time value the kernel does NOT refresh, so a connection outliving - it (~1h) would need to be reopened — we log a warning to make that explicit. - Routing uses dbt's resolved ``config.auth_type`` because an Azure SP may be - supplied either natively (azure_client_id/secret) or, legacily, through the - M2M client_id/client_secret fields — indistinguishable by input shape alone. + The kernel (use_kernel=True) routes over SEA and its auth bridge rejects the + connector credentials_provider dbt passes on the Thrift path. It accepts + PAT, Databricks OAuth M2M, and OAuth U2M, and has no Azure-AD flow. + + We forward the raw credentials the kernel understands so the connector + handles the token lifecycle and refresh itself. The Azure service principal + is the exception (the kernel cannot mint Azure tokens): dbt performs the + Azure exchange and forwards the resulting bearer token to the kernel's PAT + path. Whether an SP is Azure is not always visible from the input fields — + it may be supplied natively (azure_client_id/secret) or legacily through the + client_id/client_secret fields — so the M2M-vs-Azure split uses dbt's + resolved ``config.auth_type``. """ - # PAT: forward the token directly (long-lived; no refresh concern). + # Native Azure SP has no kernel flow; the only route is the bearer fallback. + if creds.azure_client_id and creds.azure_client_secret: + SqlUtils._add_kernel_azure_bearer_token(args, creds_manager) + return + + # PAT: forward the token directly. if creds.token: args["access_token"] = creds.token return - # OAuth U2M (browser): the kernel runs the interactive flow itself, so there - # is no token to pre-resolve. Translate dbt's auth_type "oauth" to the - # kernel's "databricks-oauth"; client_id is optional (default dbt app). - if ( - creds.auth_type == "oauth" - and not creds.client_secret - and not (creds.azure_client_id and creds.azure_client_secret) - ): + # OAuth U2M (browser): auth_type "oauth" with no client secret. The kernel + # runs the interactive flow itself, so there is no token to pre-resolve — + # translate to the kernel's "databricks-oauth" (client_id is optional). + if creds.auth_type == "oauth" and not creds.client_secret: args["auth_type"] = "databricks-oauth" if creds.client_id: args["oauth_client_id"] = creds.client_id @@ -475,38 +468,43 @@ def _add_kernel_auth_arguments( args["oauth_scopes"] = creds.oauth_scopes return - # Resolve what dbt actually authenticates as. For OAuth M2M we forward raw - # creds so the kernel owns refresh; for Azure (native or legacy-via-M2M- - # fields) the kernel has no flow, so we fall back to a resolved bearer token. - resolved_auth_type = creds_manager.config.auth_type - - if resolved_auth_type == "oauth-m2m": + # M2M: forward raw creds when it is genuine Databricks OAuth (the kernel + # refreshes it); fall back to a bearer token when the resolved auth is an + # Azure SP supplied through the client_id/client_secret fields. + if creds_manager.config.auth_type == "oauth-m2m": args["oauth_client_id"] = creds.client_id args["oauth_client_secret"] = creds.client_secret if creds.oauth_scopes: args["oauth_scopes"] = creds.oauth_scopes return - if resolved_auth_type == "azure-client-secret": - logger.warning( - "use_kernel=True with an Azure service principal: the kernel has no " - "Azure-AD auth flow, so dbt forwards a pre-resolved bearer token that " - "the kernel will NOT refresh. A connection open for longer than the " - "token's lifetime (~1h) will fail to re-authenticate; use a personal " - "access token or Databricks OAuth, or the default backend, for " - "long-lived connections." - ) - headers: dict[str, str] = {} - headers.update(creds_manager.header_factory()) - authorization = headers.get("Authorization", "") - prefix = "Bearer " - if authorization[: len(prefix)].lower() == prefix.lower(): - args["access_token"] = authorization[len(prefix) :] - return - - raise DbtConfigError( - "use_kernel=True could not resolve usable credentials. The kernel backend " - "supports personal access tokens, Databricks OAuth (M2M/U2M), and Azure " - "service principals; other auth methods must use the default backend " - "(drop use_kernel)." + SqlUtils._add_kernel_azure_bearer_token(args, creds_manager) + + @staticmethod + def _add_kernel_azure_bearer_token( + args: dict[str, Any], creds_manager: DatabricksCredentialManager + ) -> None: + """Forward an Azure SP's resolved bearer token to the kernel's PAT path. + + The kernel does not refresh this token, so a connection outliving it (~1h) + would need to be reopened; warn so that is explicit. + """ + logger.warning( + "use_kernel=True with an Azure service principal forwards a pre-resolved " + "bearer token that the kernel does not refresh; a connection open longer " + "than the token's lifetime will fail to re-authenticate. Use a personal " + "access token, Databricks OAuth, or the default backend for long-lived " + "connections." ) + headers: dict[str, str] = {} + headers.update(creds_manager.header_factory()) + authorization = headers.get("Authorization", "") + prefix = "Bearer " + if not authorization[: len(prefix)].lower() == prefix.lower(): + raise DbtConfigError( + "use_kernel=True could not resolve a bearer token from the configured " + "credentials. The kernel backend supports personal access tokens, " + "Databricks OAuth (M2M/U2M), and Azure service principals; other auth " + "methods must use the default backend (drop use_kernel)." + ) + args["access_token"] = authorization[len(prefix) :] diff --git a/tests/unit/test_handle.py b/tests/unit/test_handle.py index 89b0d99aa..0648eb975 100644 --- a/tests/unit/test_handle.py +++ b/tests/unit/test_handle.py @@ -220,11 +220,12 @@ def test_prepare_connection_arguments__default_uses_credentials_provider(self): assert "oauth_client_id" not in args assert "access_token" not in args - def _kernel_azure_args(self, bearer, **creds_kwargs): - """Build kernel connect() kwargs for an Azure SP: the credential manager's - resolved config.auth_type is stubbed to 'azure-client-secret' and its - header_factory yields the given Authorization header, so the Azure bearer - fallback is exercised without live authentication.""" + def _kernel_azure_args(self, bearer, resolved_auth_type="azure-client-secret", **creds_kwargs): + """Build kernel connect() kwargs for an Azure SP with the credential + manager's header_factory stubbed to yield the given Authorization header, + exercising the Azure bearer fallback without live authentication. The + resolved config.auth_type is stubbed (only read for the legacy-via-M2M- + fields path; the native-fields path routes before config is consulted).""" creds = DatabricksCredentials( host="yourorg.databricks.com", http_path=_KERNEL_HTTP_PATH, @@ -234,7 +235,7 @@ def _kernel_azure_args(self, bearer, **creds_kwargs): ) manager = DatabricksCredentialManager.create_from(creds) header = {"Authorization": bearer} if bearer is not None else {} - fake_config = Mock(auth_type="azure-client-secret") + fake_config = Mock(auth_type=resolved_auth_type) with ( patch.object( type(manager), @@ -285,12 +286,11 @@ def test_prepare_connection_arguments__kernel_m2m_forwards_raw_creds(self): assert args.get("credentials_provider") is None assert "access_token" not in args - def test_prepare_connection_arguments__kernel_azure_sp_forwards_resolved_token(self): - """The Azure SP (the peco CI credential) has no kernel-native flow, so as a - fallback dbt resolves it to a Databricks workspace token via - azure-client-secret auth and forwards that token to the kernel's PAT path. - Forwarding the raw creds to the kernel's Databricks-OAuth M2M would fail with - invalid_client.""" + def test_prepare_connection_arguments__kernel_native_azure_sp_forwards_token(self): + """A native Azure SP (azure_client_id/secret) has no kernel-native flow, so + dbt resolves it to a Databricks workspace token and forwards that to the + kernel's PAT path. This path routes on the azure fields before config is + consulted.""" args = self._kernel_azure_args( "Bearer azure-exchanged-token", azure_client_id="acid", @@ -300,6 +300,22 @@ def test_prepare_connection_arguments__kernel_azure_sp_forwards_resolved_token(s assert args.get("credentials_provider") is None assert "oauth_client_id" not in args + def test_prepare_connection_arguments__kernel_legacy_azure_via_m2m_fields_forwards_token(self): + """The peco CI credential is an Azure SP supplied through the + client_id/client_secret fields, which dbt resolves as azure-client-secret. + It must take the bearer fallback (forwarding the raw creds to the kernel's + Databricks-OAuth M2M fails with invalid_client).""" + args = self._kernel_azure_args( + "Bearer azure-exchanged-token", + resolved_auth_type="azure-client-secret", + client_id="cid", + client_secret="azure-secret", + ) + assert args["access_token"] == "azure-exchanged-token" + assert args.get("credentials_provider") is None + assert "oauth_client_id" not in args + assert "oauth_client_secret" not in args + def test_prepare_connection_arguments__kernel_azure_non_bearer_raises(self): """If the Azure exchange yields no usable Bearer header, fail loudly rather than hand the kernel unusable auth.""" From 81aff5fc107bf1192fa281dcf1bb38e1e45ffba0 Mon Sep 17 00:00:00 2001 From: Shubham Dhal Date: Tue, 7 Jul 2026 16:40:00 +0530 Subject: [PATCH 5/8] Restrict kernel auth to PAT/Databricks-OAuth; resolve Azure SP in test setup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review, the production kernel path should support only what the kernel natively does: PAT, Databricks OAuth M2M (raw creds so the kernel owns refresh), and OAuth U2M. Any other auth — notably an Azure service principal, which the kernel has no flow for — is rejected with a DbtConfigError directing the user to the default backend. Removes the in-adapter Azure->bearer-token conversion that was only needed to run CI on a peco Azure SP. That conversion now lives in test setup (tests/profiles.py): when DBT_DATABRICKS_USE_KERNEL=1 and the credential is an SP, resolve it to a workspace token and hand the kernel a PAT. Marked temporary — the resolved token does not auto-refresh, and it goes away once the kernel gains native Azure support. Also restores test concurrency to -n 10 to match the other integration lanes. --- .github/workflows/integration-kernel.yml | 15 +++-- CHANGELOG.md | 2 +- dbt/adapters/databricks/handle.py | 77 +++++++---------------- tests/profiles.py | 35 +++++++++++ tests/unit/test_handle.py | 79 ++++++++++-------------- 5 files changed, 96 insertions(+), 112 deletions(-) diff --git a/.github/workflows/integration-kernel.yml b/.github/workflows/integration-kernel.yml index dc5b208fb..4eec604f1 100644 --- a/.github/workflows/integration-kernel.yml +++ b/.github/workflows/integration-kernel.yml @@ -12,13 +12,12 @@ # cluster profiles from the main matrix cannot run on the kernel and are not # included here. # -# Auth: the kernel owns the token lifecycle itself and its auth bridge rejects -# the connector credentials_provider dbt passes on the Thrift path. When -# DBT_DATABRICKS_USE_KERNEL=1, tests/profiles.py sets use_kernel=True and the -# adapter forwards the raw credentials the kernel understands (PAT or Databricks -# OAuth) so the kernel owns refresh. The peco CI credential is an Azure service -# principal, which has no kernel-native flow, so it falls back to a pre-resolved -# bearer token (no auto-refresh) — see SqlUtils.prepare_connection_arguments. +# 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. name: Kernel Integration Tests on: workflow_dispatch: @@ -171,7 +170,7 @@ jobs: xargs -r hatch -v run pytest \ --color=yes -v \ --profile databricks_uc_sql_endpoint \ - -n 4 --dist=loadfile \ + -n 10 --dist=loadfile \ --reruns 2 --reruns-delay 120 \ --junitxml=logs/junit-kernel-sql-endpoint-shard-${{ matrix.shard }}.xml \ < "$SHARD_TESTS_FILE" diff --git a/CHANGELOG.md b/CHANGELOG.md index 883db0889..495dd6493 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +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 (requires `databricks-sql-connector[kernel]` on Python 3.10+). The kernel's auth bridge rejects the connector `credentials_provider` dbt uses on the default path, so the adapter forwards the raw credentials the kernel understands — personal access token or Databricks OAuth (M2M/U2M) — letting the kernel own the token lifecycle, including OAuth refresh. Azure service principals have no kernel-native flow, so they fall back to a pre-resolved bearer token (logged with a warning) that does not auto-refresh; use a PAT or Databricks OAuth, or the default backend, for long-lived Azure connections ([#1576](https://github.com/databricks/dbt-databricks/pull/1576)) +- 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 - Honor the `expression` field on `primary_key` constraints on the V1 materialization path. A primary key declared with `expression: RELY` (or any trailing clause) previously had its expression silently dropped. ([#1551](https://github.com/databricks/dbt-databricks/pull/1551)) diff --git a/dbt/adapters/databricks/handle.py b/dbt/adapters/databricks/handle.py index 2a66e9530..0dec6ba90 100644 --- a/dbt/adapters/databricks/handle.py +++ b/dbt/adapters/databricks/handle.py @@ -434,40 +434,24 @@ def _add_kernel_auth_arguments( ) -> None: """Populate connection kwargs for the connector's Rust kernel backend. - The kernel (use_kernel=True) routes over SEA and owns the token lifecycle - itself, so its auth bridge rejects the connector credentials_provider that - dbt passes on the Thrift path. It accepts only three flows — PAT, Databricks - OAuth M2M, and OAuth U2M — and has no Azure-AD service-principal flow. - - We forward the raw credentials the kernel understands so that IT owns the - token lifecycle, including refresh: for OAuth M2M the kernel re-mints tokens - via workspace OIDC, so a long-running connection stays authenticated. This - is the important property for real dbt runs, which are not bounded to a - token's lifetime. - - The Azure service principal is the exception: the kernel cannot mint Azure - tokens, so there is no raw-cred path for it. As a fallback we let dbt's - credential manager perform the Azure exchange and forward the resulting - Databricks workspace bearer token to the kernel's PAT path. That token is a - point-in-time value the kernel does NOT refresh, so a connection outliving - it (~1h) would need to be reopened — we log a warning to make that explicit. - Routing uses dbt's resolved ``config.auth_type`` because an Azure SP may be - supplied either natively (azure_client_id/secret) or, legacily, through the - M2M client_id/client_secret fields — indistinguishable by input shape alone. + 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 (long-lived; no refresh concern). + # PAT: forward the token directly. if creds.token: args["access_token"] = creds.token return - # OAuth U2M (browser): the kernel runs the interactive flow itself, so there - # is no token to pre-resolve. Translate dbt's auth_type "oauth" to the - # kernel's "databricks-oauth"; client_id is optional (default dbt app). - if ( - creds.auth_type == "oauth" - and not creds.client_secret - and not (creds.azure_client_id and creds.azure_client_secret) - ): + # 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. + if creds.auth_type == "oauth" and not creds.client_secret and not creds.azure_client_secret: args["auth_type"] = "databricks-oauth" if creds.client_id: args["oauth_client_id"] = creds.client_id @@ -475,38 +459,19 @@ def _add_kernel_auth_arguments( args["oauth_scopes"] = creds.oauth_scopes return - # Resolve what dbt actually authenticates as. For OAuth M2M we forward raw - # creds so the kernel owns refresh; for Azure (native or legacy-via-M2M- - # fields) the kernel has no flow, so we fall back to a resolved bearer token. - resolved_auth_type = creds_manager.config.auth_type - - if resolved_auth_type == "oauth-m2m": + # Databricks OAuth M2M: forward raw 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.azure_client_secret and creds_manager.config.auth_type == "oauth-m2m": args["oauth_client_id"] = creds.client_id args["oauth_client_secret"] = creds.client_secret if creds.oauth_scopes: args["oauth_scopes"] = creds.oauth_scopes return - if resolved_auth_type == "azure-client-secret": - logger.warning( - "use_kernel=True with an Azure service principal: the kernel has no " - "Azure-AD auth flow, so dbt forwards a pre-resolved bearer token that " - "the kernel will NOT refresh. A connection open for longer than the " - "token's lifetime (~1h) will fail to re-authenticate; use a personal " - "access token or Databricks OAuth, or the default backend, for " - "long-lived connections." - ) - headers: dict[str, str] = {} - headers.update(creds_manager.header_factory()) - authorization = headers.get("Authorization", "") - prefix = "Bearer " - if authorization[: len(prefix)].lower() == prefix.lower(): - args["access_token"] = authorization[len(prefix) :] - return - raise DbtConfigError( - "use_kernel=True could not resolve usable credentials. The kernel backend " - "supports personal access tokens, Databricks OAuth (M2M/U2M), and Azure " - "service principals; other auth methods must use the default backend " - "(drop use_kernel)." + "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/tests/profiles.py b/tests/profiles.py index c04894e7c..4ceb5f81f 100644 --- a/tests/profiles.py +++ b/tests/profiles.py @@ -56,10 +56,45 @@ def _build_databricks_cluster_target( # 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 89b0d99aa..f1372fa42 100644 --- a/tests/unit/test_handle.py +++ b/tests/unit/test_handle.py @@ -220,35 +220,6 @@ def test_prepare_connection_arguments__default_uses_credentials_provider(self): assert "oauth_client_id" not in args assert "access_token" not in args - def _kernel_azure_args(self, bearer, **creds_kwargs): - """Build kernel connect() kwargs for an Azure SP: the credential manager's - resolved config.auth_type is stubbed to 'azure-client-secret' and its - header_factory yields the given Authorization header, so the Azure bearer - fallback is exercised without live authentication.""" - creds = DatabricksCredentials( - host="yourorg.databricks.com", - http_path=_KERNEL_HTTP_PATH, - schema="dbt", - connection_parameters={"use_kernel": True}, - **creds_kwargs, - ) - manager = DatabricksCredentialManager.create_from(creds) - header = {"Authorization": bearer} if bearer is not None else {} - fake_config = Mock(auth_type="azure-client-secret") - with ( - patch.object( - type(manager), - "config", - new_callable=lambda: property(lambda self: fake_config), - ), - patch.object( - type(manager), - "header_factory", - new_callable=lambda: property(lambda self: (lambda: header)), - ), - ): - return SqlUtils.prepare_connection_arguments(creds, manager, _KERNEL_HTTP_PATH, {}) - 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.""" @@ -285,28 +256,42 @@ def test_prepare_connection_arguments__kernel_m2m_forwards_raw_creds(self): assert args.get("credentials_provider") is None assert "access_token" not in args - def test_prepare_connection_arguments__kernel_azure_sp_forwards_resolved_token(self): - """The Azure SP (the peco CI credential) has no kernel-native flow, so as a - fallback dbt resolves it to a Databricks workspace token via - azure-client-secret auth and forwards that token to the kernel's PAT path. - Forwarding the raw creds to the kernel's Databricks-OAuth M2M would fail with - invalid_client.""" - args = self._kernel_azure_args( - "Bearer azure-exchanged-token", + 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}, ) - assert args["access_token"] == "azure-exchanged-token" - assert args.get("credentials_provider") is None - assert "oauth_client_id" not in args + manager = DatabricksCredentialManager.create_from(creds) + with pytest.raises(DbtConfigError, match="use_kernel"): + SqlUtils.prepare_connection_arguments(creds, manager, _KERNEL_HTTP_PATH, {}) - def test_prepare_connection_arguments__kernel_azure_non_bearer_raises(self): - """If the Azure exchange yields no usable Bearer header, fail loudly rather - than hand the kernel unusable auth.""" - with pytest.raises(DbtConfigError): - self._kernel_azure_args( - "Basic abc123", azure_client_id="acid", azure_client_secret="asecret" - ) + 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 From 8c2628dd6f01eb8504ba982206d36f485b0eb182 Mon Sep 17 00:00:00 2001 From: Shubham Dhal Date: Tue, 7 Jul 2026 21:09:09 +0530 Subject: [PATCH 6/8] Tune kernel lane: -n 4 across 5 shards to avoid session eviction + token TTL MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The kernel evicts SEA sessions under sustained -n 10 load, so late in-test verification queries intermittently fail with "Cannot execute_command without an open session". Drop the kernel lane to -n 4 to keep the concurrent session count low. At -n 4 the 3-shard run took ~60min/shard — right at the ~1h lifetime of the Azure token the tests use — so split into 5 shards (~35min each) to keep every job well clear of that boundary. Both are kernel-lane workarounds documented in the workflow; the other lanes stay at -n 10 / 3 shards. --- .github/workflows/integration-kernel.yml | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/.github/workflows/integration-kernel.yml b/.github/workflows/integration-kernel.yml index 4eec604f1..2d01e6dea 100644 --- a/.github/workflows/integration-kernel.yml +++ b/.github/workflows/integration-kernel.yml @@ -18,6 +18,12 @@ # 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: @@ -71,7 +77,7 @@ jobs: > "shard-assignments/databricks_uc_sql_endpoint-collected.txt" python3 scripts/shard_assign.py \ --profile databricks_uc_sql_endpoint \ - --num-shards 3 \ + --num-shards 5 \ --input "shard-assignments/databricks_uc_sql_endpoint-collected.txt" \ --output-dir shard-assignments \ --algo lpt_historical_time \ @@ -89,7 +95,7 @@ jobs: strategy: fail-fast: false matrix: - shard: [0, 1, 2] + shard: [0, 1, 2, 3, 4] runs-on: group: databricks-protected-runner-group labels: linux-ubuntu-latest @@ -164,13 +170,21 @@ jobs: 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 10 --dist=loadfile \ + -n 4 --dist=loadfile \ --reruns 2 --reruns-delay 120 \ --junitxml=logs/junit-kernel-sql-endpoint-shard-${{ matrix.shard }}.xml \ < "$SHARD_TESTS_FILE" From 30962fa55ef517c6e49f81b94e803d3fd467c9e8 Mon Sep 17 00:00:00 2001 From: Shubham Dhal Date: Tue, 7 Jul 2026 22:38:47 +0530 Subject: [PATCH 7/8] Remove temp pull_request trigger from kernel workflow The temp trigger was added to validate the workflow on PR #1576; CI is green across all shards (run 28882012926), so revert to the intended workflow_dispatch + weekly schedule triggers. --- .github/workflows/integration-kernel.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/integration-kernel.yml b/.github/workflows/integration-kernel.yml index 2d01e6dea..0d6fc9012 100644 --- a/.github/workflows/integration-kernel.yml +++ b/.github/workflows/integration-kernel.yml @@ -27,7 +27,6 @@ name: Kernel Integration Tests on: workflow_dispatch: - pull_request: # TEMP: verify the workflow runs on PR #1576; reverted before merge. 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, From 1417cbc0db36fe5bd8f8c97ef8d7b205ea4abc04 Mon Sep 17 00:00:00 2001 From: Shubham Dhal Date: Wed, 8 Jul 2026 19:00:04 +0530 Subject: [PATCH 8/8] fix: use resolved kernel auth defaults Forward kernel OAuth fields from the credential manager so default client IDs and scopes match the normal auth path while avoiding SDK auth initialization on short-circuit flows. --- dbt/adapters/databricks/handle.py | 28 ++++++++------- tests/unit/test_handle.py | 60 +++++++++++++++++++++++-------- 2 files changed, 60 insertions(+), 28 deletions(-) diff --git a/dbt/adapters/databricks/handle.py b/dbt/adapters/databricks/handle.py index 0dec6ba90..3da799c9c 100644 --- a/dbt/adapters/databricks/handle.py +++ b/dbt/adapters/databricks/handle.py @@ -444,29 +444,31 @@ def _add_kernel_auth_arguments( error directing the user to the default backend. """ # PAT: forward the token directly. - if creds.token: - args["access_token"] = creds.token + 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. - if creds.auth_type == "oauth" and not creds.client_secret and not creds.azure_client_secret: + # 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" - if creds.client_id: - args["oauth_client_id"] = creds.client_id - if creds.oauth_scopes: - args["oauth_scopes"] = creds.oauth_scopes + args["oauth_client_id"] = creds_manager.client_id + args["oauth_scopes"] = creds_manager.oauth_scopes return - # Databricks OAuth M2M: forward raw creds so the kernel owns refresh. The + # 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.azure_client_secret and creds_manager.config.auth_type == "oauth-m2m": - args["oauth_client_id"] = creds.client_id - args["oauth_client_secret"] = creds.client_secret - if creds.oauth_scopes: - args["oauth_scopes"] = creds.oauth_scopes + 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( diff --git a/tests/unit/test_handle.py b/tests/unit/test_handle.py index f1372fa42..8f58c25a6 100644 --- a/tests/unit/test_handle.py +++ b/tests/unit/test_handle.py @@ -75,6 +75,26 @@ def _prepare_connection_args(**creds_kwargs: Any) -> dict[str, Any]: 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])] @@ -223,7 +243,7 @@ def test_prepare_connection_arguments__default_uses_credentials_provider(self): 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( + args = _prepare_connection_args_without_config_auth( token="dapiabc123", connection_parameters={"use_kernel": True}, ) @@ -232,10 +252,10 @@ def test_prepare_connection_arguments__kernel_pat_forwards_access_token(self): assert args.get("credentials_provider") is None assert "oauth_client_id" not in args - def test_prepare_connection_arguments__kernel_m2m_forwards_raw_creds(self): - """OAuth M2M forwards the raw client id/secret so the KERNEL owns the token - lifecycle (acquire + refresh via workspace OIDC) — not a pre-resolved token - that would not refresh.""" + 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, @@ -253,6 +273,8 @@ def test_prepare_connection_arguments__kernel_m2m_forwards_raw_creds(self): 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 @@ -270,8 +292,13 @@ def test_prepare_connection_arguments__kernel_azure_sp_raises(self): connection_parameters={"use_kernel": True}, ) manager = DatabricksCredentialManager.create_from(creds) - with pytest.raises(DbtConfigError, match="use_kernel"): - SqlUtils.prepare_connection_arguments(creds, manager, _KERNEL_HTTP_PATH, {}) + 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 @@ -294,33 +321,36 @@ def test_prepare_connection_arguments__kernel_legacy_azure_sp_raises(self): 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 - oauth_client_id and translates dbt's auth_type='oauth' to the kernel's + """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( + 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_no_client_id(self): + 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 (SDK uses the default dbt app). The kernel bridge routes U2M off - auth_type alone, so we must translate auth_type even without a client_id.""" - args = _prepare_connection_args( + 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_id" not in args assert "oauth_client_secret" not in args assert "access_token" not in args