Skip to content

Commit ef3c482

Browse files
committed
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.
1 parent 13cfb8d commit ef3c482

11 files changed

Lines changed: 397 additions & 7 deletions

File tree

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,182 @@
1+
# Kernel Integration Tests for dbt-databricks.
2+
#
3+
# Runs the functional suite against the SQL-warehouse profile with the
4+
# databricks-sql-connector Rust kernel backend (use_kernel=True) instead of
5+
# the default Thrift backend, on a weekly Saturday schedule. Like
6+
# integration-spog.yml this is a scheduled smoke, not a PR gate — PR-targeting
7+
# and status-reporting are omitted.
8+
#
9+
# Why warehouse-only: the kernel routes over the SEA (Statement Execution API)
10+
# HTTP transport, which is warehouse-only — the connector rejects an
11+
# all-purpose-cluster http_path ("SEA only works for warehouses"). So the two
12+
# cluster profiles from the main matrix cannot run on the kernel and are not
13+
# included here.
14+
#
15+
# Auth: the kernel owns the OAuth token lifecycle itself and its auth bridge
16+
# rejects the connector credentials_provider dbt passes on the Thrift path.
17+
# When DBT_DATABRICKS_USE_KERNEL=1, tests/profiles.py sets use_kernel=True and
18+
# the adapter forwards the raw OAuth M2M client id/secret (or a PAT) instead —
19+
# see SqlUtils.prepare_connection_arguments.
20+
name: Kernel Integration Tests
21+
on:
22+
workflow_dispatch:
23+
schedule:
24+
- cron: "17 9 * * 6" # Weekly: Saturday 09:17 UTC (14:47 IST). Clear of the
25+
# daily 21:30 UTC nightly and the Sunday SPOG slot,
26+
# which share the peco warehouse.
27+
28+
permissions:
29+
id-token: write
30+
contents: read
31+
32+
concurrency:
33+
group: ${{ github.workflow }}-${{ github.ref }}
34+
cancel-in-progress: true
35+
36+
jobs:
37+
prepare-shards:
38+
runs-on:
39+
group: databricks-protected-runner-group
40+
labels: linux-ubuntu-latest
41+
env:
42+
UV_FROZEN: "1"
43+
steps:
44+
- name: Check out the repository
45+
uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
46+
47+
- name: Setup JFrog PyPI Proxy
48+
uses: ./.github/actions/setup-jfrog-pypi
49+
50+
- name: Set up Python
51+
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
52+
with:
53+
python-version: "3.10"
54+
55+
- name: Install uv
56+
uses: astral-sh/setup-uv@38f3f104447c67c051c4a08e39b64a148898af3a # v4
57+
with:
58+
cache-local-path: ~/.cache/uv
59+
60+
- name: Install Hatch
61+
uses: pypa/hatch@257e27e51a6a5616ed08a39a408a21c35c9931bc # install
62+
63+
- name: Collect tests and assign to shards
64+
run: |
65+
set -euo pipefail
66+
mkdir -p shard-assignments
67+
hatch run pytest --collect-only -q --profile databricks_uc_sql_endpoint tests/functional 2>&1 \
68+
| grep "::" \
69+
> "shard-assignments/databricks_uc_sql_endpoint-collected.txt"
70+
python3 scripts/shard_assign.py \
71+
--profile databricks_uc_sql_endpoint \
72+
--num-shards 3 \
73+
--input "shard-assignments/databricks_uc_sql_endpoint-collected.txt" \
74+
--output-dir shard-assignments \
75+
--algo lpt_historical_time \
76+
--timings .github/test_timings.json
77+
78+
- name: Upload shard assignments
79+
uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4
80+
with:
81+
name: shard-assignments-kernel
82+
path: shard-assignments/
83+
retention-days: 5
84+
85+
uc-sql-endpoint:
86+
needs: prepare-shards
87+
strategy:
88+
fail-fast: false
89+
matrix:
90+
shard: [0, 1, 2]
91+
runs-on:
92+
group: databricks-protected-runner-group
93+
labels: linux-ubuntu-latest
94+
environment: azure-prod
95+
env:
96+
DBT_DATABRICKS_HOST_NAME: ${{ secrets.DATABRICKS_HOST }}
97+
DBT_DATABRICKS_CLIENT_ID: ${{ secrets.TEST_PECO_SP_ID }}
98+
DBT_DATABRICKS_CLIENT_SECRET: ${{ secrets.TEST_PECO_SP_SECRET }}
99+
DBT_DATABRICKS_HTTP_PATH: ${{ secrets.TEST_PECO_WAREHOUSE_HTTP_PATH }}
100+
DBT_DATABRICKS_UC_INITIAL_CATALOG: peco
101+
DBT_DATABRICKS_LOCATION_ROOT: ${{ secrets.TEST_PECO_EXTERNAL_LOCATION }}test
102+
# Wire the UC cluster id so build_cluster_http_path.py can populate
103+
# DBT_DATABRICKS_UC_CLUSTER_HTTP_PATH. The kernel-incompatible tests that
104+
# pin a model to that cluster path are @pytest.mark.skip_kernel and skip
105+
# here; this keeps the env var resolvable so their fixtures still parse.
106+
TEST_PECO_UC_CLUSTER_ID: ${{ secrets.TEST_PECO_UC_CLUSTER_ID }}
107+
TEST_PECO_SPOG_HOST: ${{ secrets.TEST_PECO_SPOG_HOST }}
108+
TEST_PECO_SPOG_WORKSPACE_ID: ${{ secrets.TEST_PECO_SPOG_WORKSPACE_ID }}
109+
# Route the whole suite through the connector's Rust kernel (SEA) backend.
110+
DBT_DATABRICKS_USE_KERNEL: "1"
111+
UV_FROZEN: "1"
112+
steps:
113+
- name: Check out repository
114+
uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
115+
116+
- name: Setup JFrog PyPI Proxy
117+
uses: ./.github/actions/setup-jfrog-pypi
118+
119+
- name: Set up python
120+
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
121+
with:
122+
python-version: "3.10"
123+
124+
- name: Get http path from environment
125+
run: python .github/workflows/build_cluster_http_path.py
126+
shell: sh
127+
128+
- name: Install uv
129+
uses: astral-sh/setup-uv@38f3f104447c67c051c4a08e39b64a148898af3a # v4
130+
with:
131+
cache-local-path: ~/.cache/uv
132+
133+
- name: Install Hatch
134+
uses: pypa/hatch@257e27e51a6a5616ed08a39a408a21c35c9931bc # install
135+
136+
- name: Install the connector kernel extra
137+
# The kernel backend ships in the optional databricks-sql-connector
138+
# [kernel] extra (a PyO3 wheel), which the frozen lock does not include.
139+
# Add it into the hatch env so use_kernel=True can import the kernel.
140+
run: |
141+
set -euo pipefail
142+
HATCH_PY="$(hatch run python -c 'import sys; print(sys.executable)')"
143+
hatch run python -c "import databricks.sql as s; print('connector version:', s.__version__)"
144+
uv pip install --python "$HATCH_PY" "databricks-sql-connector[kernel]"
145+
hatch run python -c "import databricks_sql_kernel; print('kernel extra installed')"
146+
147+
- name: Download shard assignments
148+
uses: actions/download-artifact@65a9edc5881444af0b9093a5e628f2fe47ea3b2e # v4.1.7
149+
with:
150+
name: shard-assignments-kernel
151+
path: shard-assignments/
152+
153+
- name: Resolve test list for this shard
154+
run: |
155+
set -euo pipefail
156+
SHARD_FILE="shard-assignments/databricks_uc_sql_endpoint-shard-${{ matrix.shard }}.txt"
157+
if [ ! -s "$SHARD_FILE" ]; then
158+
echo "::error::Shard file missing or empty: $SHARD_FILE"
159+
exit 1
160+
fi
161+
echo "SHARD_TESTS_FILE=$SHARD_FILE" >> "$GITHUB_ENV"
162+
echo "Files in shard ${{ matrix.shard }}: $(wc -l < "$SHARD_FILE")"
163+
164+
- name: Run Sql Endpoint Functional Tests on kernel (shard ${{ matrix.shard }})
165+
run: |
166+
mkdir -p logs
167+
DBT_TEST_USER=notnecessaryformosttests@example.com \
168+
xargs -r hatch -v run pytest \
169+
--color=yes -v \
170+
--profile databricks_uc_sql_endpoint \
171+
-n 10 --dist=loadfile \
172+
--reruns 2 --reruns-delay 120 \
173+
--junitxml=logs/junit-kernel-sql-endpoint-shard-${{ matrix.shard }}.xml \
174+
< "$SHARD_TESTS_FILE"
175+
176+
- name: Upload Kernel SQL Endpoint Test Logs
177+
if: always()
178+
uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4
179+
with:
180+
name: kernel-sql-endpoint-logs-shard-${{ matrix.shard }}
181+
path: logs/
182+
retention-days: 14

CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
- Add catalogs.yml v2 support (requires `use_catalogs_v2: true` in dbt-core) ([1440](https://github.com/databricks/dbt-databricks/pull/1440))
66
- 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)).
7+
- Support the connector's Rust kernel backend via `connection_parameters: {use_kernel: true}` for SQL warehouses with PAT or Databricks OAuth auth (requires `databricks-sql-connector[kernel]` on Python 3.10+; native Azure service-principal auth is not supported) ([#1576](https://github.com/databricks/dbt-databricks/pull/1576))
78

89
### Fixes
910
- 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 @@
2021

2122
### Under the Hood
2223

24+
- 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)).
2325
- 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).
2426
- 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))
2527
- 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))

dbt/adapters/databricks/handle.py

Lines changed: 52 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
from typing import TYPE_CHECKING, Any, Optional, TypeVar
1111

1212
from dbt.adapters.contracts.connection import AdapterResponse
13-
from dbt_common.exceptions import DbtRuntimeError
13+
from dbt_common.exceptions import DbtConfigError, DbtRuntimeError
1414

1515
import databricks.sql as dbsql
1616
from databricks.sql.client import Connection, Cursor
@@ -407,16 +407,64 @@ def prepare_connection_arguments(
407407
query_tags
408408
)
409409

410-
return {
410+
args: dict[str, Any] = {
411411
"server_hostname": creds.host,
412412
"http_path": http_path,
413-
"credentials_provider": creds_manager.credentials_provider,
414413
"http_headers": http_headers if http_headers else None,
415414
"session_configuration": session_config,
416415
"catalog": creds.database,
417416
"use_inline_params": "silent",
418417
"schema": creds.schema,
419418
"_user_agent_entry": user_agent_entry,
420419
"user_agent_entry": user_agent_entry,
421-
**connection_parameters,
422420
}
421+
422+
if connection_parameters.get("use_kernel"):
423+
SqlUtils._add_kernel_auth_arguments(args, creds)
424+
else:
425+
args["credentials_provider"] = creds_manager.credentials_provider
426+
427+
return {**args, **connection_parameters}
428+
429+
@staticmethod
430+
def _add_kernel_auth_arguments(args: dict[str, Any], creds: DatabricksCredentials) -> None:
431+
"""Populate connection kwargs for the connector's Rust kernel backend.
432+
433+
The kernel (use_kernel=True) routes over SEA and owns the token lifecycle
434+
itself, so its auth bridge rejects the connector credentials_provider that
435+
dbt passes on the Thrift path. We instead forward the raw credentials the
436+
bridge understands: PAT, OAuth M2M (client id + secret), or OAuth U2M
437+
(browser). The branch order mirrors DatabricksCredentialManager's own auth
438+
selection.
439+
440+
The kernel has no Azure-AD service-principal flow, so native Azure SP
441+
(azure_client_id/azure_client_secret) is rejected here with a clear error
442+
rather than routed to a path that can't serve it. A legacy Azure SP passed
443+
via client_id/client_secret is indistinguishable from Databricks OAuth M2M
444+
(the connector itself only disambiguates by trying both), so it is
445+
forwarded as M2M; if it is truly Azure it fails loudly at the kernel's
446+
token exchange — the same outcome as an invalid secret.
447+
"""
448+
if creds.token:
449+
args["access_token"] = creds.token
450+
elif creds.azure_client_id and creds.azure_client_secret:
451+
raise DbtConfigError(
452+
"use_kernel=True does not support Azure service-principal auth "
453+
"(azure_client_id/azure_client_secret). Use a personal access token "
454+
"or Databricks OAuth (client_id/client_secret) instead, or drop "
455+
"use_kernel to use the default backend."
456+
)
457+
elif creds.client_id and creds.client_secret:
458+
# OAuth M2M
459+
args["oauth_client_id"] = creds.client_id
460+
args["oauth_client_secret"] = creds.client_secret
461+
if creds.oauth_scopes:
462+
args["oauth_scopes"] = creds.oauth_scopes
463+
elif creds.auth_type == "oauth":
464+
# OAuth U2M (browser). Translate dbt's auth_type "oauth" to the kernel's
465+
# "databricks-oauth"; client_id is optional when using the default dbt app.
466+
args["auth_type"] = "databricks-oauth"
467+
if creds.client_id:
468+
args["oauth_client_id"] = creds.client_id
469+
if creds.oauth_scopes:
470+
args["oauth_scopes"] = creds.oauth_scopes

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,7 @@ markers = [
165165
"external: mark test as requiring an external location",
166166
"python: mark test as running a python model",
167167
"dlt: mark test as running a DLT model",
168+
"skip_kernel: skip test when running against the connector kernel backend (DBT_DATABRICKS_USE_KERNEL=1)",
168169
]
169170

170171
[tool.mypy]

tests/conftest.py

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,17 @@ def pytest_addoption(parser):
1313
parser.addoption("--profile", action="store", default=default_profile, type=str)
1414

1515

16-
# Using @pytest.mark.skip_profile('databricks_cluster') uses the 'skip_by_adapter_type'
17-
# autouse fixture below
16+
# @pytest.mark.skip_profile(...) and @pytest.mark.skip_kernel are enforced by the
17+
# autouse fixtures below (skip_by_profile_type, skip_on_kernel_backend).
1818
def pytest_configure(config):
1919
config.addinivalue_line(
2020
"markers",
2121
"skip_profile(profile): skip test for the given profile",
2222
)
23+
config.addinivalue_line(
24+
"markers",
25+
"skip_kernel: skip test when running against the connector kernel backend",
26+
)
2327

2428

2529
@pytest.fixture(scope="session")
@@ -37,6 +41,18 @@ def skip_by_profile_type(request):
3741
pytest.skip(f"skipped on '{profile_type}' profile")
3842

3943

44+
@pytest.fixture(autouse=True)
45+
def skip_on_kernel_backend(request):
46+
# The connector kernel backend (DBT_DATABRICKS_USE_KERNEL=1) routes over SEA,
47+
# which is warehouse-only. Tests that pin a model to an all-purpose-cluster
48+
# http_path or compute cannot connect through it, so they opt out with
49+
# @pytest.mark.skip_kernel.
50+
if os.environ.get("DBT_DATABRICKS_USE_KERNEL") == "1" and request.node.get_closest_marker(
51+
"skip_kernel"
52+
):
53+
pytest.skip("skipped on the connector kernel backend (SEA is warehouse-only)")
54+
55+
4056
# The profile dictionary, used to write out profiles.yml. It will pull in updates
4157
# from two separate sources, the 'profile_target' and 'profiles_config_update'.
4258
# The second one is useful when using alternative targets, etc.

tests/functional/adapter/incremental/test_incremental_column_tags.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from tests.functional.adapter.incremental import fixtures
88

99

10+
@pytest.mark.skip_kernel # pins model http_path to a UC cluster; SEA/kernel is warehouse-only
1011
@pytest.mark.skip_profile("databricks_cluster")
1112
class TestIncrementalColumnTags(RerunSafeMixin, MaterializationV2Mixin):
1213
@pytest.fixture(scope="class")
@@ -57,6 +58,7 @@ def test_changing_column_tags(self, project):
5758

5859

5960
@pytest.mark.python
61+
@pytest.mark.skip_kernel # Python model runs on a cluster; kernel/SEA is warehouse-only
6062
@pytest.mark.skip_profile("databricks_cluster")
6163
class TestIncrementalPythonColumnTags(TestIncrementalColumnTags):
6264
@pytest.fixture(scope="class")

tests/functional/adapter/incremental/test_incremental_strategies.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,7 @@ def test_incremental(self, project):
166166

167167

168168
# Only runs under SQL warehouse profile, but overrides compute at model level
169+
@pytest.mark.skip_kernel # overrides compute to a UC cluster; SEA/kernel is warehouse-only
169170
@pytest.mark.skip_profile("databricks_uc_cluster", "databricks_cluster")
170171
class TestInsertOverwriteWithModelComputeOverride(IncrementalBase):
171172
@pytest.fixture(scope="class")

tests/functional/adapter/python_model/test_python_model.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,7 @@ def test_changing_schema_via_incremental(self, project):
182182

183183

184184
@pytest.mark.python
185+
@pytest.mark.skip_kernel # pins model http_path to a cluster; SEA/kernel is warehouse-only
185186
@pytest.mark.skip_profile("databricks_cluster", "databricks_uc_cluster")
186187
@pytest.mark.flaky(reruns=2, reruns_delay=120)
187188
class TestSpecifyingHttpPath(PythonModelDataMixin, BasePythonModelTests):

tests/functional/adapter/spog/test_spog_connection_open.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ def test_probe_failure_run_succeeds(self, project):
3232
run_dbt(["run", "--select", "spog_smoke_model"], expect_pass=True)
3333

3434

35+
@pytest.mark.skip_kernel # pins model to a UC cluster; SEA/kernel is warehouse-only
3536
@pytest.mark.skip_profile("databricks_cluster", "databricks_uc_cluster")
3637
class TestSpogNamedCompute:
3738
@pytest.fixture(scope="class")

tests/profiles.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,12 @@ def _build_databricks_cluster_target(
5050
# If you are specifying a port for running tests, assume Docker
5151
# is being used and disable TLS verification
5252
connection_parameters["_tls_no_verify"] = True
53+
# Route the whole suite through the connector's Rust kernel backend (SEA
54+
# transport) when DBT_DATABRICKS_USE_KERNEL=1. Used by the weekend
55+
# integration-kernel workflow. SEA is warehouse-only, so this is only
56+
# meaningful on the databricks_uc_sql_endpoint profile.
57+
if os.getenv("DBT_DATABRICKS_USE_KERNEL") == "1":
58+
connection_parameters["use_kernel"] = True
5359
profile["connection_parameters"] = connection_parameters
5460
return profile
5561

0 commit comments

Comments
 (0)