feat: optimize dbt seed in session mode via DataFrame API - #1433
Closed
alexeyegorov wants to merge 26 commits into
Closed
feat: optimize dbt seed in session mode via DataFrame API#1433alexeyegorov wants to merge 26 commits into
alexeyegorov wants to merge 26 commits into
Conversation
Integration tests require azure-prod environment secrets that are not configured in this fork. Restrict to manual workflow_dispatch only. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
DAT-1062: Update github workflows
GitHub Actions runners are blocked by the Databricks org IP allowlist when calling `gh api repos/databricks/dbt-databricks/releases/latest`. - Replace gh api call with PyPI JSON API (curl + python3) to get latest dbt-databricks version — no auth or IP restrictions - Remove upstream release notes fetch (same IP issue) and all UPSTREAM_NOTES references from outputs, env, and PR body Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> ci(upstream-sync): use gh api instead of gh pr create Replace gh pr create with gh api REST call to avoid the GraphQL repository.parent lookup that gets blocked by the Databricks org IP allowlist on GitHub Actions runners. Labels are applied via a separate issues/labels API call since the pulls API doesn't accept labels on creation. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> ci(upstream-sync): restore upstream release notes in sync PR Fetch release notes via unauthenticated curl to the GitHub REST API to avoid the Databricks org IP allowlist that blocks authenticated gh api calls. Falls back gracefully if notes aren't available. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> ci(upstream-sync): replace repo variable with file for sync tag tracking GITHUB_TOKEN cannot write repo variables (requires PAT). Replace vars.LAST_UPSTREAM_SYNC_TAG with .github/last-upstream-sync-tag file that gets committed back to main after a clean sync. The workflow already has contents:write permission so this works without extra secrets. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> ci(upstream-sync): push tag file update to sync branch, not main The update step was doing git push origin HEAD:main which pushed all commits directly to main. Now it checks out the sync branch and pushes there, so the tag update is included in the PR. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Repo code already includes v1.11.6 content but tag file was stale at v1.11.5, which would cause a redundant no-op sync PR. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
The upstream repo uses `linux-ubuntu-latest` (a custom Databricks self-hosted runner label). Our fork doesn't have those runners, so CI jobs queue indefinitely. Switch to `ubuntu-latest`. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
feat: Fix upstream-sync
- Introduced `DatabricksSessionHandle` and `SessionCursorWrapper` to enable SparkSession-based execution. - Updated `DatabricksConnectionManager` to handle session mode connections and capabilities. - Enhanced `DatabricksCredentials` to auto-detect and validate connection methods. - Added session mode handling in Python model submission and execution. - Implemented cleanup for temporary views to prevent state leakage between models. This update allows dbt to run entirely within a single SparkSession on Databricks job clusters, improving execution efficiency and compatibility. fix: add validate_creds() call to SessionPythonJobHelper Match the pattern in BaseDatabricksHelper.__init__ for consistent early validation of credentials in session mode. fix: is_cluster() and _connection_keys_session() behavior in session mode - is_cluster() now returns True in session mode since session mode runs on a job cluster; prevents incorrect SQL Warehouse code paths in macros - _connection_keys_session() now respects with_aliases parameter and uses "database" (the actual field name) vs "catalog" (the alias) correctly, matching the DBSQL branch pattern Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Introduced comprehensive unit tests for session mode components, including `SessionCursorWrapper`, `DatabricksSessionHandle`, and session mode credentials. - Enhanced test coverage for session mode auto-detection and validation in `DatabricksCredentials`. - Implemented tests for session mode Python model submission and execution, ensuring proper handling of temporary views and execution errors. These additions improve the reliability and robustness of session mode features in the Databricks adapter. Add higher test coverage for the new session mode behaviour - test_session.py: update list_schemas/list_tables assertions to expect backtick-quoted identifiers (matching _quote_identifier() output) - test_connection_manager.py: mock is_session_mode=False so is_cluster() tests exercise the http_path logic, not the session mode short-circuit - escape underscore in list_schemas LIKE pattern assertion - cover _validate_session_mode error paths - cover unique_field and _connection_keys_session in session mode - cover session cursor cancel, dbr_version fallbacks, closed handle guard - cover DatabricksSessionHandle cancel, close, rollback, __del__ - add connection manager session mode tests - refactor TestSubmitPythonJobSessionMode to use fixture - bypass BehaviorFlag descriptor in adapter fixture: Use object.__setattr__ to set the behavior attribute directly on the adapter instance, bypassing the BehaviorFlag class-level descriptor that raises CompilationError when assigned via normal attribute access. - Use PropertyMock to patch behavior at class level, bypassing the BehaviorFlag data descriptor that raises CompilationError on __get__ - Mock connections with a string name to satisfy log_code_execution protobuf serialization - Return a real AdapterResponse from the patched super() call to satisfy CodeExecutionStatus logging event serialization All 799 unit tests now pass (3 previously failing). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
SessionCursorWrapper.execute() now renders bindings into the SQL string instead of rejecting them. This enables dbt seed in session mode, where SparkSession.sql() only accepts fully-formed SQL strings. Type rendering matches DBSQL connector behavior: - None → NULL - bool → TRUE/FALSE (checked before int) - str → single-quoted with quote escaping - Decimal → float (matches SqlUtils.translate_bindings) - int/float → bare numeric literals Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Update README to explain this is an enhanced fork of dbt-databricks, not a copy. Add session mode documentation, versioning scheme, and updated installation instructions. Rename package to dbt-databricks-enhanced in pyproject.toml with updated URLs. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
feat: implement session mode support for Databricks connections
- add timeout and thread-based execution to SessionPythonSubmitter - pass model timeout config to SessionPythonSubmitter - extract CANCEL_GRACE_PERIOD constant and add review comments - add job group isolation tests for SessionPythonSubmitter - add exception propagation test for SessionPythonSubmitter - Document setJobGroup race window in worker thread - Document mock cancellation limitation in test Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
fix(session): add timeout and cancellation to Python model execution
Single-service devcontainer using mcr.microsoft.com/devcontainers/python:3.10 with hatch + uv + Databricks CLI pre-installed. HATCH_DATA_DIR pinned to workspace for predictable interpreter path; named Docker volumes persist hatch envs and uv/pip caches across rebuilds. Includes VS Code extensions and settings matching docs/dbt-databricks-dev.md guidance. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Documents how enhancements, hotfixes, upstream syncs, and shared docs/CI propagate between main and X.Y.latest branches, and calls out the backward-hotfix rule (ship on X.Y.latest then merge back to main). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
docs: release strategy, CI workflow, and devcontainer
Maps dbt SQL type strings to PySpark DataType objects. Handles string, bigint, int, double, float, boolean, date, timestamp, decimal(p,s). Unknown types fall back to StringType. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Uses spark.createDataFrame() + df.write.mode('overwrite').insertInto()
to bypass SQL string rendering for seed loading in session mode.
Converts Decimal to float matching existing SqlUtils behavior.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Thin @available methods exposed to Jinja macros. is_session_mode() delegates to connection manager. load_seed_data() extracts table/column info and delegates to session handle. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
When is_session_mode() is true, databricks__load_csv_rows bypasses INSERT SQL generation and delegates to adapter.load_seed_data() which uses spark.createDataFrame() + insertInto(). DBSQL mode unchanged. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Guard against None from convert_type() (falls back to StringType) - Handle varchar(N) with length parameter (e.g. varchar(255)) - Add time type mapping to StringType (no PySpark TimeType exists) - Add tests for all three edge cases Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
alexeyegorov
requested review from
benc-db,
jprakash-db,
sd-db and
tejassp-db
as code owners
April 30, 2026 14:01
Author
|
Opened against wrong repo by mistake. Recreating against viadkt/dbt-databricks-enhanced. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
spark.createDataFrame()+df.write.mode("overwrite").insertInto()instead_build_spark_schematype mapping (dbt SQL types → PySpark types) with full coverage including edge cases (None, varchar(N), time, decimal(p,s))dbt seedexperience, no config changes. DBSQL mode completely untouched.Architecture
Files changed
session.pyload_seed_data(),_build_spark_schema(),_dbt_type_to_spark_type()— bulk of logic, 100% custom codeimpl.py@availablewrappers:is_session_mode(),load_seed_data()helpers.sqldatabricks__load_csv_rowstest_session.pyTest plan
dbt seedon Databricks job cluster in session mode🤖 Generated with Claude Code