diff --git a/generated/known_airflow_exceptions.txt b/generated/known_airflow_exceptions.txt index bfb87136d2ee0..f5d82fa18f85e 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 729d84af78311..8ab7d327af3e3 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 @@ def __init__( 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 @@ def _get_hook(self, caller: str) -> DatabricksHook: ) 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 08517a17f884a..c0a40738246a8 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 @@ 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 @@ def test_init_statement_id(self): 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): """