From a947255593b5629f6696fce04ac20c31b9ada48c Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Thu, 23 Jul 2026 15:14:41 -0700 Subject: [PATCH 1/3] Validate DatabricksSQLStatementsSensor statement fields after rendering statement and statement_id are template fields, rendered after __init__ runs. The constructor enforced their mutual exclusivity and presence there, acting on the un-rendered Jinja expressions. Move both checks into execute(); the warehouse_id check reads no template field and stays in __init__. related: #70296 Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../providers/databricks/sensors/databricks.py | 11 ++++++----- .../unit/databricks/sensors/test_databricks.py | 14 ++++++++++++++ .../ci/prek/validate_operators_init_exemptions.txt | 1 - 3 files changed, 20 insertions(+), 6 deletions(-) diff --git a/providers/databricks/src/airflow/providers/databricks/sensors/databricks.py b/providers/databricks/src/airflow/providers/databricks/sensors/databricks.py index 9497cf8624d35..55e64106d113c 100644 --- a/providers/databricks/src/airflow/providers/databricks/sensors/databricks.py +++ b/providers/databricks/src/airflow/providers/databricks/sensors/databricks.py @@ -70,11 +70,6 @@ def __init__( **kwargs, ): # Handle the scenario where either both statement and statement_id are set/not set - if statement and statement_id: - raise AirflowException("Cannot provide both statement and statement_id.") - if not statement and not statement_id: - raise AirflowException("One of either statement or statement_id must be provided.") - if not warehouse_id: raise AirflowException("warehouse_id must be provided.") @@ -112,6 +107,12 @@ def _get_hook(self, caller: str) -> DatabricksHook: ) def execute(self, context: Context): + # statement/statement_id are template fields; validate their combination here, after + # rendering, rather than in __init__ where they are still un-rendered Jinja expressions. + if self.statement and self.statement_id: + raise AirflowException("Cannot provide both statement and statement_id.") + if not self.statement and not self.statement_id: + raise AirflowException("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 fe04781f65d1d..c8d263a867ed1 100644 --- a/providers/databricks/tests/unit/databricks/sensors/test_databricks.py +++ b/providers/databricks/tests/unit/databricks/sensors/test_databricks.py @@ -68,6 +68,20 @@ def test_init_statement_id(self): assert op.statement_id == STATEMENT_ID assert op.warehouse_id == WAREHOUSE_ID + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"statement": STATEMENT, "statement_id": STATEMENT_ID}, "Cannot provide both"), + ({}, "One of either statement or statement_id"), + ], + ) + def test_statement_combination_validated_at_execute(self, kwargs, match): + # statement/statement_id are template fields: the combination check must run at execute + # (after rendering), so it no longer raises in __init__. + op = DatabricksSQLStatementsSensor(task_id=TASK_ID, warehouse_id=WAREHOUSE_ID, **kwargs) + with pytest.raises(AirflowException, match=match): + op.execute(None) + @mock.patch("airflow.providers.databricks.sensors.databricks.DatabricksHook") def test_exec_success(self, db_mock_class): """ diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt b/scripts/ci/prek/validate_operators_init_exemptions.txt index 01fd5bc56dbb7..adedbe2f820f6 100644 --- a/scripts/ci/prek/validate_operators_init_exemptions.txt +++ b/scripts/ci/prek/validate_operators_init_exemptions.txt @@ -40,7 +40,6 @@ providers/databricks/src/airflow/providers/databricks/operators/databricks_repos providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py::DatabricksReposDeleteOperator providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py::DatabricksReposUpdateOperator providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py::DatabricksCopyIntoOperator -providers/databricks/src/airflow/providers/databricks/sensors/databricks.py::DatabricksSQLStatementsSensor providers/dbt/cloud/src/airflow/providers/dbt/cloud/operators/dbt.py::DbtCloudGetJobRunArtifactOperator providers/docker/src/airflow/providers/docker/operators/docker.py::DockerOperator providers/google/src/airflow/providers/google/cloud/operators/bigquery.py::BigQueryInsertJobOperator From 10c8b9b0fb367eb94ba53d63df5ec0fcc586d0fe Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Thu, 23 Jul 2026 23:13:10 -0700 Subject: [PATCH 2/3] Tighten DatabricksSQLStatementsSensor validation comments Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../src/airflow/providers/databricks/sensors/databricks.py | 3 +-- .../tests/unit/databricks/sensors/test_databricks.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/providers/databricks/src/airflow/providers/databricks/sensors/databricks.py b/providers/databricks/src/airflow/providers/databricks/sensors/databricks.py index 55e64106d113c..8d162b11708cb 100644 --- a/providers/databricks/src/airflow/providers/databricks/sensors/databricks.py +++ b/providers/databricks/src/airflow/providers/databricks/sensors/databricks.py @@ -107,8 +107,7 @@ def _get_hook(self, caller: str) -> DatabricksHook: ) def execute(self, context: Context): - # statement/statement_id are template fields; validate their combination here, after - # rendering, rather than in __init__ where they are still un-rendered Jinja expressions. + # statement/statement_id are template fields; validate their combination after rendering. if self.statement and self.statement_id: raise AirflowException("Cannot provide both statement and statement_id.") if not self.statement and not self.statement_id: diff --git a/providers/databricks/tests/unit/databricks/sensors/test_databricks.py b/providers/databricks/tests/unit/databricks/sensors/test_databricks.py index c8d263a867ed1..51093cf33222f 100644 --- a/providers/databricks/tests/unit/databricks/sensors/test_databricks.py +++ b/providers/databricks/tests/unit/databricks/sensors/test_databricks.py @@ -76,8 +76,7 @@ def test_init_statement_id(self): ], ) def test_statement_combination_validated_at_execute(self, kwargs, match): - # statement/statement_id are template fields: the combination check must run at execute - # (after rendering), so it no longer raises in __init__. + # statement/statement_id are template fields; the combination check runs at execute. op = DatabricksSQLStatementsSensor(task_id=TASK_ID, warehouse_id=WAREHOUSE_ID, **kwargs) with pytest.raises(AirflowException, match=match): op.execute(None) From 3e11bd9eb73a67e70aab63674ea043585af65eaf Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Fri, 24 Jul 2026 11:29:57 -0700 Subject: [PATCH 3/3] Drop narrating comments from DatabricksSQLStatementsSensor Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../src/airflow/providers/databricks/sensors/databricks.py | 1 - .../databricks/tests/unit/databricks/sensors/test_databricks.py | 1 - 2 files changed, 2 deletions(-) diff --git a/providers/databricks/src/airflow/providers/databricks/sensors/databricks.py b/providers/databricks/src/airflow/providers/databricks/sensors/databricks.py index 8d162b11708cb..729d84af78311 100644 --- a/providers/databricks/src/airflow/providers/databricks/sensors/databricks.py +++ b/providers/databricks/src/airflow/providers/databricks/sensors/databricks.py @@ -107,7 +107,6 @@ def _get_hook(self, caller: str) -> DatabricksHook: ) def execute(self, context: Context): - # statement/statement_id are template fields; validate their combination after rendering. if self.statement and self.statement_id: raise AirflowException("Cannot provide both statement and statement_id.") if not self.statement and not self.statement_id: diff --git a/providers/databricks/tests/unit/databricks/sensors/test_databricks.py b/providers/databricks/tests/unit/databricks/sensors/test_databricks.py index 51093cf33222f..08517a17f884a 100644 --- a/providers/databricks/tests/unit/databricks/sensors/test_databricks.py +++ b/providers/databricks/tests/unit/databricks/sensors/test_databricks.py @@ -76,7 +76,6 @@ def test_init_statement_id(self): ], ) def test_statement_combination_validated_at_execute(self, kwargs, match): - # statement/statement_id are template fields; the combination check runs at execute. op = DatabricksSQLStatementsSensor(task_id=TASK_ID, warehouse_id=WAREHOUSE_ID, **kwargs) with pytest.raises(AirflowException, match=match): op.execute(None)