diff --git a/providers/databricks/docs/operators/jobs_create.rst b/providers/databricks/docs/operators/jobs_create.rst index 115dcc39d09fc..b5946c45b14f5 100644 --- a/providers/databricks/docs/operators/jobs_create.rst +++ b/providers/databricks/docs/operators/jobs_create.rst @@ -71,7 +71,8 @@ If ``parameters`` is not set in ``json`` and the operator's ``params`` dict is n each key/value pair in ``params`` is converted into one such ``{"name": key, "default": value}`` entry, so that Airflow Dag params can be forwarded as Databricks job parameters without hardcoding the API shape in ``json``. If ``json`` already contains ``parameters``, -it is left untouched. +it is left untouched. Params whose value resolves to ``None`` are skipped, since Databricks +has no value to receive for them. .. code-block:: python diff --git a/providers/databricks/docs/operators/run_now.rst b/providers/databricks/docs/operators/run_now.rst index 0b53f659b8afa..20260286040f2 100644 --- a/providers/databricks/docs/operators/run_now.rst +++ b/providers/databricks/docs/operators/run_now.rst @@ -60,8 +60,9 @@ for this run. If ``job_parameters`` is not set in ``json`` and the operator's ``params`` dict is non-empty, ``params`` is forwarded as ``job_parameters`` as-is, so Airflow Dag params can be passed dynamically to a run without hardcoding them in ``json``. If ``json`` already -contains ``job_parameters``, it is left untouched. You can set ``forward_dag_params=False`` to -disable this parameter forwarding behavior. +contains ``job_parameters``, it is left untouched. Params whose value resolves to ``None`` are +skipped, since Databricks has no value to receive for them. You can set +``forward_dag_params=False`` to disable this parameter forwarding behavior. .. note:: The Databricks API does not permit ``job_parameters`` to be used in combination with diff --git a/providers/databricks/docs/operators/submit_run.rst b/providers/databricks/docs/operators/submit_run.rst index 0e85b7dd8f663..92588d06ea531 100644 --- a/providers/databricks/docs/operators/submit_run.rst +++ b/providers/databricks/docs/operators/submit_run.rst @@ -125,6 +125,9 @@ Tasks whose only parameter slot is ``List[str]`` (``spark_jar_task``, ``spark_py dict to a positional argument list — pass those parameters explicitly via the ``json`` or ``tasks`` argument. +Params whose value resolves to ``None`` are skipped, since Databricks has no value to +receive for them. + .. code-block:: python notebook_run = DatabricksSubmitRunOperator( diff --git a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py index 6f587f398245e..7377dfefd7300 100644 --- a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py +++ b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py @@ -329,6 +329,14 @@ def _handle_deferrable_databricks_operator_completion(event: dict, log: Logger) ) +def _get_forwardable_dag_params(params: Mapping[str, Any]) -> dict[str, Any]: + """Return the Dag params that can be forwarded to Databricks.""" + # A Param that resolves to None has no value to forward, and the Databricks payload has no + # slot for null. The dict() round-trip is what resolves each Param: ParamsDict.items() yields + # the Param objects themselves, none of which would ever compare equal to None. + return {key: value for key, value in dict(params).items() if value is not None} + + def _inject_airflow_params_into_task(task: dict, params: dict) -> None: """Set dict-shaped per-task parameter fields from ``params`` if they are not already set.""" for task_key, field in _DICT_PARAM_FIELD_BY_TASK.items(): @@ -468,7 +476,8 @@ class DatabricksCreateJobsOperator(BaseOperator): If ``parameters`` is not set in ``json`` and the operator's ``params`` dict is non-empty, the operator's ``params`` are automatically converted to job-level ``parameters`` (a list of ``{"name": k, "default": v}`` entries) so that Airflow Dag params can be forwarded as - Databricks job parameters without hardcoding them in ``json``. + Databricks job parameters without hardcoding them in ``json``. Params whose value is + ``None`` are skipped. """ @@ -575,8 +584,8 @@ def execute(self, context: Context) -> int: if "name" not in json: raise AirflowException("Missing required parameter: name") job_id = self._hook.find_job_id_by_name(json["name"]) - if not json.get("parameters") and self.params: - json["parameters"] = [{"name": k, "default": v} for k, v in dict(self.params).items()] + if not json.get("parameters") and (forwardable_params := _get_forwardable_dag_params(self.params)): + json["parameters"] = [{"name": k, "default": v} for k, v in forwardable_params.items()] if job_id is None: return self._hook.create_job(json) self._hook.reset_job(str(job_id), json) @@ -734,7 +743,7 @@ class DatabricksSubmitRunOperator(ResumableJobMixin, BaseOperator): ``sql_task.parameters``, ``run_job_task.job_parameters``. Tasks whose only parameter field is ``List[str]`` (``spark_jar_task``, ``spark_python_task``, ``spark_submit_task``) are skipped because there is no canonical mapping from a key/value dict to a positional - argument list. + argument list. Params whose value is ``None`` are skipped. """ external_id_key = "databricks_run_id" @@ -921,15 +930,14 @@ def _prepare_submit_json(self, context: Context) -> dict[str, Any]: json["pipeline_task"]["pipeline_id"] = self._hook.find_pipeline_id_by_name(pipeline_name) del json["pipeline_task"]["pipeline_name"] - if self.params: - params_dump = dict(self.params) + if forwardable_params := _get_forwardable_dag_params(self.params): tasks = json.get("tasks") if isinstance(tasks, list): for task in tasks: if isinstance(task, dict): - _inject_airflow_params_into_task(task, params_dump) + _inject_airflow_params_into_task(task, forwardable_params) else: - _inject_airflow_params_into_task(json, params_dump) + _inject_airflow_params_into_task(json, forwardable_params) if self.openlineage_inject_parent_job_info or self.openlineage_inject_transport_info: self.log.info("Automatic injection of OpenLineage information into Spark properties is enabled.") @@ -1221,7 +1229,8 @@ class DatabricksRunNowOperator(ResumableJobMixin, BaseOperator): If ``job_parameters`` is not set in ``json`` and the operator's ``params`` dict is non-empty, the operator's ``params`` are automatically forwarded as ``job_parameters`` so that Airflow Dag params can be passed dynamically to Databricks runs without - hardcoding them in ``json``. Set ``forward_dag_params=False`` to disable this. + hardcoding them in ``json``. Params whose value is ``None`` are skipped. Set + ``forward_dag_params=False`` to disable this. Note that the Databricks API does not permit ``job_parameters`` to be used in combination with ``notebook_params``, ``python_params``, ``jar_params``, ``spark_submit_params``, ``python_named_params``, or ``dbt_commands``; auto-forwarding is automatically skipped @@ -1378,10 +1387,10 @@ def _build_run_now_payload(self) -> dict[str, Any]: if ( self.forward_dag_params and not json.get("job_parameters") - and self.params and not any(k in json for k in _RUN_NOW_PARAM_SLOTS_CONFLICTING_WITH_JOB_PARAMETERS) + and (forwardable_params := _get_forwardable_dag_params(self.params)) ): - json["job_parameters"] = dict(self.params) + json["job_parameters"] = forwardable_params return json diff --git a/providers/databricks/tests/unit/databricks/operators/test_databricks.py b/providers/databricks/tests/unit/databricks/operators/test_databricks.py index 4a5dc4bd43fc7..3da1440daef83 100644 --- a/providers/databricks/tests/unit/databricks/operators/test_databricks.py +++ b/providers/databricks/tests/unit/databricks/operators/test_databricks.py @@ -728,6 +728,32 @@ def test_does_not_override_existing_parameters(self, db_mock_class, found_job_id settings = call_args[0] if hook_method == "create_job" else call_args[1] assert settings["parameters"] == JOB_PARAMS + @pytest.mark.parametrize( + ("params", "expected_parameters"), + [ + pytest.param( + {"env": "prod", "start_date_str": None}, + [{"name": "env", "default": "prod"}], + id="some-params-none", + ), + pytest.param({"start_date_str": None}, None, id="all-params-none"), + ], + ) + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + def test_skips_airflow_params_whose_value_is_none(self, db_mock_class, params, expected_parameters): + op = DatabricksCreateJobsOperator( + task_id=TASK_ID, + json={"name": JOB_NAME, "tasks": TASKS}, + params=params, + ) + db_mock = db_mock_class.return_value + db_mock.find_job_id_by_name.return_value = None + + op.execute({}) + + settings = db_mock.create_job.call_args.args[0] + assert settings.get("parameters") == expected_parameters + class TestDatabricksSubmitRunOperator: @staticmethod @@ -1510,6 +1536,33 @@ def test_submit_run_does_not_override_existing_task_parameters(self, db_mock_cla actual = db_mock.submit_run.call_args.args[0] assert actual["notebook_task"]["base_parameters"] == {"explicit": "value"} + @pytest.mark.parametrize( + ("params", "expected_named_parameters"), + [ + pytest.param({"env": "prod", "start_date_str": None}, {"env": "prod"}, id="some-params-none"), + pytest.param({"start_date_str": None}, None, id="all-params-none"), + ], + ) + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + def test_submit_run_skips_airflow_params_whose_value_is_none( + self, db_mock_class, params, expected_named_parameters + ): + op = DatabricksSubmitRunOperator( + durable=False, + task_id=TASK_ID, + json={"python_wheel_task": {"package_name": "my_package", "entry_point": "main"}}, + new_cluster=NEW_CLUSTER, + params=params, + ) + db_mock = db_mock_class.return_value + db_mock.submit_run.return_value = RUN_ID + db_mock.get_run = make_run_with_state_mock("TERMINATED", "SUCCESS") + + op.execute(None) + + actual = db_mock.submit_run.call_args.args[0] + assert actual["python_wheel_task"].get("named_parameters") == expected_named_parameters + @pytest.mark.parametrize( ("json", "exception_message"), [ @@ -2955,6 +3008,32 @@ def test_run_now_does_not_override_existing_job_parameters(self, db_mock_class): actual = db_mock.run_now.call_args.args[0] assert actual["job_parameters"] == {"explicit": "value"} + @pytest.mark.parametrize( + ("params", "expected_job_parameters"), + [ + pytest.param({"env": "prod", "start_date_str": None}, {"env": "prod"}, id="some-params-none"), + pytest.param({"start_date_str": None}, None, id="all-params-none"), + ], + ) + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + def test_run_now_skips_airflow_params_whose_value_is_none( + self, db_mock_class, params, expected_job_parameters + ): + op = DatabricksRunNowOperator( + durable=False, + task_id=TASK_ID, + job_id=JOB_ID, + params=params, + ) + db_mock = db_mock_class.return_value + db_mock.run_now.return_value = RUN_ID + db_mock.get_run = make_run_with_state_mock("TERMINATED", "SUCCESS") + + op.execute(None) + + actual = db_mock.run_now.call_args.args[0] + assert actual.get("job_parameters") == expected_job_parameters + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") def test_run_now_does_not_inject_airflow_params_when_forward_dag_params_is_false(self, db_mock_class): """