diff --git a/providers/docker/src/airflow/providers/docker/operators/docker.py b/providers/docker/src/airflow/providers/docker/operators/docker.py index e87f8d08e1664..e20f2cade37a2 100644 --- a/providers/docker/src/airflow/providers/docker/operators/docker.py +++ b/providers/docker/src/airflow/providers/docker/operators/docker.py @@ -308,10 +308,7 @@ def __init__( self.mount_tmp_dir = mount_tmp_dir self.tmp_dir = tmp_dir self.user = user - mounts = [mount if isinstance(mount, Mount) else Mount(**mount) for mount in (mounts or [])] - self.mounts: list[Mount] = mounts - for mount in self.mounts: - mount.template_fields = ("Source", "Target", "Type") + self.mounts = mounts or [] self.entrypoint = entrypoint self.working_dir = working_dir self.xcom_all = xcom_all @@ -489,7 +486,14 @@ def _copy_from_docker(self, container_id, src): lib = getattr(self, "pickling_library", pickle) return lib.load(file) + def _normalize_mounts(self) -> None: + # A user dict is rebuilt into a Mount, but rendering flattens a Mount into a + # plain dict with API-cased keys (Target/Source/Type) that docker-py already + # accepts, so leave those (and any real Mount) alone. + self.mounts = [m if isinstance(m, Mount) or "Target" in m else Mount(**m) for m in self.mounts] + def execute(self, context: Context) -> list[str] | str | None: + self._normalize_mounts() # Pull the docker image if `force_pull` is set or image does not exist locally if self.force_pull or not self.cli.images(name=self.image): self.log.info("::group::Pulling docker image %s", self.image) diff --git a/providers/docker/src/airflow/providers/docker/operators/docker_swarm.py b/providers/docker/src/airflow/providers/docker/operators/docker_swarm.py index 539cb1c2f7ad1..0e844f9323afe 100644 --- a/providers/docker/src/airflow/providers/docker/operators/docker_swarm.py +++ b/providers/docker/src/airflow/providers/docker/operators/docker_swarm.py @@ -174,6 +174,7 @@ def __init__( self.log_driver_config = None def execute(self, context: Context) -> None: + self._normalize_mounts() self.environment["AIRFLOW_TMP_DIR"] = self.tmp_dir return self._run_service() diff --git a/providers/docker/tests/unit/docker/operators/test_docker.py b/providers/docker/tests/unit/docker/operators/test_docker.py index 81fd757ddf765..324c593782718 100644 --- a/providers/docker/tests/unit/docker/operators/test_docker.py +++ b/providers/docker/tests/unit/docker/operators/test_docker.py @@ -874,12 +874,16 @@ def test_dict_mounts_are_normalized_to_mount_objects(self): Mount(target="/logs", source="logs", type="volume"), ], ) - assert all(isinstance(m, Mount) for m in op.mounts) - assert op.mounts[0]["Target"] == "/data" - assert op.mounts[0]["Source"] == "workspace" - assert op.mounts[0]["Type"] == "volume" - assert op.mounts[0]["ReadOnly"] is False - assert op.mounts[1]["Target"] == "/logs" + assert not isinstance(op.mounts[0], Mount) + + op.execute(None) + + passed_mounts = self.client_mock.create_host_config.call_args.kwargs["mounts"] + assert all(isinstance(m, Mount) for m in passed_mounts) + assert passed_mounts[0]["Target"] == "/data" + assert passed_mounts[0]["Source"] == "workspace" + assert passed_mounts[0]["ReadOnly"] is False + assert passed_mounts[1]["Target"] == "/logs" @pytest.mark.db_test def test_dict_mounts_are_templated(self, create_task_instance_of_operator): @@ -893,4 +897,24 @@ def test_dict_mounts_are_templated(self, create_task_instance_of_operator): ], ) rendered = ti.render_templates() - assert rendered.mounts[0]["Target"] == f"/{ti.run_id}" + assert rendered.mounts[0]["target"] == f"/{ti.run_id}" + + @pytest.mark.db_test + def test_mount_objects_survive_render_then_execute(self, create_task_instance_of_operator): + # A user-supplied Mount is a dict subclass, so the templater flattens it to a + # plain API-cased dict during rendering. Drive the real render -> execute path + # to prove execute() forwards it to docker-py instead of raising on it. + ti = create_task_instance_of_operator( + operator_class=DockerOperator, + dag_id="test", + task_id="test", + image="test", + mount_tmp_dir=False, + mounts=[Mount(source="workspace", target="/{{task_instance.run_id}}", type="volume")], + ) + task = ti.render_templates() + task.execute(None) + + passed_mounts = self.client_mock.create_host_config.call_args.kwargs["mounts"] + assert passed_mounts[0]["Target"] == f"/{ti.run_id}" + assert passed_mounts[0]["Source"] == "workspace" diff --git a/providers/docker/tests/unit/docker/operators/test_docker_swarm.py b/providers/docker/tests/unit/docker/operators/test_docker_swarm.py index 65e83d45ddd27..9f8fdabd91048 100644 --- a/providers/docker/tests/unit/docker/operators/test_docker_swarm.py +++ b/providers/docker/tests/unit/docker/operators/test_docker_swarm.py @@ -209,6 +209,25 @@ def test_no_auto_remove(self, types_mock, docker_api_client_patcher): "Docker service being removed even when `auto_remove` set to `never`" ) + @mock.patch("airflow.providers.docker.operators.docker_swarm.types") + def test_dict_mounts_converted_at_execute(self, types_mock, docker_api_client_patcher): + client_mock = mock.Mock(spec=APIClient) + client_mock.create_service.return_value = {"ID": "some_id"} + client_mock.images.return_value = [] + client_mock.pull.return_value = [b'{"status":"pull log"}'] + client_mock.tasks.return_value = [{"ServiceID": "some_id", "Status": {"State": "complete"}}] + docker_api_client_patcher.return_value = client_mock + + operator = DockerSwarmOperator( + image="", + task_id="unittest", + enable_logging=False, + mounts=[{"source": "/host", "target": "/container", "type": "bind"}], + ) + operator.execute(None) + + assert all(isinstance(m, types.Mount) for m in operator.mounts) + @pytest.mark.parametrize("status", ["failed", "shutdown", "rejected", "orphaned", "remove"]) @mock.patch("airflow.providers.docker.operators.docker_swarm.types") def test_non_complete_service_raises_error(self, types_mock, docker_api_client_patcher, status): diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt b/scripts/ci/prek/validate_operators_init_exemptions.txt index 29ee3dd263922..0890990686d76 100644 --- a/scripts/ci/prek/validate_operators_init_exemptions.txt +++ b/scripts/ci/prek/validate_operators_init_exemptions.txt @@ -22,7 +22,6 @@ providers/amazon/src/airflow/providers/amazon/aws/transfers/gcs_to_s3.py::GCSToS providers/amazon/src/airflow/providers/amazon/aws/transfers/s3_to_redshift.py::S3ToRedshiftOperator providers/anthropic/src/airflow/providers/anthropic/operators/agent.py::AnthropicAgentSessionOperator providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py::KubernetesPodOperator -providers/docker/src/airflow/providers/docker/operators/docker.py::DockerOperator providers/google/src/airflow/providers/google/cloud/operators/bigquery.py::BigQueryInsertJobOperator providers/google/src/airflow/providers/google/cloud/operators/cloud_batch.py::CloudBatchSubmitJobOperator providers/google/src/airflow/providers/google/cloud/operators/cloud_build.py::CloudBuildCreateBuildOperator