diff --git a/dbt/adapters/databricks/impl.py b/dbt/adapters/databricks/impl.py index 6640f0d76..474168c16 100644 --- a/dbt/adapters/databricks/impl.py +++ b/dbt/adapters/databricks/impl.py @@ -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 @@ -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 @@ -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 @@ -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) diff --git a/dbt/include/databricks/macros/materializations/incremental/incremental.sql b/dbt/include/databricks/macros/materializations/incremental/incremental.sql index e79859822..cc4ec4cfd 100644 --- a/dbt/include/databricks/macros/materializations/incremental/incremental.sql +++ b/dbt/include/databricks/macros/materializations/incremental/incremental.sql @@ -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 %} @@ -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) --#} @@ -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 --#} @@ -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) -%} @@ -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) %} diff --git a/docs/flow/incremental_flow.md b/docs/flow/incremental_flow.md index a255948e5..67f875a4d 100644 --- a/docs/flow/incremental_flow.md +++ b/docs/flow/incremental_flow.md @@ -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 diff --git a/tests/functional/adapter/microbatch/fixtures.py b/tests/functional/adapter/microbatch/fixtures.py index c09ca9bf3..479a518dc 100644 --- a/tests/functional/adapter/microbatch/fixtures.py +++ b/tests/functional/adapter/microbatch/fixtures.py @@ -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') }} +""" diff --git a/tests/functional/adapter/microbatch/test_microbatch_concurrent.py b/tests/functional/adapter/microbatch/test_microbatch_concurrent.py new file mode 100644 index 000000000..3d7a2adb0 --- /dev/null +++ b/tests/functional/adapter/microbatch/test_microbatch_concurrent.py @@ -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", "") diff --git a/tests/unit/test_adapter.py b/tests/unit/test_adapter.py index 92520bca8..4df033569 100644 --- a/tests/unit/test_adapter.py +++ b/tests/unit/test_adapter.py @@ -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