diff --git a/providers/databricks/docs/operators/delta_table_sensor.rst b/providers/databricks/docs/operators/delta_table_sensor.rst new file mode 100644 index 0000000000000..12aea773f63db --- /dev/null +++ b/providers/databricks/docs/operators/delta_table_sensor.rst @@ -0,0 +1,100 @@ +.. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _howto/operator:DatabricksDeltaTableVersionSensor: + +DatabricksDeltaTableVersionSensor +================================= + +Use the :class:`~airflow.providers.databricks.sensors.databricks_delta_table.DatabricksDeltaTableVersionSensor` +to monitor a Delta table in Databricks and detect new commit versions. When a newer version is observed, +the sensor succeeds, pushes provenance metadata to XCom, and optionally emits an Airflow Asset event. + + +Using the Sensor +---------------- + +The sensor requires a target Delta table and a Databricks SQL endpoint. + +To identify the table, specify either: + +* ``table_name`` - String representing the 1-, 2-, or 3-part Delta table name (``table``, ``schema.table``, or ``catalog.schema.table``). +* ``unity_table`` - An instance of :class:`~airflow.providers.databricks.assets.databricks.UnityTableIdentity`. When provided without explicit ``outlets``, the sensor automatically registers the corresponding Unity Catalog asset (``databricks://///``) as an outlet. + +To connect to Databricks SQL, specify either ``sql_warehouse_name`` or ``http_path``. + +Other parameters, such as connection credentials and timeout configuration, are documented in the class reference. + + +Waiting for a new version +------------------------- + +By default, omitting ``baseline_version`` and ``target_version`` enrolls the table by capturing its current version on the first check and waiting for subsequent commits. You can also specify an explicit ``baseline_version`` to succeed when the table's version is strictly greater: + +.. code-block:: python + + from airflow.providers.databricks.sensors.databricks_delta_table import ( + DatabricksDeltaTableVersionSensor, + ) + + wait_for_updates = DatabricksDeltaTableVersionSensor( + task_id="wait_for_table_updates", + table_name="main.default.customers", + sql_warehouse_name="my_warehouse", + baseline_version=42, + poke_interval=30, + ) + + +Waiting for a specific version +------------------------------ + +To wait for an exact target Delta version, set ``target_version``: + +.. code-block:: python + + from airflow.providers.databricks.sensors.databricks_delta_table import ( + DatabricksDeltaTableVersionSensor, + ) + + wait_for_v50 = DatabricksDeltaTableVersionSensor( + task_id="wait_for_v50", + table_name="main.default.customers", + sql_warehouse_name="my_warehouse", + target_version=50, + ) + + +Using the observed version downstream +------------------------------------- + +When a new version is detected, the sensor pushes metadata to XCom under the key ``delta_table_version``, including ``version``, ``timestamp``, ``operation``, and ``table_identity``. If multiple commits occurred between polling checks, the sensor observes and reports the latest commit. + +Downstream tasks requiring reproducible reads should pin their queries to the observed version using ``VERSION AS OF``: + +.. code-block:: sql + + SELECT * FROM main.default.customers + VERSION AS OF {{ ti.xcom_pull(task_ids="wait_for_table_updates", key="delta_table_version")["version"] }} + +When an outlet asset is configured or inferred from ``unity_table``, the sensor also attaches the observed version and table identity to the asset event extra metadata. + + +Deferrable mode +--------------- + +To avoid consuming worker slots while polling, run the sensor in deferrable mode by setting ``deferrable=True`` (or globally via ``operators.default_deferrable``). The sensor yields execution to the triggerer until a newer version is observed. diff --git a/providers/databricks/provider.yaml b/providers/databricks/provider.yaml index 2b66f546bf2a6..2c777e3ca681d 100644 --- a/providers/databricks/provider.yaml +++ b/providers/databricks/provider.yaml @@ -131,6 +131,7 @@ integrations: - /docs/apache-airflow-providers-databricks/operators/sql_statements.rst - /docs/apache-airflow-providers-databricks/operators/warehouse.rst - /docs/apache-airflow-providers-databricks/operators/copy_into.rst + - /docs/apache-airflow-providers-databricks/operators/delta_table_sensor.rst tags: [service] - integration-name: Databricks Repos external-doc-url: https://docs.databricks.com/repos/index.html @@ -187,6 +188,7 @@ triggers: - integration-name: Databricks python-modules: - airflow.providers.databricks.triggers.databricks + - airflow.providers.databricks.triggers.databricks_delta_table sensors: - integration-name: Databricks @@ -194,6 +196,7 @@ sensors: - airflow.providers.databricks.sensors.databricks - airflow.providers.databricks.sensors.databricks_sql - airflow.providers.databricks.sensors.databricks_partition + - airflow.providers.databricks.sensors.databricks_delta_table connection-types: - hook-class-name: airflow.providers.databricks.hooks.databricks.DatabricksHook diff --git a/providers/databricks/src/airflow/providers/databricks/get_provider_info.py b/providers/databricks/src/airflow/providers/databricks/get_provider_info.py index a218ce3546417..e2f4c85fee40f 100644 --- a/providers/databricks/src/airflow/providers/databricks/get_provider_info.py +++ b/providers/databricks/src/airflow/providers/databricks/get_provider_info.py @@ -121,7 +121,10 @@ def get_provider_info(): "triggers": [ { "integration-name": "Databricks", - "python-modules": ["airflow.providers.databricks.triggers.databricks"], + "python-modules": [ + "airflow.providers.databricks.triggers.databricks", + "airflow.providers.databricks.triggers.databricks_delta_table", + ], } ], "sensors": [ @@ -131,6 +134,7 @@ def get_provider_info(): "airflow.providers.databricks.sensors.databricks", "airflow.providers.databricks.sensors.databricks_sql", "airflow.providers.databricks.sensors.databricks_partition", + "airflow.providers.databricks.sensors.databricks_delta_table", ], } ], diff --git a/providers/databricks/src/airflow/providers/databricks/sensors/databricks_delta_table.py b/providers/databricks/src/airflow/providers/databricks/sensors/databricks_delta_table.py new file mode 100644 index 0000000000000..6a5600cce53e3 --- /dev/null +++ b/providers/databricks/src/airflow/providers/databricks/sensors/databricks_delta_table.py @@ -0,0 +1,419 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import time +from collections.abc import Sequence +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.compat.sdk import ( + AirflowFailException, + AirflowSensorTimeout, + BaseSensorOperator, + conf, +) +from airflow.providers.common.sql.hooks.handlers import fetch_all_handler +from airflow.providers.databricks.hooks.databricks_sql import DatabricksSqlHook +from airflow.providers.databricks.triggers.databricks_delta_table import DatabricksDeltaTableVersionTrigger +from airflow.providers.databricks.utils.query_tags import build_query_tags + +if TYPE_CHECKING: + from pydantic import JsonValue + + from airflow.providers.common.compat.sdk import Context + from airflow.providers.databricks.assets.databricks import UnityTableIdentity + + +class DatabricksDeltaTableVersionSensor(BaseSensorOperator): + """ + Sensor to observe Databricks Delta table commits and detect version changes. + + Monitors a Delta table by polling its commit history via ``DESCRIBE HISTORY``. + Succeeds when a commit version newer than ``baseline_version`` is detected, + or when ``target_version`` is reached. + + If ``baseline_version`` is not specified and ``target_version`` is None, the sensor + captures the current version on its first poke as the baseline and waits for a newer + commit on subsequent evaluations. + + :param table_name: Name of the Delta table (e.g. ``catalog.schema.table``, ``schema.table``, or ``table``). + :param unity_table: Optional static :class:`~airflow.providers.databricks.assets.databricks.UnityTableIdentity`. + If specified and ``outlets`` is not provided, defaults the outlet asset to ``unity_table.to_asset()``. + :param baseline_version: Baseline Delta version to observe changes against. + :param target_version: Optional explicit target Delta version to wait for. + :param allow_recreation: Whether to treat a lower version (table dropped and recreated) as a change. + Defaults to True. + :param databricks_conn_id: Reference to Databricks connection ID, defaults to DatabricksSqlHook.default_conn_name. + :param sql_warehouse_name: Optional name of Databricks SQL warehouse. + :param http_path: Optional HTTP path of Databricks SQL warehouse or cluster. + :param catalog: Initial catalog name, defaults to "". + :param schema: Initial schema name, defaults to "default". + :param session_configuration: Optional dictionary of Spark session parameters. + :param http_headers: Optional list of (key, value) pairs for HTTP headers. + :param client_parameters: Additional parameters passed to the Databricks SQL connector. + :param query_tags: Optional dictionary of query tags. + :param include_airflow_query_tags: Whether to attach Airflow metadata query tags. Defaults to True. + :param deferrable: Whether to run the sensor in deferrable mode. + """ + + template_fields: Sequence[str] = ( + "databricks_conn_id", + "catalog", + "schema", + "table_name", + "baseline_version", + "target_version", + "http_headers", + "query_tags", + ) + + def __init__( + self, + *, + table_name: str | None = None, + unity_table: UnityTableIdentity | None = None, + baseline_version: int | None = None, + target_version: int | None = None, + allow_recreation: bool = True, + databricks_conn_id: str = DatabricksSqlHook.default_conn_name, + sql_warehouse_name: str | None = None, + http_path: str | None = None, + catalog: str = "", + schema: str = "default", + session_configuration: dict | None = None, + http_headers: list[tuple[str, str]] | None = None, + client_parameters: dict[str, Any] | None = None, + query_tags: dict[str, str | None] | None = None, + include_airflow_query_tags: bool = True, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + outlets: list | None = None, + **kwargs, + ) -> None: + if table_name is None and unity_table is None: + raise ValueError("One of 'table_name' or 'unity_table' must be provided.") + + self.unity_table = unity_table + if outlets is None and unity_table is not None: + outlets = [unity_table.to_asset()] + + hook_params = kwargs.pop("hook_params", {}) + super().__init__(outlets=outlets, **kwargs) + + self.table_name = table_name + self.baseline_version = baseline_version + self._baseline_version: int | None = None + self._baseline_initialized: bool = False + self.target_version = target_version + self.allow_recreation = allow_recreation + self.databricks_conn_id = databricks_conn_id + self._sql_warehouse_name = sql_warehouse_name + self._http_path = http_path + self.catalog = catalog + self.schema = schema + self.session_config = session_configuration + self.session_configuration = session_configuration + self.http_headers = http_headers + self.client_parameters = client_parameters or {} + self.hook_params = hook_params + self.query_tags = query_tags or {} + self.include_airflow_query_tags = include_airflow_query_tags + self.deferrable = deferrable + self.caller = "DatabricksDeltaTableVersionSensor" + + def _ensure_baseline_initialized(self) -> None: + """Initialize runtime baseline version from templated baseline_version if not yet initialized.""" + if not self._baseline_initialized: + if self.baseline_version is not None: + self._baseline_version = int(self.baseline_version) + else: + self._baseline_version = None + self._baseline_initialized = True + + @cached_property + def hook(self) -> DatabricksSqlHook: + return DatabricksSqlHook( + self.databricks_conn_id, + self._http_path, + self._sql_warehouse_name, + self.session_config, + self.http_headers, + self.catalog, + self.schema, + caller=self.caller, + **self.client_parameters, + **self.hook_params, + ) + + def _validate_workspace_host(self) -> None: + """Validate connection host matches unity_table host when unity_table is specified.""" + if self.unity_table is not None: + if (self.hook.host or "").lower() != self.unity_table.host: + raise ValueError( + f"Databricks connection host {self.hook.host!r} does not match " + f"unity_table host {self.unity_table.host!r}." + ) + + def _resolve_table_name(self) -> str: + """Resolve and validate the effective table name.""" + if self.unity_table is not None: + unity_canonical = f"{self.unity_table.catalog}.{self.unity_table.schema}.{self.unity_table.table}" + if self.table_name: + parts = [p.lower() for p in self.table_name.split(".")] + if len(parts) == 3: + if (parts[0], parts[1], parts[2]) != ( + self.unity_table.catalog, + self.unity_table.schema, + self.unity_table.table, + ): + raise ValueError( + f"table_name {self.table_name!r} does not match unity_table {unity_canonical!r}." + ) + elif len(parts) == 2: + if (parts[0], parts[1]) != (self.unity_table.schema, self.unity_table.table): + raise ValueError( + f"table_name {self.table_name!r} does not match unity_table {unity_canonical!r}." + ) + elif len(parts) == 1: + if parts[0] != self.unity_table.table: + raise ValueError( + f"table_name {self.table_name!r} does not match unity_table {unity_canonical!r}." + ) + else: + raise ValueError(f"Invalid table_name format {self.table_name!r}.") + return unity_canonical + + if not self.table_name: + raise ValueError("table_name cannot be empty.") + return self.table_name + + def _get_latest_version(self, table_name: str) -> tuple[int, str | None, str | None]: + sql = f"DESCRIBE HISTORY {table_name} LIMIT 1" + try: + result = self.hook.run(sql, handler=fetch_all_handler) + except Exception as e: + if DatabricksDeltaTableVersionTrigger._is_permanent_error(e): + raise AirflowFailException( + f"Failed to fetch Delta table history for '{table_name}'. " + "Ensure the table exists, is a Delta table, and the caller has read permissions: " + f"{e}" + ) from e + raise + + if not isinstance(result, Sequence) or not result: + raise AirflowFailException(f"Delta table '{table_name}' returned empty history.") + + row = result[0] + if not isinstance(row, Sequence) or not row: + raise AirflowFailException(f"Delta table '{table_name}' returned malformed history.") + + version_raw = row[0] + if not isinstance(version_raw, (int, str)): + raise AirflowFailException( + f"Delta table '{table_name}' returned invalid version {version_raw!r}." + ) + version = int(version_raw) + timestamp = str(row[1]) if len(row) > 1 and row[1] is not None else None + operation = str(row[4]) if len(row) > 4 and row[4] is not None else None + return version, timestamp, operation + + def _evaluate_version(self, current_version: int, table_name: str) -> bool: + """Evaluate if the current version meets sensor criteria.""" + if self._baseline_version is None and self.target_version is None: + self._baseline_version = current_version + self.log.info( + "Initialized baseline version to %s for table '%s'; waiting for subsequent commits.", + current_version, + table_name, + ) + return False + + if self.target_version is not None: + if current_version >= self.target_version: + self.log.info( + "Target version %s reached for table '%s' (current version: %s).", + self.target_version, + table_name, + current_version, + ) + return True + return False + + if self._baseline_version is not None: + if current_version > self._baseline_version: + self.log.info( + "Observed new Delta version %s (baseline: %s) for table '%s'.", + current_version, + self._baseline_version, + table_name, + ) + return True + if current_version < self._baseline_version: + if self.allow_recreation: + self.log.info( + "Observed table recreation: current version %s is less than baseline %s for table '%s'.", + current_version, + self._baseline_version, + table_name, + ) + return True + raise AirflowFailException( + f"Delta table '{table_name}' was recreated: current version {current_version} " + f"is less than baseline version {self._baseline_version}." + ) + + return False + + def _push_provenance( + self, + context: Context, + current_version: int, + table_name: str, + timestamp: str | None, + operation: str | None, + ) -> dict[str, Any]: + table_identity: dict[str, JsonValue] + if self.unity_table: + table_identity = { + "host": self.unity_table.host, + "catalog": self.unity_table.catalog, + "schema": self.unity_table.schema, + "table": self.unity_table.table, + } + else: + parts = table_name.split(".") + if len(parts) == 3: + table_identity = { + "catalog": parts[0], + "schema": parts[1], + "table": parts[2], + } + elif len(parts) == 2: + table_identity = { + "schema": parts[0], + "table": parts[1], + } + else: + table_identity = {"table": table_name} + + provenance = { + "version": current_version, + "baseline_version": self._baseline_version, + "table_name": table_name, + "table_identity": table_identity, + "timestamp": timestamp, + "operation": operation, + } + if self.do_xcom_push and context is not None and "ti" in context and context["ti"]: + context["ti"].xcom_push(key="delta_table_version", value=provenance) + + # Emit asset event metadata if outlet_events is supported in context. + # As established on #74196 and #74195, the task-success asset event means + # "a newer Delta version was observed," with version + table identity in extra. + # It is not a vague "table refreshed" and not a promise that every intermediate + # commit was delivered. Authors needing pinned reads must VERSION AS OF that observed version. + if context is not None and "outlet_events" in context: + extra: dict[str, JsonValue] = { + "version": current_version, + "observed_version": current_version, + "table_identity": table_identity, + "operation": operation, + } + if timestamp is not None: + extra["timestamp"] = timestamp + for outlet in self.outlets: + context["outlet_events"][outlet].extra = extra + return provenance + + def poke(self, context: Context) -> bool: + self._ensure_baseline_initialized() + self._validate_workspace_host() + self.hook.query_tags = build_query_tags(context, self.query_tags, self.include_airflow_query_tags) + table_name = self._resolve_table_name() + current_version, timestamp, operation = self._get_latest_version(table_name) + + if self._evaluate_version(current_version, table_name): + self._push_provenance(context, current_version, table_name, timestamp, operation) + return True + return False + + def execute(self, context: Context) -> Any: + self._ensure_baseline_initialized() + self._validate_workspace_host() + table_name = self._resolve_table_name() + if not self.deferrable: + return super().execute(context) + + query_tags = build_query_tags(context, self.query_tags, self.include_airflow_query_tags) + + # If baseline is not provided and target_version is not set, capture baseline synchronously first + if self._baseline_version is None and self.target_version is None: + self.hook.query_tags = query_tags + current_version, _, _ = self._get_latest_version(table_name) + self._baseline_version = current_version + self.log.info( + "Initialized baseline version to %s for table '%s' before deferring.", + current_version, + table_name, + ) + + end_time = time.time() + self.timeout + self.defer( + trigger=DatabricksDeltaTableVersionTrigger( + table_name=table_name, + databricks_conn_id=self.databricks_conn_id, + baseline_version=self._baseline_version, + target_version=self.target_version, + allow_recreation=self.allow_recreation, + sql_warehouse_name=self._sql_warehouse_name, + http_path=self._http_path, + catalog=self.catalog, + schema=self.schema, + session_configuration=self.session_config, + http_headers=self.http_headers, + client_parameters=self.client_parameters, + hook_params=self.hook_params, + query_tags=query_tags, + polling_period_seconds=int(self.poke_interval), + end_time=end_time, + caller=self.caller, + ), + method_name="execute_complete", + ) + + def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> Any: + self._ensure_baseline_initialized() + if not event: + raise AirflowFailException("Trigger did not return an event.") + if event.get("status") == "success": + table_name = event.get("table_name") or self._resolve_table_name() + version = event["version"] + if "baseline_version" in event: + self._baseline_version = event["baseline_version"] + self._push_provenance( + context, + version, + table_name, + event.get("timestamp"), + event.get("operation"), + ) + self.log.info("Successfully detected Delta table '%s' version: %s.", table_name, version) + return version + if event.get("status") == "timeout": + raise AirflowSensorTimeout(event.get("message", "Sensor timed out waiting for Delta table version.")) + raise AirflowFailException(event.get("message", "Sensor deferred execution failed.")) diff --git a/providers/databricks/src/airflow/providers/databricks/triggers/databricks_delta_table.py b/providers/databricks/src/airflow/providers/databricks/triggers/databricks_delta_table.py new file mode 100644 index 0000000000000..ee2baeea53ef1 --- /dev/null +++ b/providers/databricks/src/airflow/providers/databricks/triggers/databricks_delta_table.py @@ -0,0 +1,251 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import asyncio +import time +from collections.abc import Sequence +from typing import Any + +from asgiref.sync import sync_to_async + +from airflow.providers.common.sql.hooks.handlers import fetch_all_handler +from airflow.providers.databricks.hooks.databricks_sql import DatabricksSqlHook +from airflow.triggers.base import BaseTrigger, TriggerEvent + + +class DatabricksDeltaTableVersionTrigger(BaseTrigger): + """ + Asynchronously poll a Databricks Delta table until a newer committed version is observed. + + :param table_name: Fully qualified or local Delta table name. + :param databricks_conn_id: Reference to Databricks connection ID. + :param baseline_version: The baseline Delta version to compare against. + :param target_version: Optional explicit target Delta version. + :param allow_recreation: Whether to treat a lower version (table recreation) as a version change. + :param sql_warehouse_name: Optional name of Databricks SQL warehouse. + :param http_path: Optional HTTP path of Databricks SQL warehouse. + :param catalog: Initial catalog name. + :param schema: Initial schema name. + :param polling_period_seconds: Polling interval in seconds. + :param end_time: Absolute epoch timestamp when the trigger must timeout. + :param caller: Name of the caller for hook tracing. + """ + + def __init__( + self, + table_name: str, + databricks_conn_id: str, + baseline_version: int | None = None, + target_version: int | None = None, + allow_recreation: bool = True, + sql_warehouse_name: str | None = None, + http_path: str | None = None, + catalog: str = "", + schema: str = "default", + session_configuration: dict[str, str] | None = None, + http_headers: list[tuple[str, str]] | None = None, + client_parameters: dict[str, Any] | None = None, + hook_params: dict[str, Any] | None = None, + query_tags: dict[str, str | None] | None = None, + polling_period_seconds: int = 30, + end_time: float | None = None, + caller: str = "DatabricksDeltaTableVersionTrigger", + ) -> None: + super().__init__() + self.table_name = table_name + self.databricks_conn_id = databricks_conn_id + self.baseline_version = baseline_version + self._baseline_version = baseline_version + self.target_version = target_version + self.allow_recreation = allow_recreation + self.sql_warehouse_name = sql_warehouse_name + self.http_path = http_path + self.catalog = catalog + self.schema = schema + self.session_configuration = session_configuration + self.http_headers = http_headers + self.client_parameters = client_parameters or {} + self.hook_params = hook_params or {} + self.query_tags = query_tags + self.polling_period_seconds = polling_period_seconds + self.end_time = end_time if end_time is not None else (time.time() + 3600) + self.caller = caller + + def serialize(self) -> tuple[str, dict[str, Any]]: + return ( + "airflow.providers.databricks.triggers.databricks_delta_table.DatabricksDeltaTableVersionTrigger", + { + "table_name": self.table_name, + "databricks_conn_id": self.databricks_conn_id, + "baseline_version": self._baseline_version, + "target_version": self.target_version, + "allow_recreation": self.allow_recreation, + "sql_warehouse_name": self.sql_warehouse_name, + "http_path": self.http_path, + "catalog": self.catalog, + "schema": self.schema, + "session_configuration": self.session_configuration, + "http_headers": self.http_headers, + "client_parameters": self.client_parameters, + "hook_params": self.hook_params, + "query_tags": self.query_tags, + "polling_period_seconds": self.polling_period_seconds, + "end_time": self.end_time, + "caller": self.caller, + }, + ) + + def _get_hook(self) -> DatabricksSqlHook: + return DatabricksSqlHook( + databricks_conn_id=self.databricks_conn_id, + http_path=self.http_path, + sql_endpoint_name=self.sql_warehouse_name, + session_configuration=self.session_configuration, + http_headers=self.http_headers, + catalog=self.catalog, + schema=self.schema, + caller=self.caller, + query_tags=self.query_tags, + **self.client_parameters, + **self.hook_params, + ) + + def _get_version(self) -> tuple[int | None, str | None, str | None]: + hook = self._get_hook() + sql = f"DESCRIBE HISTORY {self.table_name} LIMIT 1" + result = hook.run(sql, handler=fetch_all_handler) + if not isinstance(result, Sequence) or not result: + return None, None, None + row = result[0] + if not isinstance(row, Sequence) or not row: + return None, None, None + version_raw = row[0] + if not isinstance(version_raw, (int, str)): + return None, None, None + version = int(version_raw) + timestamp = str(row[1]) if len(row) > 1 and row[1] is not None else None + operation = str(row[4]) if len(row) > 4 and row[4] is not None else None + return version, timestamp, operation + + @staticmethod + def _is_permanent_error(exc: Exception) -> bool: + """ + Check if an error querying Delta table version is permanent. + + Permanent failures include missing table/view errors and authorization/permission failures. + Transient failures (e.g. connection drops, network timeouts) should be retried. + """ + if isinstance(exc, PermissionError): + return True + msg = str(exc).upper() + permanent_patterns = ( + "TABLE_OR_VIEW_NOT_FOUND", + "TABLE OR VIEW NOT FOUND", + "PERMISSION_DENIED", + "PERMISSION DENIED", + "ACCESS_DENIED", + "ACCESS DENIED", + "DOES NOT EXIST", + "NOT FOUND", + "UNAUTHORIZED", + "FORBIDDEN", + "ANALYSISEXCEPTION", + "PARSE_SYNTAX_ERROR", + ) + return any(pattern in msg for pattern in permanent_patterns) + + async def run(self): + while time.time() < self.end_time: + try: + version, timestamp, operation = await sync_to_async(self._get_version)() + if version is not None: + if self._baseline_version is None and self.target_version is None: + self._baseline_version = version + self.log.info( + "Initialized baseline version to %s for table %s; waiting for newer version.", + version, + self.table_name, + ) + else: + satisfied = False + if self.target_version is not None: + satisfied = version >= self.target_version + elif self._baseline_version is not None: + if version > self._baseline_version: + satisfied = True + elif version < self._baseline_version: + if self.allow_recreation: + self.log.info( + "Detected table recreation for %s: current version %s is less than baseline %s", + self.table_name, + version, + self._baseline_version, + ) + satisfied = True + else: + yield TriggerEvent( + { + "status": "error", + "message": ( + f"Table '{self.table_name}' was recreated: current version " + f"{version} is less than baseline version {self._baseline_version}" + ), + } + ) + return + + if satisfied: + yield TriggerEvent( + { + "status": "success", + "version": version, + "baseline_version": self._baseline_version, + "table_name": self.table_name, + "timestamp": timestamp, + "operation": operation, + } + ) + return + except Exception as e: + if self._is_permanent_error(e): + yield TriggerEvent( + { + "status": "error", + "message": ( + f"Failed to query Delta table '{self.table_name}' version due to " + f"permanent error: {e}" + ), + } + ) + return + self.log.warning( + "Transient error querying Delta table version for %s: %s", self.table_name, e + ) + + now = time.time() + if now >= self.end_time: + break + await asyncio.sleep(min(self.polling_period_seconds, max(0.1, self.end_time - now))) + + yield TriggerEvent( + { + "status": "timeout", + "message": f"Timed out waiting for new Delta table version on '{self.table_name}'.", + } + ) diff --git a/providers/databricks/tests/unit/databricks/sensors/test_databricks_delta_table.py b/providers/databricks/tests/unit/databricks/sensors/test_databricks_delta_table.py new file mode 100644 index 0000000000000..b56634c0f8e08 --- /dev/null +++ b/providers/databricks/tests/unit/databricks/sensors/test_databricks_delta_table.py @@ -0,0 +1,496 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from unittest import mock +from unittest.mock import MagicMock, patch + +import pytest + +from airflow.exceptions import TaskDeferred +from airflow.providers.common.compat.assets import Asset +from airflow.providers.common.compat.sdk import AirflowException +from airflow.providers.databricks.assets.databricks import UnityTableIdentity +from airflow.providers.databricks.sensors.databricks_delta_table import DatabricksDeltaTableVersionSensor +from airflow.providers.databricks.triggers.databricks_delta_table import DatabricksDeltaTableVersionTrigger + +TASK_ID = "test-delta-table-sensor" +DEFAULT_CONN_ID = "databricks_default" +TABLE_NAME = "main.default.users" +UNITY_TABLE = UnityTableIdentity( + host="https://my-workspace.cloud.databricks.com/", catalog="main", schema="default", table="users" +) +ASSET_URI = "databricks://my-workspace.cloud.databricks.com/main/default/users" + + +def _make_history_row(version=5, timestamp="2026-10-05 12:00:00", operation="WRITE"): + return [[version, timestamp, "user_id", "user_name", operation]] + + +class TestDatabricksDeltaTableVersionSensor: + def test_init_validation(self): + with pytest.raises(ValueError, match="One of 'table_name' or 'unity_table' must be provided"): + DatabricksDeltaTableVersionSensor(task_id=TASK_ID) + + def test_init_with_unity_table_sets_outlet_asset(self): + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + unity_table=UNITY_TABLE, + ) + assert sensor.outlets == [Asset(uri=ASSET_URI)] + assert sensor._resolve_table_name() == "main.default.users" + + def test_init_with_custom_outlets_preserved(self): + custom_asset = Asset(uri="custom://my-custom-uri") + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + unity_table=UNITY_TABLE, + outlets=[custom_asset], + ) + assert sensor.outlets == [custom_asset] + + @pytest.mark.parametrize( + ("table_name", "valid"), + [ + ("main.default.users", True), + ("Main.Default.Users", True), + ("MAIN.DEFAULT.USERS", True), + ("default.users", True), + ("Default.Users", True), + ("DEFAULT.USERS", True), + ("users", True), + ("Users", True), + ("USERS", True), + ("other.default.users", False), + ("Other.Default.Users", False), + ("main.other.users", False), + ("main.default.orders", False), + ("a.b.c.d", False), + ], + ) + def test_table_name_validation_against_unity_table(self, table_name, valid): + if valid: + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=table_name, + unity_table=UNITY_TABLE, + ) + assert sensor._resolve_table_name() == "main.default.users" + else: + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=table_name, + unity_table=UNITY_TABLE, + ) + with pytest.raises(ValueError, match="does not match unity_table|Invalid table_name"): + sensor._resolve_table_name() + + @pytest.mark.parametrize( + ("table_name", "catalog", "schema"), + [ + ("MAIN.Default.Users", None, None), + ("Default.Users", "MAIN", None), + ("Users", "MAIN", "Default"), + ], + ) + @patch("airflow.providers.databricks.sensors.databricks_delta_table.DatabricksSqlHook") + def test_sensor_matches_table_case_insensitively(self, mock_hook_cls, table_name, catalog, schema): + mock_hook = mock_hook_cls.return_value + mock_hook.host = UNITY_TABLE.host + mock_hook.run.return_value = _make_history_row(version=6) + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=table_name, + catalog=catalog or "", + schema=schema or "default", + unity_table=UNITY_TABLE, + baseline_version=5, + ) + context = {"ti": MagicMock()} + assert sensor.poke(context) is True + mock_hook.run.assert_called_once_with("DESCRIBE HISTORY main.default.users LIMIT 1", handler=mock.ANY) + + @pytest.mark.parametrize("host", ["other-workspace.cloud.databricks.com", None, ""]) + @patch("airflow.providers.databricks.sensors.databricks_delta_table.DatabricksSqlHook") + def test_sensor_rejects_connection_for_other_workspace_before_sql(self, mock_hook_cls, host): + mock_hook = mock_hook_cls.return_value + mock_hook.host = host + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + unity_table=UNITY_TABLE, + baseline_version=5, + ) + context = {"ti": MagicMock()} + with pytest.raises(ValueError, match="connection host .* does not match unity_table host"): + sensor.poke(context) + assert mock_hook.run.call_args_list == [] + + with pytest.raises(ValueError, match="connection host .* does not match unity_table host"): + sensor.execute(context) + + @pytest.mark.parametrize( + "host", + [ + "my-workspace.cloud.databricks.com", + "My-Workspace.cloud.Databricks.com", + "MY-WORKSPACE.CLOUD.DATABRICKS.COM", + ], + ) + @patch("airflow.providers.databricks.sensors.databricks_delta_table.DatabricksSqlHook") + def test_sensor_matches_workspace_case_insensitively(self, mock_hook_cls, host): + mock_hook = mock_hook_cls.return_value + mock_hook.host = host + mock_hook.run.return_value = _make_history_row(version=6) + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + unity_table=UNITY_TABLE, + baseline_version=5, + ) + context = {"ti": MagicMock()} + assert sensor.poke(context) is True + + @patch("airflow.providers.databricks.sensors.databricks_delta_table.DatabricksSqlHook") + def test_poke_unchanged_version_returns_false(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = _make_history_row(version=5) + + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=TABLE_NAME, + baseline_version=5, + sql_warehouse_name="test_wh", + ) + context = {"ti": MagicMock()} + assert sensor.poke(context) is False + mock_hook.run.assert_called_once_with(f"DESCRIBE HISTORY {TABLE_NAME} LIMIT 1", handler=mock.ANY) + + @patch("airflow.providers.databricks.sensors.databricks_delta_table.DatabricksSqlHook") + def test_poke_newer_version_returns_true_and_pushes_xcom(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = _make_history_row(version=6, operation="MERGE") + + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=TABLE_NAME, + baseline_version=5, + sql_warehouse_name="test_wh", + ) + ti_mock = MagicMock() + context = {"ti": ti_mock} + assert sensor.poke(context) is True + expected_table_identity = {"catalog": "main", "schema": "default", "table": "users"} + ti_mock.xcom_push.assert_called_once_with( + key="delta_table_version", + value={ + "version": 6, + "baseline_version": 5, + "table_name": TABLE_NAME, + "table_identity": expected_table_identity, + "timestamp": "2026-10-05 12:00:00", + "operation": "MERGE", + }, + ) + + @patch("airflow.providers.databricks.sensors.databricks_delta_table.DatabricksSqlHook") + def test_poke_newer_version_emits_outlet_events_with_version_and_table_identity(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.host = UNITY_TABLE.host + mock_hook.run.return_value = _make_history_row(version=7, operation="WRITE") + + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + unity_table=UNITY_TABLE, + baseline_version=5, + sql_warehouse_name="test_wh", + ) + outlet_asset = sensor.outlets[0] + outlet_event_mock = MagicMock() + + class _FakeOutletEvents: + """Minimal accessor double supporting unhashable Asset keys across Airflow versions.""" + + def __init__(self, pairs): + self._pairs = list(pairs) + + def __getitem__(self, key): + for k, v in self._pairs: + if k is key or k == key: + return v + raise KeyError(key) + + context = { + "ti": MagicMock(), + "outlet_events": _FakeOutletEvents([(outlet_asset, outlet_event_mock)]), + } + + assert sensor.poke(context) is True + expected_identity = { + "host": UNITY_TABLE.host, + "catalog": UNITY_TABLE.catalog, + "schema": UNITY_TABLE.schema, + "table": UNITY_TABLE.table, + } + assert outlet_event_mock.extra == { + "version": 7, + "observed_version": 7, + "table_identity": expected_identity, + "operation": "WRITE", + "timestamp": "2026-10-05 12:00:00", + } + + @patch("airflow.providers.databricks.sensors.databricks_delta_table.DatabricksSqlHook") + def test_templated_baseline_not_copied_in_init_and_initialized_on_poke(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.host = UNITY_TABLE.host + mock_hook.run.return_value = _make_history_row(version=10) + + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=TABLE_NAME, + baseline_version="{{ params.baseline }}", + sql_warehouse_name="test_wh", + ) + assert sensor.baseline_version == "{{ params.baseline }}" + assert sensor._baseline_version is None + assert sensor._baseline_initialized is False + + # Rendered by Airflow engine prior to poke/execute + sensor.baseline_version = 8 + + context = {"ti": MagicMock()} + assert sensor.poke(context) is True + assert sensor._baseline_version == 8 + assert sensor._baseline_initialized is True + + @patch("airflow.providers.databricks.sensors.databricks_delta_table.DatabricksSqlHook") + def test_user_baseline_retained_across_pokes(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.host = UNITY_TABLE.host + mock_hook.run.return_value = _make_history_row(version=5) + + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=TABLE_NAME, + baseline_version=5, + sql_warehouse_name="test_wh", + ) + context = {"ti": MagicMock()} + assert sensor.poke(context) is False + assert sensor._baseline_version == 5 + + mock_hook.run.return_value = _make_history_row(version=6) + assert sensor.poke(context) is True + assert sensor._baseline_version == 5 + + @patch("airflow.providers.databricks.sensors.databricks_delta_table.DatabricksSqlHook") + def test_poke_initial_enrollment_captures_baseline(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = _make_history_row(version=10) + + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=TABLE_NAME, + baseline_version=None, + sql_warehouse_name="test_wh", + ) + context = {"ti": MagicMock()} + # First poke captures baseline=10 and returns False + assert sensor.poke(context) is False + assert sensor._baseline_version == 10 + + # Second poke with version 10 returns False + assert sensor.poke(context) is False + + # Third poke with version 11 returns True + mock_hook.run.return_value = _make_history_row(version=11) + assert sensor.poke(context) is True + + @patch("airflow.providers.databricks.sensors.databricks_delta_table.DatabricksSqlHook") + def test_poke_table_recreation_handling(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + # Recreated table resets history to version 0 + mock_hook.run.return_value = _make_history_row(version=0) + + # allow_recreation=True (default) succeeds + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=TABLE_NAME, + baseline_version=15, + allow_recreation=True, + sql_warehouse_name="test_wh", + ) + assert sensor.poke({"ti": MagicMock()}) is True + + # allow_recreation=False raises AirflowException + strict_sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=TABLE_NAME, + baseline_version=15, + allow_recreation=False, + sql_warehouse_name="test_wh", + ) + with pytest.raises(AirflowException, match="was recreated"): + strict_sensor.poke({"ti": MagicMock()}) + + @patch("airflow.providers.databricks.sensors.databricks_delta_table.DatabricksSqlHook") + def test_poke_target_version(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = _make_history_row(version=7) + + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=TABLE_NAME, + target_version=10, + sql_warehouse_name="test_wh", + ) + assert sensor.poke({"ti": MagicMock()}) is False + + mock_hook.run.return_value = _make_history_row(version=10) + assert sensor.poke({"ti": MagicMock()}) is True + + @patch("airflow.providers.databricks.sensors.databricks_delta_table.DatabricksSqlHook") + def test_poke_transient_error_reraised(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.side_effect = ConnectionResetError("Connection reset by peer") + + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=TABLE_NAME, + baseline_version=1, + sql_warehouse_name="test_wh", + ) + with pytest.raises(ConnectionResetError, match="Connection reset by peer"): + sensor.poke({"ti": MagicMock()}) + + @patch("airflow.providers.databricks.sensors.databricks_delta_table.DatabricksSqlHook") + def test_poke_non_delta_or_inaccessible_table_raises(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.side_effect = Exception("TABLE_OR_VIEW_NOT_FOUND") + + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=TABLE_NAME, + baseline_version=1, + sql_warehouse_name="test_wh", + ) + with pytest.raises(AirflowException, match="Failed to fetch Delta table history"): + sensor.poke({"ti": MagicMock()}) + + @patch("airflow.providers.databricks.sensors.databricks_delta_table.DatabricksSqlHook") + def test_deferrable_execution_defers_to_trigger(self, mock_hook_cls): + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=TABLE_NAME, + baseline_version=5, + sql_warehouse_name="test_wh", + session_configuration={"spark.sql.shuffle.partitions": "10"}, + http_headers=[("X-Header", "1")], + client_parameters={"client_param": "abc"}, + hook_params={"hook_param": "xyz"}, + query_tags={"tag": "val"}, + deferrable=True, + ) + context = {"ti": MagicMock(), "dag": MagicMock(), "task": MagicMock()} + with pytest.raises(TaskDeferred) as exc_info: + sensor.execute(context) + + trigger = exc_info.value.trigger + assert isinstance(trigger, DatabricksDeltaTableVersionTrigger) + assert trigger.table_name == TABLE_NAME + assert trigger.baseline_version == 5 + assert trigger.session_configuration == {"spark.sql.shuffle.partitions": "10"} + assert trigger.http_headers == [("X-Header", "1")] + assert trigger.client_parameters == {"client_param": "abc"} + assert trigger.hook_params == {"hook_param": "xyz"} + assert trigger.query_tags is not None + assert exc_info.value.method_name == "execute_complete" + + def test_execute_complete_success(self): + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=TABLE_NAME, + baseline_version=5, + sql_warehouse_name="test_wh", + ) + ti_mock = MagicMock() + context = {"ti": ti_mock} + event = { + "status": "success", + "version": 8, + "baseline_version": 5, + "table_name": TABLE_NAME, + "timestamp": "2026-10-05 12:00:00", + "operation": "WRITE", + } + res = sensor.execute_complete(context, event) + assert res == 8 + ti_mock.xcom_push.assert_called_once_with( + key="delta_table_version", + value={ + "version": 8, + "baseline_version": 5, + "table_name": TABLE_NAME, + "table_identity": {"catalog": "main", "schema": "default", "table": "users"}, + "timestamp": "2026-10-05 12:00:00", + "operation": "WRITE", + }, + ) + + def test_execute_complete_failure(self): + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=TABLE_NAME, + baseline_version=5, + sql_warehouse_name="test_wh", + ) + event = {"status": "error", "message": "Table was recreated"} + with pytest.raises(AirflowException, match="Table was recreated"): + sensor.execute_complete({}, event) + + def test_execute_complete_restores_baseline_version_for_deferred_enrollment(self): + sensor = DatabricksDeltaTableVersionSensor( + task_id=TASK_ID, + table_name=TABLE_NAME, + baseline_version=None, + sql_warehouse_name="test_wh", + ) + assert sensor._baseline_version is None + ti_mock = MagicMock() + context = {"ti": ti_mock} + event = { + "status": "success", + "version": 8, + "baseline_version": 5, + "table_name": TABLE_NAME, + "timestamp": "2026-10-05 12:00:00", + "operation": "WRITE", + } + res = sensor.execute_complete(context, event) + assert res == 8 + assert sensor._baseline_version == 5 + ti_mock.xcom_push.assert_called_once_with( + key="delta_table_version", + value={ + "version": 8, + "baseline_version": 5, + "table_name": TABLE_NAME, + "table_identity": {"catalog": "main", "schema": "default", "table": "users"}, + "timestamp": "2026-10-05 12:00:00", + "operation": "WRITE", + }, + ) diff --git a/providers/databricks/tests/unit/databricks/triggers/test_databricks_delta_table.py b/providers/databricks/tests/unit/databricks/triggers/test_databricks_delta_table.py new file mode 100644 index 0000000000000..0e8542ed671c9 --- /dev/null +++ b/providers/databricks/tests/unit/databricks/triggers/test_databricks_delta_table.py @@ -0,0 +1,272 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import time +from unittest.mock import patch + +import pytest + +from airflow.providers.databricks.triggers.databricks_delta_table import DatabricksDeltaTableVersionTrigger +from airflow.triggers.base import TriggerEvent + +TABLE_NAME = "main.default.users" +CONN_ID = "databricks_default" + + +class TestDatabricksDeltaTableVersionTrigger: + def test_serialization(self): + end_time = time.time() + 100 + trigger = DatabricksDeltaTableVersionTrigger( + table_name=TABLE_NAME, + databricks_conn_id=CONN_ID, + baseline_version=5, + target_version=10, + allow_recreation=True, + sql_warehouse_name="test_wh", + http_path="/sql/1.0/warehouses/abc", + catalog="main", + schema="default", + session_configuration={"spark.sql.shuffle.partitions": "10"}, + http_headers=[("X-Custom-Header", "value")], + client_parameters={"client_arg": "value"}, + hook_params={"hook_arg": "value"}, + query_tags={"tag_k": "tag_v"}, + polling_period_seconds=15, + end_time=end_time, + caller="test_caller", + ) + classpath, kwargs = trigger.serialize() + assert ( + classpath + == "airflow.providers.databricks.triggers.databricks_delta_table.DatabricksDeltaTableVersionTrigger" + ) + assert kwargs == { + "table_name": TABLE_NAME, + "databricks_conn_id": CONN_ID, + "baseline_version": 5, + "target_version": 10, + "allow_recreation": True, + "sql_warehouse_name": "test_wh", + "http_path": "/sql/1.0/warehouses/abc", + "catalog": "main", + "schema": "default", + "session_configuration": {"spark.sql.shuffle.partitions": "10"}, + "http_headers": [("X-Custom-Header", "value")], + "client_parameters": {"client_arg": "value"}, + "hook_params": {"hook_arg": "value"}, + "query_tags": {"tag_k": "tag_v"}, + "polling_period_seconds": 15, + "end_time": end_time, + "caller": "test_caller", + } + + @patch("airflow.providers.databricks.triggers.databricks_delta_table.DatabricksSqlHook") + def test_get_hook_passes_all_configuration(self, mock_hook_cls): + trigger = DatabricksDeltaTableVersionTrigger( + table_name=TABLE_NAME, + databricks_conn_id=CONN_ID, + sql_warehouse_name="test_wh", + http_path="/sql/1.0/warehouses/abc", + catalog="main", + schema="default", + session_configuration={"spark.sql.shuffle.partitions": "10"}, + http_headers=[("X-Custom-Header", "value")], + client_parameters={"client_arg": "value"}, + hook_params={"hook_arg": "value"}, + query_tags={"tag_k": "tag_v"}, + caller="custom_caller", + ) + hook = trigger._get_hook() + mock_hook_cls.assert_called_once_with( + databricks_conn_id=CONN_ID, + http_path="/sql/1.0/warehouses/abc", + sql_endpoint_name="test_wh", + session_configuration={"spark.sql.shuffle.partitions": "10"}, + http_headers=[("X-Custom-Header", "value")], + catalog="main", + schema="default", + caller="custom_caller", + query_tags={"tag_k": "tag_v"}, + client_arg="value", + hook_arg="value", + ) + assert hook == mock_hook_cls.return_value + + @pytest.mark.asyncio + @patch.object(DatabricksDeltaTableVersionTrigger, "_get_version") + async def test_run_success_newer_version(self, mock_get_version): + mock_get_version.return_value = (6, "2026-10-05 12:00:00", "WRITE") + trigger = DatabricksDeltaTableVersionTrigger( + table_name=TABLE_NAME, + databricks_conn_id=CONN_ID, + baseline_version=5, + polling_period_seconds=0.01, + ) + + generator = trigger.run() + event = await generator.asend(None) + assert isinstance(event, TriggerEvent) + assert event.payload["status"] == "success" + assert event.payload["version"] == 6 + assert event.payload["baseline_version"] == 5 + assert event.payload["table_name"] == TABLE_NAME + + @pytest.mark.asyncio + @patch.object(DatabricksDeltaTableVersionTrigger, "_get_version") + async def test_run_initial_enrollment(self, mock_get_version): + mock_get_version.side_effect = [ + (5, "2026-10-05 12:00:00", "WRITE"), + (6, "2026-10-05 12:01:00", "WRITE"), + ] + trigger = DatabricksDeltaTableVersionTrigger( + table_name=TABLE_NAME, + databricks_conn_id=CONN_ID, + baseline_version=None, + polling_period_seconds=0.01, + ) + + generator = trigger.run() + event = await generator.asend(None) + assert event.payload["status"] == "success" + assert event.payload["version"] == 6 + assert event.payload["baseline_version"] == 5 + + @pytest.mark.asyncio + @patch.object(DatabricksDeltaTableVersionTrigger, "_get_version") + async def test_run_recreation_allowed(self, mock_get_version): + mock_get_version.return_value = (0, "2026-10-05 12:00:00", "CREATE TABLE") + trigger = DatabricksDeltaTableVersionTrigger( + table_name=TABLE_NAME, + databricks_conn_id=CONN_ID, + baseline_version=10, + allow_recreation=True, + polling_period_seconds=0.01, + ) + + generator = trigger.run() + event = await generator.asend(None) + assert event.payload["status"] == "success" + assert event.payload["version"] == 0 + + @pytest.mark.asyncio + @patch.object(DatabricksDeltaTableVersionTrigger, "_get_version") + async def test_run_recreation_disallowed(self, mock_get_version): + mock_get_version.return_value = (0, "2026-10-05 12:00:00", "CREATE TABLE") + trigger = DatabricksDeltaTableVersionTrigger( + table_name=TABLE_NAME, + databricks_conn_id=CONN_ID, + baseline_version=10, + allow_recreation=False, + polling_period_seconds=0.01, + ) + + generator = trigger.run() + event = await generator.asend(None) + assert event.payload["status"] == "error" + assert "was recreated" in event.payload["message"] + + @pytest.mark.asyncio + @patch.object(DatabricksDeltaTableVersionTrigger, "_get_version") + async def test_run_target_version(self, mock_get_version): + mock_get_version.side_effect = [ + (8, "2026-10-05 12:00:00", "WRITE"), + (10, "2026-10-05 12:05:00", "WRITE"), + ] + trigger = DatabricksDeltaTableVersionTrigger( + table_name=TABLE_NAME, + databricks_conn_id=CONN_ID, + target_version=10, + polling_period_seconds=0.01, + ) + + generator = trigger.run() + event = await generator.asend(None) + assert event.payload["status"] == "success" + assert event.payload["version"] == 10 + + @pytest.mark.asyncio + @patch.object(DatabricksDeltaTableVersionTrigger, "_get_version") + async def test_run_timeout(self, mock_get_version): + mock_get_version.return_value = (5, "2026-10-05 12:00:00", "WRITE") + trigger = DatabricksDeltaTableVersionTrigger( + table_name=TABLE_NAME, + databricks_conn_id=CONN_ID, + baseline_version=5, + end_time=time.time() - 1, # Already timed out + polling_period_seconds=0.01, + ) + + generator = trigger.run() + event = await generator.asend(None) + assert event.payload["status"] == "timeout" + + @pytest.mark.asyncio + @patch.object(DatabricksDeltaTableVersionTrigger, "_get_version") + async def test_run_permanent_error_missing_table(self, mock_get_version): + mock_get_version.side_effect = Exception( + "TABLE_OR_VIEW_NOT_FOUND: The table or view cannot be found." + ) + trigger = DatabricksDeltaTableVersionTrigger( + table_name=TABLE_NAME, + databricks_conn_id=CONN_ID, + baseline_version=5, + polling_period_seconds=0.01, + ) + + generator = trigger.run() + event = await generator.asend(None) + assert event.payload["status"] == "error" + assert "permanent error" in event.payload["message"] + assert "TABLE_OR_VIEW_NOT_FOUND" in event.payload["message"] + + @pytest.mark.asyncio + @patch.object(DatabricksDeltaTableVersionTrigger, "_get_version") + async def test_run_permanent_error_permission_denied(self, mock_get_version): + mock_get_version.side_effect = Exception("[PERMISSION_DENIED] User lacks SELECT privilege on table.") + trigger = DatabricksDeltaTableVersionTrigger( + table_name=TABLE_NAME, + databricks_conn_id=CONN_ID, + baseline_version=5, + polling_period_seconds=0.01, + ) + + generator = trigger.run() + event = await generator.asend(None) + assert event.payload["status"] == "error" + assert "permanent error" in event.payload["message"] + assert "PERMISSION_DENIED" in event.payload["message"] + + @pytest.mark.asyncio + @patch.object(DatabricksDeltaTableVersionTrigger, "_get_version") + async def test_run_transient_error_retries_and_succeeds(self, mock_get_version): + mock_get_version.side_effect = [ + Exception("Connection reset by peer"), + (6, "2026-10-05 12:00:00", "WRITE"), + ] + trigger = DatabricksDeltaTableVersionTrigger( + table_name=TABLE_NAME, + databricks_conn_id=CONN_ID, + baseline_version=5, + polling_period_seconds=0.01, + ) + + generator = trigger.run() + event = await generator.asend(None) + assert event.payload["status"] == "success" + assert event.payload["version"] == 6