Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 22 additions & 0 deletions dbt/adapters/databricks/impl.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import json
import posixpath
import re
import threading
from abc import ABC, abstractmethod
from collections import defaultdict
from collections.abc import Iterable, Iterator
Expand Down Expand Up @@ -32,6 +33,7 @@
from dbt_common.behavior_flags import BehaviorFlag
from dbt_common.contracts.config.base import BaseConfig, MergeBehavior
from dbt_common.exceptions import DbtConfigError, DbtInternalError, DbtRuntimeError
from dbt_common.invocation import get_invocation_id
from dbt_common.record import auto_record_function, record_function
from dbt_common.utils import executor
from dbt_common.utils.dict import AttrDict
Expand Down Expand Up @@ -303,6 +305,10 @@ def __init__(self, config: Any, mp_context: SpawnContext) -> None:
"has_dbr_capability": self._has_dbr_capability_parse,
}

# State for `claim_first_batch_operation` (first-microbatch-batch gating).
self._first_batch_lock = threading.Lock()
self._first_batch_claims: set[tuple[str, str, str]] = set()

def _has_dbr_capability_parse(self, capability_name: str) -> bool:
"""Parse-time stub: True only on SQL warehouses for warehouse-supported capabilities."""
creds = self.config.credentials
Expand Down Expand Up @@ -1128,6 +1134,22 @@ def is_cluster(self) -> bool:
"""Check if the current connection is a cluster."""
return self.connections.is_cluster()

@available
def claim_first_batch_operation(self, relation_name: str, operation: str) -> bool:
"""Return True only for the first caller of (invocation, relation, operation), else False.

dbt-core runs the first microbatch batch alone before the parallel ones, so the first
caller here is that batch — letting callers confine per-model metadata writes (CLUSTER BY,
SET TBLPROPERTIES) to it and avoid colliding ALTERs. Thread-safe shared state, so a Python
method not a macro (Fusion parity owed). Keyed on invocation id so `dbt retry` re-claims.
"""
key = (get_invocation_id(), relation_name, operation)
with self._first_batch_lock:
if key in self._first_batch_claims:
return False
self._first_batch_claims.add(key)
return True

@available.parse(lambda *a, **k: {})
def clean_sql(self, sql: str) -> str:
return SqlUtils.clean_sql(sql)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,10 @@

{{ run_pre_hooks() }}

{#-- Confine per-model config work to the first batch (dbt-core runs it alone before the
parallel batches) so concurrent batches' metadata ALTERs don't collide. --#}
{%- set is_first_batch = not model.batch or adapter.claim_first_batch_operation(target_relation.render(), 'config_changes') -%}

{% call statement('main', language=language) %}
{{ get_create_intermediate_table(intermediate_relation, compiled_code, language) }}
{% endcall %}
Expand Down Expand Up @@ -59,7 +63,9 @@
{%- endif -%}
{#-- Relation must be merged --#}
{%- do process_schema_changes(on_schema_change, intermediate_relation, existing_relation) -%}
{{ process_config_changes(target_relation, existing_relation) }}
{%- if is_first_batch -%}
{{ process_config_changes(target_relation, existing_relation) }}
{%- endif -%}
{% set build_sql = get_build_sql(incremental_strategy, target_relation, intermediate_relation) %}
{%- if language == 'sql' -%}
{#-- Check if build_sql is a list (multi-statement strategy) or a string (single statement) --#}
Expand Down Expand Up @@ -100,6 +106,11 @@

{#-- Run pre-hooks --#}
{{ run_hooks(pre_hooks) }}

{#-- Confine per-model config work to the first batch (dbt-core runs it alone before the
parallel batches) so concurrent batches' metadata ALTERs don't collide. --#}
{%- set is_first_batch = not model.batch or adapter.claim_first_batch_operation(target_relation.render(), 'config_changes') -%}

{#-- Incremental run logic --#}
{%- if existing_relation is none -%}
{#-- Relation must be created --#}
Expand Down Expand Up @@ -142,7 +153,7 @@
{%- endif -%}
{#-- Relation must be merged --#}
{%- set _configuration_changes = none -%}
{%- if config.get('incremental_apply_config_changes', True) | as_bool -%}
{%- if is_first_batch and config.get('incremental_apply_config_changes', True) | as_bool -%}
{%- set model_config = adapter.get_config_from_model(config.model) -%}
{%- set _existing_config = adapter.get_relation_config(existing_relation, model_config) -%}
{%- set _configuration_changes = model_config.get_changeset(_existing_config) -%}
Expand Down Expand Up @@ -212,7 +223,9 @@
{{ apply_constraints(target_relation, constraints) }}
{% endif %}
{%- endif -%}
{% do persist_docs(target_relation, model, for_relation=True) %}
{%- if is_first_batch -%}
{% do persist_docs(target_relation, model, for_relation=True) %}
{%- endif -%}
{%- endif -%}

{% set should_revoke = should_revoke(existing_relation, full_refresh_mode) %}
Expand Down
9 changes: 8 additions & 1 deletion docs/flow/incremental_flow.md
Original file line number Diff line number Diff line change
@@ -1,12 +1,19 @@
# Incremental Flow

_Last updated: 2026-08-10_
_Last updated: 2026-08-12_

> Two diagrams follow: **Existing** is the default path, **New** is used when the
> `use_materialization_v2` behavior flag is enabled. See [flow/README.md](README.md) for what the
> flag is and how the selection works. Source:
> `dbt/include/databricks/macros/materializations/incremental/incremental.sql`.

> **Concurrent microbatch:** dbt-core runs this materialization once per batch. The
> configuration-change application (`CLUSTER BY`, `SET TBLPROPERTIES`, tags, comments, masks,
> constraints — and, in the Existing path, `persist_docs`) is claimed by the first batch via
> `adapter.claim_first_batch_operation`, which dbt-core runs alone before the parallel batches, so
> those metadata statements run once instead of colliding with the other batches' concurrent
> writes. The incremental write and `optimize` still run on every batch.

## Existing Incremental Flow

```mermaid
Expand Down
31 changes: 31 additions & 0 deletions tests/functional/adapter/microbatch/fixtures.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,3 +56,34 @@
- name: event_time
- name: amount
"""

# Five days so a re-run produces first + parallel-middle + last batches.
concurrent_input_model_sql = """
{{ config(materialized='table', event_time='event_time') }}
select 1 as id, TIMESTAMP '2020-01-01 00:00:00-0' as event_time
union all
select 2 as id, TIMESTAMP '2020-01-02 00:00:00-0' as event_time
union all
select 3 as id, TIMESTAMP '2020-01-03 00:00:00-0' as event_time
union all
select 4 as id, TIMESTAMP '2020-01-04 00:00:00-0' as event_time
union all
select 5 as id, TIMESTAMP '2020-01-05 00:00:00-0' as event_time
"""

# Carries both #1443 collision triggers: liquid_clustered_by (CLUSTER BY) and tblproperties
# (SET TBLPROPERTIES), which must run on the first batch only under concurrency.
concurrent_microbatch_model_sql = """
{{ config(
materialized='incremental',
incremental_strategy='microbatch',
unique_key='id',
event_time='event_time',
batch_size='day',
begin=modules.datetime.datetime(2020, 1, 1, 0, 0, 0),
concurrent_batches=true,
liquid_clustered_by=['id'],
tblproperties={'delta.columnMapping.mode': 'name'}
) }}
select * from {{ ref('concurrent_input_model') }}
"""
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
from importlib import metadata

import pytest
from dbt.tests import util
from packaging import version

from tests.functional.adapter.microbatch import fixtures

try:
from dbt.tests.util import patch_microbatch_end_time
except ImportError:
from freezegun import freeze_time as patch_microbatch_end_time

dbt_version = metadata.version("dbt-core")


@pytest.mark.skipif(
version.parse(dbt_version) < version.parse("1.9.0b1"),
reason="Microbatch is not supported with this version of core",
)
@pytest.mark.skip_profile("databricks_cluster")
class TestConcurrentMicrobatchConfigChanges:
"""Concurrent microbatch: config changes (CLUSTER BY, SET TBLPROPERTIES) must run on the
first batch only, else they collide with concurrent batch writes. See issue #1443."""

@pytest.fixture(scope="class")
def project_config_update(self):
return {"flags": {"use_concurrent_microbatch": True}}

@pytest.fixture(scope="class")
def models(self):
return {
"concurrent_input_model.sql": fixtures.concurrent_input_model_sql,
"concurrent_microbatch_model.sql": fixtures.concurrent_microbatch_model_sql,
}

def test_all_batches_succeed_with_config_changes(self, project):
# Backfill: relation doesn't exist yet, so batches run sequentially.
with patch_microbatch_end_time("2020-01-05 13:57:00"):
util.run_dbt(["run"])

# Re-run: middle batches now run in parallel; each must land despite the config changes.
with patch_microbatch_end_time("2020-01-05 13:57:00"):
util.run_dbt(["run", "--select", "concurrent_microbatch_model"])

rows = project.run_sql(
"select count(*) from "
f"{project.database}.{project.test_schema}.concurrent_microbatch_model",
fetch="all",
)
assert rows[0][0] == 5

properties = project.run_sql(
"show tblproperties "
f"{project.database}.{project.test_schema}.concurrent_microbatch_model",
fetch="all",
)
prop = {row[0]: row[1] for row in properties}
assert prop.get("delta.columnMapping.mode") == "name"
assert "id" in prop.get("clusteringColumns", "")
52 changes: 52 additions & 0 deletions tests/unit/test_adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -2051,3 +2051,55 @@ def test_no_op_on_empty_host(self):
DatabricksAdapter.debug_emit_spog_block(host="", http_path="")
mock_probe.assert_not_called()
mock_logger.info.assert_not_called()


class TestClaimFirstBatchOperation(DatabricksAdapterBase):
"""`claim_first_batch_operation` confines per-model microbatch work to the first batch."""

@pytest.fixture
def adapter(self, setUp) -> DatabricksAdapter:
return DatabricksAdapter(self._get_config(), get_context("spawn"))

def test_first_claim_wins_and_rest_lose(self, adapter):
assert adapter.claim_first_batch_operation("`cat`.`sch`.`tbl`", "config_changes") is True
assert adapter.claim_first_batch_operation("`cat`.`sch`.`tbl`", "config_changes") is False
assert adapter.claim_first_batch_operation("`cat`.`sch`.`tbl`", "config_changes") is False

def test_claims_are_independent_per_operation(self, adapter):
assert adapter.claim_first_batch_operation("`cat`.`sch`.`tbl`", "config_changes") is True
assert adapter.claim_first_batch_operation("`cat`.`sch`.`tbl`", "optimize") is True
assert adapter.claim_first_batch_operation("`cat`.`sch`.`tbl`", "optimize") is False

def test_claims_are_independent_per_relation(self, adapter):
assert adapter.claim_first_batch_operation("`cat`.`sch`.`tbl`", "config_changes") is True
assert adapter.claim_first_batch_operation("`cat`.`sch`.`other`", "config_changes") is True

def test_claim_resets_across_invocations(self, adapter):
from dbt_common.invocation import reset_invocation_id

assert adapter.claim_first_batch_operation("`cat`.`sch`.`tbl`", "config_changes") is True
assert adapter.claim_first_batch_operation("`cat`.`sch`.`tbl`", "config_changes") is False
# `dbt retry` runs under a fresh invocation id, which must re-open the claim.
reset_invocation_id()
assert adapter.claim_first_batch_operation("`cat`.`sch`.`tbl`", "config_changes") is True

def test_exactly_one_winner_under_concurrency(self, adapter):
import threading

results: list[bool] = []
results_lock = threading.Lock()
start = threading.Barrier(20)

def worker():
start.wait()
won = adapter.claim_first_batch_operation("`cat`.`sch`.`tbl`", "config_changes")
with results_lock:
results.append(won)

threads = [threading.Thread(target=worker) for _ in range(20)]
for t in threads:
t.start()
for t in threads:
t.join()

assert sum(results) == 1
Loading