This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 83149aaf222 Validate DatabricksSQLStatementsSensor exclusivity at
__init__ (#70831)
83149aaf222 is described below
commit 83149aaf22233bd6be2b644ab3ce8a262b57ba2e
Author: ahilashsasidharan <[email protected]>
AuthorDate: Tue Sep 8 20:19:05 2026 -0400
Validate DatabricksSQLStatementsSensor exclusivity at __init__ (#70831)
* Validate DatabricksSQLStatementsSensor exclusivity at __init__
* Change DatabricksSQLStatementsSensor init exclusivity check to raise
ValueError over AirflowException
* Enhance tests and replace nearby AirflowExceptions with ValueError
* Add additional test coverage for render_template_as_native_obj=True
scenario and update error message to hold for all failure scenarios
---
generated/known_airflow_exceptions.txt | 2 +-
.../providers/databricks/sensors/databricks.py | 13 +--
.../unit/databricks/sensors/test_databricks.py | 93 ++++++++++++++++++++--
3 files changed, 95 insertions(+), 13 deletions(-)
diff --git a/generated/known_airflow_exceptions.txt
b/generated/known_airflow_exceptions.txt
index fd54d247222..cd9b8e3c99c 100644
--- a/generated/known_airflow_exceptions.txt
+++ b/generated/known_airflow_exceptions.txt
@@ -179,7 +179,7 @@
providers/databricks/src/airflow/providers/databricks/operators/databricks_repos
providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py::8
providers/databricks/src/airflow/providers/databricks/operators/databricks_workflow.py::4
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py::7
-providers/databricks/src/airflow/providers/databricks/sensors/databricks.py::4
+providers/databricks/src/airflow/providers/databricks/sensors/databricks.py::1
providers/databricks/src/airflow/providers/databricks/sensors/databricks_partition.py::4
providers/databricks/src/airflow/providers/databricks/sensors/databricks_sql.py::1
providers/databricks/src/airflow/providers/databricks/utils/databricks.py::3
diff --git
a/providers/databricks/src/airflow/providers/databricks/sensors/databricks.py
b/providers/databricks/src/airflow/providers/databricks/sensors/databricks.py
index 729d84af783..8ab7d327af3 100644
---
a/providers/databricks/src/airflow/providers/databricks/sensors/databricks.py
+++
b/providers/databricks/src/airflow/providers/databricks/sensors/databricks.py
@@ -69,9 +69,11 @@ class
DatabricksSQLStatementsSensor(DatabricksSQLStatementsMixin, BaseSensorOper
include_airflow_query_tags: bool = True,
**kwargs,
):
- # Handle the scenario where either both statement and statement_id are
set/not set
+ if statement is not None and statement_id is not None:
+ raise ValueError("Provide exactly one of statement or
statement_id.")
+
if not warehouse_id:
- raise AirflowException("warehouse_id must be provided.")
+ raise ValueError("warehouse_id must be provided.")
super().__init__(**kwargs)
@@ -107,10 +109,11 @@ class
DatabricksSQLStatementsSensor(DatabricksSQLStatementsMixin, BaseSensorOper
)
def execute(self, context: Context):
- if self.statement and self.statement_id:
- raise AirflowException("Cannot provide both statement and
statement_id.")
+ # Both fields are templated, so "neither resolves to a value" is only
knowable
+ # after rendering — __init__ cannot catch it. The both-provided case
is a pure
+ # provision error and is checked there instead.
if not self.statement and not self.statement_id:
- raise AirflowException("One of either statement or statement_id
must be provided.")
+ raise ValueError("One of either statement or statement_id must be
provided.")
if not self.statement_id:
# Otherwise, we'll go ahead and "submit" the statement
tags = build_query_tags(context, self.query_tags,
self.include_airflow_query_tags)
diff --git
a/providers/databricks/tests/unit/databricks/sensors/test_databricks.py
b/providers/databricks/tests/unit/databricks/sensors/test_databricks.py
index 08517a17f88..c0a40738246 100644
--- a/providers/databricks/tests/unit/databricks/sensors/test_databricks.py
+++ b/providers/databricks/tests/unit/databricks/sensors/test_databricks.py
@@ -22,7 +22,8 @@ from unittest import mock
import pytest
from tenacity import stop_after_attempt, wait_incrementing
-from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred
+from airflow.models.dag import DAG
+from airflow.providers.common.compat.sdk import AirflowException,
TaskDeferred, timezone
from airflow.providers.databricks.hooks.databricks import SQLStatementState
from airflow.providers.databricks.sensors.databricks import
DatabricksSQLStatementsSensor
from airflow.providers.databricks.triggers.databricks import
DatabricksSQLStatementExecutionTrigger
@@ -69,17 +70,95 @@ class TestDatabricksSQLStatementsSensor:
assert op.warehouse_id == WAREHOUSE_ID
@pytest.mark.parametrize(
- ("kwargs", "match"),
+ ("statement", "statement_id"),
[
- ({"statement": STATEMENT, "statement_id": STATEMENT_ID}, "Cannot
provide both"),
- ({}, "One of either statement or statement_id"),
+ (STATEMENT, STATEMENT_ID),
+ (STATEMENT, ""),
+ ("", ""),
],
)
- def test_statement_combination_validated_at_execute(self, kwargs, match):
- op = DatabricksSQLStatementsSensor(task_id=TASK_ID,
warehouse_id=WAREHOUSE_ID, **kwargs)
- with pytest.raises(AirflowException, match=match):
+ def test_both_statements_included_validated_at_init(self, statement,
statement_id):
+ with pytest.raises(ValueError, match="Provide exactly one of"):
+ DatabricksSQLStatementsSensor(
+ statement=statement,
+ statement_id=statement_id,
+ task_id=TASK_ID,
+ warehouse_id=WAREHOUSE_ID,
+ )
+
+ @pytest.mark.parametrize(
+ ("statement", "statement_id"),
+ [
+ ("{{ None }}", STATEMENT_ID),
+ (STATEMENT, "{{ None }}"),
+ ("{{ None }}", "{{ None }}"),
+ ],
+ )
+ def
test_both_provided_with_template_renders_to_none_validated_at_init(self,
statement, statement_id):
+ """
+ Both statement and statement_id are provided; at least one is a
template that
+ would render to None under render_template_as_native_obj=True. The
constructor
+ must raise before any rendering occurs, so the check is pinned to
__init__.
+ If the exclusivity check were moved to execute(), this test would fail
because
+ the constructor would succeed and if rendered before execute() the
template
+ would render to None making it appear only one value was provided so
the moved
+ check would not catch the exclusivity violation
+ """
+ dag = DAG(
+ dag_id="test_native_obj_dag",
+ start_date=timezone.datetime(2025, 1, 1),
+ schedule=None,
+ render_template_as_native_obj=True,
+ )
+ with pytest.raises(ValueError, match="Provide exactly one of"):
+ DatabricksSQLStatementsSensor(
+ task_id=TASK_ID,
+ warehouse_id=WAREHOUSE_ID,
+ statement=statement,
+ statement_id=statement_id,
+ dag=dag,
+ )
+
+ @pytest.mark.parametrize(
+ ("statement", "statement_id"),
+ [
+ (None, None),
+ ("", None),
+ ],
+ )
+ def test_both_statements_missing_validated_at_execute(self, statement,
statement_id):
+ op = DatabricksSQLStatementsSensor(
+ task_id=TASK_ID, warehouse_id=WAREHOUSE_ID, statement=statement,
statement_id=statement_id
+ )
+ with pytest.raises(ValueError, match="One of either statement or
statement_id"):
op.execute(None)
+ @pytest.mark.parametrize(
+ ("statement", "statement_id"),
+ [
+ (None, "{{ None }}"),
+ ("{{ None }}", None),
+ ],
+ )
+ def test_both_missing_after_template_rendered_validated_at_execute(self,
statement, statement_id):
+ dag = DAG(
+ dag_id="test_native_obj_dag",
+ start_date=timezone.datetime(2025, 1, 1),
+ schedule=None,
+ render_template_as_native_obj=True,
+ )
+ op = DatabricksSQLStatementsSensor(
+ task_id=TASK_ID,
+ warehouse_id=WAREHOUSE_ID,
+ statement=statement,
+ statement_id=statement_id,
+ dag=dag,
+ )
+ context = {"dag": dag}
+ op.render_template_fields(context)
+ with pytest.raises(ValueError, match="One of either statement or
statement_id"):
+ op.execute(context)
+
@mock.patch("airflow.providers.databricks.sensors.databricks.DatabricksHook")
def test_exec_success(self, db_mock_class):
"""
