From e8423e552b528f9c66758d4df32b361d8e1e1944 Mon Sep 17 00:00:00 2001 From: axelray-dev <110029405+axelray-dev@users.noreply.github.com> Date: Tue, 14 Jul 2026 00:48:50 +0800 Subject: [PATCH 1/4] Fix custom SSH remote job cleanup paths (#69813) --- .../providers/ssh/operators/ssh_remote_job.py | 5 ++-- .../airflow/providers/ssh/utils/remote_job.py | 23 +++++++++++++------ .../tests/unit/ssh/utils/test_remote_job.py | 15 ++++++++++++ 3 files changed, 34 insertions(+), 9 deletions(-) diff --git a/providers/ssh/src/airflow/providers/ssh/operators/ssh_remote_job.py b/providers/ssh/src/airflow/providers/ssh/operators/ssh_remote_job.py index dc3414eca686d..96a310be9cf42 100644 --- a/providers/ssh/src/airflow/providers/ssh/operators/ssh_remote_job.py +++ b/providers/ssh/src/airflow/providers/ssh/operators/ssh_remote_job.py @@ -451,10 +451,11 @@ def _cleanup_remote_job(self, job_dir: str, remote_os: str) -> None: leave the directory rather than failing the (already finished) task. """ self.log.info("Cleaning up remote job directory: %s", job_dir) + expected_base = self._paths.base_dir if self._paths else self.remote_base_dir if remote_os == "posix": - cleanup_cmd = build_posix_cleanup_command(job_dir) + cleanup_cmd = build_posix_cleanup_command(job_dir, expected_base=expected_base) else: - cleanup_cmd = build_windows_cleanup_command(job_dir) + cleanup_cmd = build_windows_cleanup_command(job_dir, expected_base=expected_base) last_error: Exception | None = None for attempt in range(1, self.cleanup_retries + 1): diff --git a/providers/ssh/src/airflow/providers/ssh/utils/remote_job.py b/providers/ssh/src/airflow/providers/ssh/utils/remote_job.py index 496b0179d8913..35bd63da6e816 100644 --- a/providers/ssh/src/airflow/providers/ssh/utils/remote_job.py +++ b/providers/ssh/src/airflow/providers/ssh/utils/remote_job.py @@ -30,18 +30,25 @@ WINDOWS_DEFAULT_BASE_DIR = "$env:TEMP\\airflow-ssh-jobs" -def _validate_job_dir(job_dir: str, remote_os: Literal["posix", "windows"]) -> None: +def _validate_job_dir( + job_dir: str, + remote_os: Literal["posix", "windows"], + expected_base: str | None = None, +) -> None: """ Validate that job_dir is under the expected base directory. :param job_dir: The job directory path to validate :param remote_os: Operating system type + :param expected_base: Expected base directory for the job :raises ValueError: If job_dir doesn't start with the expected base path """ if remote_os == "posix": - expected_prefix = POSIX_DEFAULT_BASE_DIR + "/" + expected_base = expected_base or POSIX_DEFAULT_BASE_DIR + expected_prefix = expected_base + "/" else: - expected_prefix = WINDOWS_DEFAULT_BASE_DIR + "\\" + expected_base = expected_base or WINDOWS_DEFAULT_BASE_DIR + expected_prefix = expected_base + "\\" if not job_dir.startswith(expected_prefix): raise ValueError( @@ -452,27 +459,29 @@ def build_windows_kill_command(pid_file: str) -> str: return f"powershell.exe -NoProfile -NonInteractive -EncodedCommand {encoded_script}" -def build_posix_cleanup_command(job_dir: str) -> str: +def build_posix_cleanup_command(job_dir: str, expected_base: str | None = None) -> str: """ Build a POSIX command to clean up the job directory. :param job_dir: Path to the job directory + :param expected_base: Expected base directory for the job :return: Shell command to remove the directory :raises ValueError: If job_dir is not under the expected base directory """ - _validate_job_dir(job_dir, "posix") + _validate_job_dir(job_dir, "posix", expected_base) return f"rm -rf '{job_dir}'" -def build_windows_cleanup_command(job_dir: str) -> str: +def build_windows_cleanup_command(job_dir: str, expected_base: str | None = None) -> str: """ Build a PowerShell command to clean up the job directory. :param job_dir: Path to the job directory + :param expected_base: Expected base directory for the job :return: PowerShell command to remove the directory :raises ValueError: If job_dir is not under the expected base directory """ - _validate_job_dir(job_dir, "windows") + _validate_job_dir(job_dir, "windows", expected_base) escaped_path = job_dir.replace("'", "''") script = f"Remove-Item -Recurse -Force -Path '{escaped_path}' -ErrorAction SilentlyContinue" script_bytes = script.encode("utf-16-le") diff --git a/providers/ssh/tests/unit/ssh/utils/test_remote_job.py b/providers/ssh/tests/unit/ssh/utils/test_remote_job.py index 82b285d4587fc..b79d2f6129012 100644 --- a/providers/ssh/tests/unit/ssh/utils/test_remote_job.py +++ b/providers/ssh/tests/unit/ssh/utils/test_remote_job.py @@ -394,6 +394,14 @@ def test_posix_cleanup(self): assert "rm -rf" in cmd assert "/tmp/airflow-ssh-jobs/job_123" in cmd + def test_posix_cleanup_with_custom_base(self): + """Test POSIX cleanup under a custom base directory.""" + cmd = build_posix_cleanup_command( + "/tmp-data/airflow-ssh-jobs/job_123", expected_base="/tmp-data/airflow-ssh-jobs" + ) + assert "rm -rf" in cmd + assert "/tmp-data/airflow-ssh-jobs/job_123" in cmd + def test_windows_cleanup(self): """Test Windows cleanup command.""" cmd = build_windows_cleanup_command("$env:TEMP\\airflow-ssh-jobs\\job_123") @@ -409,6 +417,13 @@ def test_posix_cleanup_rejects_invalid_path(self): with pytest.raises(ValueError, match="Invalid job directory"): build_posix_cleanup_command("/tmp/other_dir") + def test_posix_cleanup_rejects_path_outside_custom_base(self): + """Test POSIX cleanup rejects paths outside a custom base directory.""" + with pytest.raises(ValueError, match="Invalid job directory"): + build_posix_cleanup_command( + "/tmp/airflow-ssh-jobs/job_123", expected_base="/tmp-data/airflow-ssh-jobs" + ) + def test_windows_cleanup_rejects_invalid_path(self): """Test Windows cleanup rejects paths outside expected base directory.""" with pytest.raises(ValueError, match="Invalid job directory"): From 9c101745771211330a1ec22832fa80f10897ad71 Mon Sep 17 00:00:00 2001 From: axelray-dev <110029405+axelray-dev@users.noreply.github.com> Date: Tue, 14 Jul 2026 02:54:42 +0800 Subject: [PATCH 2/4] Add Windows custom-base SSH cleanup tests (#69813) --- .../ssh/tests/unit/ssh/utils/test_remote_job.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/providers/ssh/tests/unit/ssh/utils/test_remote_job.py b/providers/ssh/tests/unit/ssh/utils/test_remote_job.py index b79d2f6129012..51a553f230520 100644 --- a/providers/ssh/tests/unit/ssh/utils/test_remote_job.py +++ b/providers/ssh/tests/unit/ssh/utils/test_remote_job.py @@ -412,6 +412,15 @@ def test_windows_cleanup(self): assert "Remove-Item" in decoded_script assert "-Recurse" in decoded_script + def test_windows_cleanup_with_custom_base(self): + """Test Windows cleanup under a custom base directory.""" + cmd = build_windows_cleanup_command( + "$env:TEMP\\custom-airflow-ssh-jobs\\job_123", + expected_base="$env:TEMP\\custom-airflow-ssh-jobs", + ) + assert "powershell.exe" in cmd + assert "-EncodedCommand" in cmd + def test_posix_cleanup_rejects_invalid_path(self): """Test POSIX cleanup rejects paths outside expected base directory.""" with pytest.raises(ValueError, match="Invalid job directory"): @@ -428,3 +437,11 @@ def test_windows_cleanup_rejects_invalid_path(self): """Test Windows cleanup rejects paths outside expected base directory.""" with pytest.raises(ValueError, match="Invalid job directory"): build_windows_cleanup_command("C:\\temp\\other_dir") + + def test_windows_cleanup_rejects_path_outside_custom_base(self): + """Test Windows cleanup rejects paths outside a custom base directory.""" + with pytest.raises(ValueError, match="Invalid job directory"): + build_windows_cleanup_command( + "$env:TEMP\\airflow-ssh-jobs\\job_123", + expected_base="$env:TEMP\\custom-airflow-ssh-jobs", + ) From b63ef25b05a622920b3b68a965268fc853627cb9 Mon Sep 17 00:00:00 2001 From: axelray-dev <110029405+axelray-dev@users.noreply.github.com> Date: Tue, 14 Jul 2026 09:04:03 +0800 Subject: [PATCH 3/4] Harden SSH remote cleanup validation for custom bases (#69813) --- .../providers/ssh/operators/ssh_remote_job.py | 14 ++++-- .../airflow/providers/ssh/utils/remote_job.py | 4 +- .../unit/ssh/operators/test_ssh_remote_job.py | 43 +++++++++++++++++++ .../tests/unit/ssh/utils/test_remote_job.py | 10 +++++ 4 files changed, 65 insertions(+), 6 deletions(-) diff --git a/providers/ssh/src/airflow/providers/ssh/operators/ssh_remote_job.py b/providers/ssh/src/airflow/providers/ssh/operators/ssh_remote_job.py index 96a310be9cf42..a2dc781c01ffc 100644 --- a/providers/ssh/src/airflow/providers/ssh/operators/ssh_remote_job.py +++ b/providers/ssh/src/airflow/providers/ssh/operators/ssh_remote_job.py @@ -452,10 +452,16 @@ def _cleanup_remote_job(self, job_dir: str, remote_os: str) -> None: """ self.log.info("Cleaning up remote job directory: %s", job_dir) expected_base = self._paths.base_dir if self._paths else self.remote_base_dir - if remote_os == "posix": - cleanup_cmd = build_posix_cleanup_command(job_dir, expected_base=expected_base) - else: - cleanup_cmd = build_windows_cleanup_command(job_dir, expected_base=expected_base) + try: + if remote_os == "posix": + cleanup_cmd = build_posix_cleanup_command(job_dir, expected_base=expected_base) + else: + cleanup_cmd = build_windows_cleanup_command(job_dir, expected_base=expected_base) + except ValueError as e: + self.log.warning( + "Cleanup validation failed; leaving remote job directory %s in place: %s", job_dir, e + ) + return last_error: Exception | None = None for attempt in range(1, self.cleanup_retries + 1): diff --git a/providers/ssh/src/airflow/providers/ssh/utils/remote_job.py b/providers/ssh/src/airflow/providers/ssh/utils/remote_job.py index 35bd63da6e816..82cd9fdaec501 100644 --- a/providers/ssh/src/airflow/providers/ssh/utils/remote_job.py +++ b/providers/ssh/src/airflow/providers/ssh/utils/remote_job.py @@ -44,10 +44,10 @@ def _validate_job_dir( :raises ValueError: If job_dir doesn't start with the expected base path """ if remote_os == "posix": - expected_base = expected_base or POSIX_DEFAULT_BASE_DIR + expected_base = POSIX_DEFAULT_BASE_DIR if expected_base is None else expected_base expected_prefix = expected_base + "/" else: - expected_base = expected_base or WINDOWS_DEFAULT_BASE_DIR + expected_base = WINDOWS_DEFAULT_BASE_DIR if expected_base is None else expected_base expected_prefix = expected_base + "\\" if not job_dir.startswith(expected_prefix): diff --git a/providers/ssh/tests/unit/ssh/operators/test_ssh_remote_job.py b/providers/ssh/tests/unit/ssh/operators/test_ssh_remote_job.py index 6ecf80baef28d..dbb9b79b821f5 100644 --- a/providers/ssh/tests/unit/ssh/operators/test_ssh_remote_job.py +++ b/providers/ssh/tests/unit/ssh/operators/test_ssh_remote_job.py @@ -333,6 +333,35 @@ def test_execute_complete_with_cleanup(self): call_args = self.mock_hook.exec_ssh_client_command.call_args assert "rm -rf" in call_args[0][1] + def test_execute_complete_with_cleanup_custom_remote_base_dir(self): + """Test execute_complete performs cleanup under a custom remote base directory.""" + op = SSHRemoteJobOperator( + task_id="test_task", + ssh_conn_id="test_conn", + command="/path/to/script.sh", + cleanup="on_success", + remote_base_dir="/tmp-data/airflow-ssh-jobs", + ) + + event = { + "done": True, + "status": "success", + "exit_code": 0, + "job_id": "test_job_123", + "job_dir": "/tmp-data/airflow-ssh-jobs/test_job_123", + "log_file": "/tmp-data/airflow-ssh-jobs/test_job_123/stdout.log", + "exit_code_file": "/tmp-data/airflow-ssh-jobs/test_job_123/exit_code", + "log_chunk": "", + "log_offset": 0, + "remote_os": "posix", + } + + op.execute_complete({}, event) + + self.mock_hook.exec_ssh_client_command.assert_called_once() + call_args = self.mock_hook.exec_ssh_client_command.call_args + assert "rm -rf '/tmp-data/airflow-ssh-jobs/test_job_123'" == call_args[0][1] + def test_on_kill(self): """Test on_kill attempts to kill remote process.""" op = SSHRemoteJobOperator( @@ -460,3 +489,17 @@ def test_cleanup_gives_up_after_retries_without_raising(self): op._cleanup_remote_job("/tmp/airflow-ssh-jobs/test_job_123", "posix") assert self.mock_hook.exec_ssh_client_command.call_count == 3 + + def test_cleanup_validation_failure_leaves_directory_without_raising(self, caplog): + """A cleanup validation failure leaves the remote directory in place.""" + op = SSHRemoteJobOperator( + task_id="test_task", + ssh_conn_id="test_conn", + command="/path/to/script.sh", + remote_base_dir="/tmp-data/airflow-ssh-jobs", + ) + + op._cleanup_remote_job("/tmp/airflow-ssh-jobs/test_job_123", "posix") + + self.mock_hook.exec_ssh_client_command.assert_not_called() + assert "Cleanup validation failed" in caplog.text diff --git a/providers/ssh/tests/unit/ssh/utils/test_remote_job.py b/providers/ssh/tests/unit/ssh/utils/test_remote_job.py index 51a553f230520..5848c39dc912d 100644 --- a/providers/ssh/tests/unit/ssh/utils/test_remote_job.py +++ b/providers/ssh/tests/unit/ssh/utils/test_remote_job.py @@ -433,6 +433,11 @@ def test_posix_cleanup_rejects_path_outside_custom_base(self): "/tmp/airflow-ssh-jobs/job_123", expected_base="/tmp-data/airflow-ssh-jobs" ) + def test_posix_cleanup_does_not_replace_empty_expected_base_with_default(self): + """Test an empty expected base does not fall back to the default base directory.""" + with pytest.raises(ValueError, match="Invalid job directory"): + build_posix_cleanup_command("/tmp/airflow-ssh-jobs/job_123", expected_base="") + def test_windows_cleanup_rejects_invalid_path(self): """Test Windows cleanup rejects paths outside expected base directory.""" with pytest.raises(ValueError, match="Invalid job directory"): @@ -445,3 +450,8 @@ def test_windows_cleanup_rejects_path_outside_custom_base(self): "$env:TEMP\\airflow-ssh-jobs\\job_123", expected_base="$env:TEMP\\custom-airflow-ssh-jobs", ) + + def test_windows_cleanup_does_not_replace_empty_expected_base_with_default(self): + """Test an empty expected base does not fall back to the default base directory.""" + with pytest.raises(ValueError, match="Invalid job directory"): + build_windows_cleanup_command("$env:TEMP\\airflow-ssh-jobs\\job_123", expected_base="") From ed8c90b2a471943c9e3b2b66a097f78a66beebee Mon Sep 17 00:00:00 2001 From: axelray-dev <110029405+axelray-dev@users.noreply.github.com> Date: Wed, 22 Jul 2026 07:30:54 +0800 Subject: [PATCH 4/4] Fix SSH cleanup empty-base regression test --- providers/ssh/tests/unit/ssh/operators/test_ssh_remote_job.py | 2 +- providers/ssh/tests/unit/ssh/utils/test_remote_job.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/providers/ssh/tests/unit/ssh/operators/test_ssh_remote_job.py b/providers/ssh/tests/unit/ssh/operators/test_ssh_remote_job.py index dbb9b79b821f5..174381f8ff437 100644 --- a/providers/ssh/tests/unit/ssh/operators/test_ssh_remote_job.py +++ b/providers/ssh/tests/unit/ssh/operators/test_ssh_remote_job.py @@ -360,7 +360,7 @@ def test_execute_complete_with_cleanup_custom_remote_base_dir(self): self.mock_hook.exec_ssh_client_command.assert_called_once() call_args = self.mock_hook.exec_ssh_client_command.call_args - assert "rm -rf '/tmp-data/airflow-ssh-jobs/test_job_123'" == call_args[0][1] + assert call_args[0][1] == "rm -rf '/tmp-data/airflow-ssh-jobs/test_job_123'" def test_on_kill(self): """Test on_kill attempts to kill remote process.""" diff --git a/providers/ssh/tests/unit/ssh/utils/test_remote_job.py b/providers/ssh/tests/unit/ssh/utils/test_remote_job.py index 5848c39dc912d..cb3483cd59360 100644 --- a/providers/ssh/tests/unit/ssh/utils/test_remote_job.py +++ b/providers/ssh/tests/unit/ssh/utils/test_remote_job.py @@ -435,8 +435,8 @@ def test_posix_cleanup_rejects_path_outside_custom_base(self): def test_posix_cleanup_does_not_replace_empty_expected_base_with_default(self): """Test an empty expected base does not fall back to the default base directory.""" - with pytest.raises(ValueError, match="Invalid job directory"): - build_posix_cleanup_command("/tmp/airflow-ssh-jobs/job_123", expected_base="") + cmd = build_posix_cleanup_command("/job_123", expected_base="") + assert cmd == "rm -rf '/job_123'" def test_windows_cleanup_rejects_invalid_path(self): """Test Windows cleanup rejects paths outside expected base directory."""