diff --git a/dlt/destinations/impl/databricks/databricks.py b/dlt/destinations/impl/databricks/databricks.py index d4d015eb11..d1d3d74fe9 100644 --- a/dlt/destinations/impl/databricks/databricks.py +++ b/dlt/destinations/impl/databricks/databricks.py @@ -486,6 +486,8 @@ def _ensure_arrow_record_batch(self, batch: pyarrow.RecordBatch) -> pyarrow.Reco class DatabricksClient(SqlJobClientWithStagingDataset, SupportsStagingDestination): + SCHEMA_UPDATE_MAX_ATTEMPTS: ClassVar[int] = 10 + def __init__( self, schema: Schema, @@ -672,6 +674,21 @@ def _get_constraints_sql( return "" + def _should_retry_schema_update(self, ex: Exception) -> bool: + # delta rejects losers of concurrent metadata commits, the statement did not apply + msg = str(ex) + return "DELTA_METADATA_CHANGED" in msg or "DELTA_CONCURRENT" in msg + + def _is_schema_update_applied(self, ex: Exception) -> bool: + # a concurrent load added the same column first + msg = str(ex) + return "FIELDS_ALREADY_EXISTS" in msg or "COLUMN_ALREADY_EXISTS" in msg + + def _make_create_table(self, qualified_name: str, table: PreparedTableSchema) -> str: + if self.capabilities.supports_create_table_if_not_exists: + return f"CREATE TABLE IF NOT EXISTS {qualified_name}" + return f"CREATE TABLE {qualified_name}" + def _get_table_update_sql( self, table_name: str, new_columns: Sequence[TColumnSchema], generate_alter: bool ) -> List[str]: @@ -755,6 +772,9 @@ def _get_table_update_sql( ) # Note: DELTA is the default format, no explicit USING clause needed + comment = table.get(TABLE_COMMENT_HINT) or table.get("description") + escaped_comment = escape_databricks_literal(comment) if comment else None + # For CREATE TABLE, we need custom generation if we have any custom clauses or non-DELTA format if not generate_alter and ( cluster_clause or partition_clause or tblproperties_clause or using_clause @@ -778,6 +798,10 @@ def _get_table_update_sql( # Add CLUSTER BY clause if cluster_clause: sql += f" {cluster_clause}" + # comment is inlined: a standalone COMMENT ON TABLE is a separate metadata write + # that conflicts with concurrent loads initializing the same schema + if escaped_comment: + sql += f" COMMENT {escaped_comment}" # Add TBLPROPERTIES clause (comes after CLUSTER BY) if tblproperties_clause: sql += f" {tblproperties_clause}" @@ -792,11 +816,12 @@ def _get_table_update_sql( qualified_name = self.sql_client.make_qualified_table_name(table_name) sql_result.append(f"ALTER TABLE {qualified_name} {cluster_clause}") + if not generate_alter and escaped_comment: + sql_result[0] += f" COMMENT {escaped_comment}" + qualified_name = self.sql_client.make_qualified_table_name(table_name) - if table.get(TABLE_COMMENT_HINT) or table.get("description"): - comment = table.get(TABLE_COMMENT_HINT) or table.get("description") - escaped_comment = escape_databricks_literal(comment) + if generate_alter and escaped_comment: sql_result.append(f"COMMENT ON TABLE {qualified_name} IS {escaped_comment}") if table.get(TABLE_TAGS_HINT): diff --git a/dlt/destinations/impl/databricks/sql_client.py b/dlt/destinations/impl/databricks/sql_client.py index 8f473b0565..0c5ac2f58c 100644 --- a/dlt/destinations/impl/databricks/sql_client.py +++ b/dlt/destinations/impl/databricks/sql_client.py @@ -105,6 +105,9 @@ def rollback_transaction(self) -> None: def native_connection(self) -> "DatabricksSqlConnection": return self._conn + def create_dataset(self) -> None: + self.execute_sql("CREATE SCHEMA IF NOT EXISTS %s" % self.fully_qualified_dataset_name()) + def drop_tables(self, *tables: str) -> None: # Tables are drop with `IF EXISTS`, but databricks raises when the schema doesn't exist. # Multi statement exec is safe and the error can be ignored since all tables are in the same schema. diff --git a/dlt/destinations/job_client_impl.py b/dlt/destinations/job_client_impl.py index 10f58e150b..321e242046 100644 --- a/dlt/destinations/job_client_impl.py +++ b/dlt/destinations/job_client_impl.py @@ -2,7 +2,9 @@ from abc import abstractmethod import base64 import contextlib +import random from copy import copy +from time import sleep from types import TracebackType from typing import ( Any, @@ -68,7 +70,7 @@ CredentialsConfiguration, ) -from dlt.destinations.exceptions import DatabaseUndefinedRelation +from dlt.destinations.exceptions import DatabaseException, DatabaseUndefinedRelation from dlt.destinations.job_impl import ( ReferenceFollowupJobRequest, ) @@ -247,6 +249,9 @@ def __init__( class SqlJobClientBase(WithSqlClient, JobClientBase, WithStateSync): + SCHEMA_UPDATE_MAX_ATTEMPTS: ClassVar[int] = 1 + """Attempts to converge the destination schema when concurrent loads modify it, 1 disables retries""" + def __init__( self, schema: Schema, @@ -649,20 +654,84 @@ def _get_storage_table_query_columns(self) -> List[str]: fields += ["numeric_precision", "numeric_scale"] return fields + def _should_retry_schema_update(self, ex: Exception) -> bool: + """Tells if a schema update statement that failed with `ex` did not apply due to a + concurrent schema modification and may converge when retried from re-read storage. + Destinations where concurrent loads modify table metadata without ddl transactions + detect their conflict errors here. + """ + return False + + def _is_schema_update_applied(self, ex: Exception) -> bool: + """Tells if a schema update statement that failed with `ex` has its effect already + present in the destination i.e. a concurrent load added the same column first. + """ + return False + + def _execute_schema_update_with_tolerance( + self, sql_scripts: Sequence[str] + ) -> Tuple[List[str], List[Exception]]: + """Executes statements one by one so a conflicting statement does not abort the rest. + Returns statements that failed on a concurrent schema modification with their errors. + """ + failed: List[str] = [] + errors: List[Exception] = [] + for sql in sql_scripts: + try: + self.sql_client.execute_sql(sql) + except DatabaseException as ex: + if self._is_schema_update_applied(ex): + continue + if not self._should_retry_schema_update(ex): + raise + failed.append(sql) + errors.append(ex) + return failed, errors + def _execute_schema_update_sql( self, only_tables: Iterable[str], store_schema: bool = True ) -> TSchemaTables: # Only `only_tables` are included, or all if None. - sql_scripts, schema_update = self._build_schema_update_sql( - list(self.get_storage_tables(only_tables or self.schema.tables.keys())) - ) - # Stay within max query size when doing DDL. - # Some DB backends use bytes not characters, so decrease the limit by half, - # assuming most of the characters in DDL encoded into single bytes. - self.sql_client.execute_many(sql_scripts) - # skip writing the version row when the schema is already stored (enforced update) + table_names = list(only_tables or self.schema.tables.keys()) + attempts = max(1, self.SCHEMA_UPDATE_MAX_ATTEMPTS) + schema_update: TSchemaTables = {} + pending: List[str] = [] + prev_round: Optional[Tuple[Tuple[str, ...], Tuple[str, ...]]] = None + for attempt in range(attempts): + sql_scripts, update = self._build_schema_update_sql( + list(self.get_storage_tables(table_names)) + ) + schema_update = schema_update or update + if attempts == 1: + self.sql_client.execute_many(sql_scripts) + break + # replay failed statements the rebuilt diff cannot regenerate + sql_scripts.extend(sql for sql in pending if sql not in sql_scripts) + if not sql_scripts: + break + pending, errors = self._execute_schema_update_with_tolerance(sql_scripts) + if pending: + this_round = (tuple(sql_scripts), tuple(str(ex) for ex in errors)) + # identical statements failing identically twice is a persistent error, not a + # concurrent modification + if this_round == prev_round or attempt == attempts - 1: + raise errors[0] + prev_round = this_round + logger.warning( + f"{len(pending)} schema update statements failed on a possible concurrent" + f" schema modification, retrying ({attempt + 1}/{attempts}): {errors[0]}" + ) + # add jitter to prevent lockstep in concurrent loads + sleep(random.uniform(0.05, 0.25 * (attempt + 1))) if store_schema: - self._update_schema_in_storage(self.schema) + # under concurrency an identical schema may have been stored by another load + if attempts > 1 and self.get_stored_schema_by_hash(self.schema.stored_version_hash): + logger.info( + f"Schema with hash {self.schema.stored_version_hash} was stored by a" + " concurrent load, skipping version row" + ) + else: + self._update_schema_in_storage(self.schema) return schema_update def _build_schema_update_sql( diff --git a/tests/load/databricks/test_databricks_concurrent_schema.py b/tests/load/databricks/test_databricks_concurrent_schema.py new file mode 100644 index 0000000000..5494248e5d --- /dev/null +++ b/tests/load/databricks/test_databricks_concurrent_schema.py @@ -0,0 +1,341 @@ +from tests.utils import skip_if_not_active + +skip_if_not_active("databricks") + +import threading +from typing import Any, Callable, Dict, Generator, Iterable, List, Tuple, cast +from unittest import mock + +import pytest + +from dlt.common.schema import Schema +from dlt.common.schema.typing import TColumnSchema, TTableSchemaColumns +from dlt.common.schema.utils import new_table +from dlt.common.utils import uniq_id +from dlt.destinations.exceptions import DatabaseTerminalException +from dlt.destinations.impl.databricks.databricks import DatabricksClient + +from tests.load.databricks.test_databricks_table_builder import create_client +from tests.load.utils import empty_schema, yield_client + +pytestmark = pytest.mark.essential + +TABLE_NAME = "event_test_table" +COL_A: TTableSchemaColumns = {"col_a": {"name": "col_a", "data_type": "text"}} +ALL_COLS: TTableSchemaColumns = { + **COL_A, + "col_b": {"name": "col_b", "data_type": "bigint"}, +} + + +def _mocked_client(empty_schema: Schema) -> DatabricksClient: + client = create_client(empty_schema, create_indexes=False) + client.schema.update_table( + new_table( + TABLE_NAME, + columns=[ + {"name": "col_a", "data_type": "text"}, + {"name": "col_b", "data_type": "bigint"}, + ], + ) + ) + return client + + +def test_schema_update_converges_under_concurrent_modification(empty_schema: Schema) -> None: + client = _mocked_client(empty_schema) + # a concurrent load created the table with col_a only, then another one conflicts on ALTER + client.get_storage_tables = mock.MagicMock( # type: ignore[method-assign] + side_effect=[ + [(TABLE_NAME, {})], + [(TABLE_NAME, dict(COL_A))], + [(TABLE_NAME, dict(COL_A))], + [(TABLE_NAME, dict(ALL_COLS))], + ] + ) + calls: List[str] = [] + + def execute_sql(sql: str, *args: Any) -> None: + calls.append("execute:" + sql) + if len(calls) == 2: + raise DatabaseTerminalException(Exception("[DELTA_METADATA_CHANGED] concurrent update")) + + client.sql_client.execute_sql = mock.MagicMock(side_effect=execute_sql) # type: ignore[method-assign] + client.get_stored_schema_by_hash = mock.MagicMock(return_value=None) # type: ignore[method-assign] + client._update_schema_in_storage = mock.MagicMock( # type: ignore[method-assign] + side_effect=lambda schema: calls.append("store") + ) + + client._execute_schema_update_sql([TABLE_NAME]) + + executed = [c for c in calls if c.startswith("execute")] + assert len(executed) == 3 + assert executed[0].startswith("execute:CREATE TABLE IF NOT EXISTS") + assert "ADD COLUMN" in executed[1] and "`col_b`" in executed[1] + assert "ADD COLUMN" in executed[2] and "`col_b`" in executed[2] + assert calls[-1] == "store" + + +def test_schema_update_replays_dropped_hint_statements(empty_schema: Schema) -> None: + """A comment failing after its ALTER applied is replayed even though the rebuilt diff no + longer generates statements for the table.""" + client = create_client(empty_schema, create_indexes=False) + table = new_table( + TABLE_NAME, + columns=[ + {"name": "col_a", "data_type": "text"}, + {"name": "col_b", "data_type": "bigint"}, + ], + ) + table["description"] = "test table comment" + client.schema.update_table(table) + client.get_storage_tables = mock.MagicMock( # type: ignore[method-assign] + side_effect=[ + [(TABLE_NAME, dict(COL_A))], + [(TABLE_NAME, dict(ALL_COLS))], + [(TABLE_NAME, dict(ALL_COLS))], + ] + ) + executed: List[str] = [] + + def execute_sql(sql: str, *args: Any) -> None: + executed.append(sql) + if sql.startswith("COMMENT ON TABLE") and executed.count(sql) == 1: + raise DatabaseTerminalException(Exception("[DELTA_METADATA_CHANGED] concurrent update")) + + client.sql_client.execute_sql = mock.MagicMock(side_effect=execute_sql) # type: ignore[method-assign] + client.get_stored_schema_by_hash = mock.MagicMock(return_value=None) # type: ignore[method-assign] + client._update_schema_in_storage = mock.MagicMock() # type: ignore[method-assign] + + client._execute_schema_update_sql([TABLE_NAME]) + + comments = [sql for sql in executed if sql.startswith("COMMENT ON TABLE")] + assert len(comments) == 2 + client._update_schema_in_storage.assert_called_once() + + +def test_schema_update_tolerates_column_added_by_concurrent_load(empty_schema: Schema) -> None: + client = _mocked_client(empty_schema) + client.get_storage_tables = mock.MagicMock( # type: ignore[method-assign] + side_effect=[ + [(TABLE_NAME, dict(COL_A))], + [(TABLE_NAME, dict(ALL_COLS))], + ] + ) + client.sql_client.execute_sql = mock.MagicMock( # type: ignore[method-assign] + side_effect=DatabaseTerminalException(Exception("[FIELDS_ALREADY_EXISTS] col_b")) + ) + client.get_stored_schema_by_hash = mock.MagicMock(return_value=None) # type: ignore[method-assign] + client._update_schema_in_storage = mock.MagicMock() # type: ignore[method-assign] + + client._execute_schema_update_sql([TABLE_NAME]) + + assert client.sql_client.execute_sql.call_count == 1 + client._update_schema_in_storage.assert_called_once() + + +def test_schema_update_raises_on_other_errors(empty_schema: Schema) -> None: + client = _mocked_client(empty_schema) + client.get_storage_tables = mock.MagicMock(return_value=[(TABLE_NAME, {})]) # type: ignore[method-assign] + client.sql_client.execute_sql = mock.MagicMock( # type: ignore[method-assign] + side_effect=DatabaseTerminalException(Exception("[PARSE_SYNTAX_ERROR] boom")) + ) + client._update_schema_in_storage = mock.MagicMock() # type: ignore[method-assign] + + with pytest.raises(DatabaseTerminalException): + client._execute_schema_update_sql([TABLE_NAME]) + + assert client.sql_client.execute_sql.call_count == 1 + client._update_schema_in_storage.assert_not_called() + + +def test_schema_update_bails_on_persistent_error(empty_schema: Schema) -> None: + """Identical statements failing identically twice raise without burning all attempts.""" + client = _mocked_client(empty_schema) + client.get_storage_tables = mock.MagicMock( # type: ignore[method-assign] + side_effect=lambda table_names: [(TABLE_NAME, {})] + ) + client.sql_client.execute_sql = mock.MagicMock( # type: ignore[method-assign] + side_effect=DatabaseTerminalException(Exception("[DELTA_METADATA_CHANGED] static failure")) + ) + client._update_schema_in_storage = mock.MagicMock() # type: ignore[method-assign] + + with pytest.raises(DatabaseTerminalException): + client._execute_schema_update_sql([TABLE_NAME]) + + assert client.sql_client.execute_sql.call_count == 2 + client._update_schema_in_storage.assert_not_called() + + +def test_schema_update_exhausts_attempts( + empty_schema: Schema, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(DatabricksClient, "SCHEMA_UPDATE_MAX_ATTEMPTS", 3) + client = _mocked_client(empty_schema) + client.get_storage_tables = mock.MagicMock( # type: ignore[method-assign] + side_effect=lambda table_names: [(TABLE_NAME, {})] + ) + counter = iter(range(100)) + + def execute_sql(sql: str, *args: Any) -> None: + # distinct conflicting commits so the no-progress bail does not trigger + raise DatabaseTerminalException( + Exception(f"[DELTA_METADATA_CHANGED] commit {next(counter)}") + ) + + client.sql_client.execute_sql = mock.MagicMock(side_effect=execute_sql) # type: ignore[method-assign] + client._update_schema_in_storage = mock.MagicMock() # type: ignore[method-assign] + + with pytest.raises(DatabaseTerminalException): + client._execute_schema_update_sql([TABLE_NAME]) + + assert client.sql_client.execute_sql.call_count == 3 + client._update_schema_in_storage.assert_not_called() + + +def test_schema_update_skips_version_row_stored_by_concurrent_load(empty_schema: Schema) -> None: + client = _mocked_client(empty_schema) + client.get_storage_tables = mock.MagicMock(return_value=[(TABLE_NAME, dict(ALL_COLS))]) # type: ignore[method-assign] + client.sql_client.execute_sql = mock.MagicMock() # type: ignore[method-assign] + client.get_stored_schema_by_hash = mock.MagicMock(return_value=object()) # type: ignore[method-assign] + client._update_schema_in_storage = mock.MagicMock() # type: ignore[method-assign] + + client._execute_schema_update_sql([TABLE_NAME]) + + client.sql_client.execute_sql.assert_not_called() + client._update_schema_in_storage.assert_not_called() + + +N_WORKERS = 4 + + +TClientGen = Generator[DatabricksClient, None, None] + + +def _new_client(dataset_name: str) -> Tuple[DatabricksClient, TClientGen]: + # connection is opened later, inside the worker thread + gen = cast( + TClientGen, yield_client("databricks", dataset_name=dataset_name, enter_client=False) + ) + return next(gen), gen + + +def _sync_after_storage_read(client: DatabricksClient, barrier: threading.Barrier) -> None: + original = client.get_storage_tables + state = {"synced": False} + + def wrapper(table_names: Iterable[str]) -> Iterable[Tuple[str, TTableSchemaColumns]]: + result = list(original(table_names)) + if not state["synced"]: + state["synced"] = True + try: + barrier.wait(timeout=120) + except threading.BrokenBarrierError: + pass + return result + + client.get_storage_tables = wrapper # type: ignore[method-assign] + + +def _run_concurrent_update( + clients_gens: List[Tuple[DatabricksClient, TClientGen]], + mutate: Callable[[Schema, int], Any], +) -> Dict[int, BaseException]: + barrier = threading.Barrier(len(clients_gens)) + errors: Dict[int, BaseException] = {} + + def run(i: int, client: DatabricksClient) -> None: + try: + with client: + mutate(client.schema, i) + client.schema._bump_version() + _sync_after_storage_read(client, barrier) + client.update_stored_schema() + except Exception as e: + errors[i] = e + + threads = [threading.Thread(target=run, args=(i, c)) for i, (c, _) in enumerate(clients_gens)] + for t in threads: + t.start() + for t in threads: + t.join() + return errors + + +def _drop_dataset(dataset_name: str) -> None: + client, gen = _new_client(dataset_name) + try: + with client: + client.sql_client.drop_dataset() + except Exception: + pass + finally: + gen.close() + + +def test_concurrent_create_table_converges() -> None: + dataset = "conc_create_" + uniq_id() + raced = "raced_" + uniq_id() + columns: List[TColumnSchema] = [ + {"name": "id", "data_type": "bigint"}, + {"name": "val", "data_type": "text"}, + ] + + setup, setup_gen = _new_client(dataset) + with setup: + setup.initialize_storage() + setup.update_stored_schema() + setup_gen.close() + + clients_gens = [_new_client(dataset) for _ in range(N_WORKERS)] + try: + errors = _run_concurrent_update( + clients_gens, lambda schema, i: schema.update_table(new_table(raced, columns=columns)) + ) + assert errors == {}, "concurrent CREATE raced: " + str( + {i: str(e)[:140] for i, e in errors.items()} + ) + verify, verify_gen = _new_client(dataset) + with verify: + _, storage_columns = list(verify.get_storage_tables([raced]))[0] + verify_gen.close() + assert {"id", "val"} <= set(storage_columns) + finally: + for _, gen in clients_gens: + gen.close() + _drop_dataset(dataset) + + +def test_concurrent_divergent_add_column_converges() -> None: + dataset = "conc_addcol_" + uniq_id() + raced = "raced_" + uniq_id() + + setup, setup_gen = _new_client(dataset) + with setup: + setup.initialize_storage() + setup.schema.update_table(new_table(raced, columns=[{"name": "id", "data_type": "bigint"}])) + setup.schema._bump_version() + setup.update_stored_schema() + setup_gen.close() + + clients_gens = [_new_client(dataset) for _ in range(N_WORKERS)] + try: + errors = _run_concurrent_update( + clients_gens, + lambda schema, i: schema.update_table( + new_table(raced, columns=[{"name": f"col_{i}", "data_type": "bigint"}]) + ), + ) + assert errors == {}, "divergent ADD COLUMN raced: " + str( + {i: str(e)[:140] for i, e in errors.items()} + ) + verify, verify_gen = _new_client(dataset) + with verify: + _, storage_columns = list(verify.get_storage_tables([raced]))[0] + verify_gen.close() + expected = {"id"} | {f"col_{i}" for i in range(N_WORKERS)} + assert expected <= set(storage_columns) + finally: + for _, gen in clients_gens: + gen.close() + _drop_dataset(dataset) diff --git a/tests/load/databricks/test_databricks_table_builder.py b/tests/load/databricks/test_databricks_table_builder.py index b86c4d676d..1a9bab1087 100644 --- a/tests/load/databricks/test_databricks_table_builder.py +++ b/tests/load/databricks/test_databricks_table_builder.py @@ -83,3 +83,35 @@ def test_primary_key_constraint_conditional_on_create_indexes( assert "PRIMARY KEY (`id`)" in sql else: assert "PRIMARY KEY" not in sql + + +@pytest.mark.parametrize("with_cluster", [False, True], ids=["plain", "clustered"]) +def test_create_table_uses_if_not_exists(empty_schema: Schema, with_cluster: bool) -> None: + # generate_alter=False takes both the base and the custom-clause CREATE paths + columns: List[TColumnSchema] = [ + {"name": "col_a", "data_type": "text", "cluster": with_cluster}, + {"name": "col_b", "data_type": "bigint"}, + ] + client = create_client(empty_schema, create_indexes=False) + + sql = client._get_table_update_sql("event_test_table", columns, generate_alter=False)[0] + + assert sql.startswith("CREATE TABLE IF NOT EXISTS") + + +@pytest.mark.parametrize("with_cluster", [False, True], ids=["plain", "clustered"]) +def test_table_comment_generation(empty_schema: Schema, with_cluster: bool) -> None: + table_name = "event_test_table" + columns: List[TColumnSchema] = [{"name": "col_a", "data_type": "text", "cluster": with_cluster}] + client = create_client(empty_schema, create_indexes=False) + table = new_table(table_name, columns=columns) + table["description"] = "test table comment" + client.schema.update_table(table) + + create_sql = client._get_table_update_sql(table_name, columns, generate_alter=False) + assert len(create_sql) == 1 + assert "COMMENT 'test table comment'" in create_sql[0] + assert "COMMENT ON TABLE" not in create_sql[0] + + alter_sql = client._get_table_update_sql(table_name, columns, generate_alter=True) + assert any(s.startswith("COMMENT ON TABLE") for s in alter_sql)