From 0095a4f4144102ac51afd81fc1331fe987a44958 Mon Sep 17 00:00:00 2001 From: Harsh Gupta Date: Wed, 15 Jul 2026 03:07:00 +0530 Subject: [PATCH] Fix SSHRemoteJobOperator cleanup failing with custom remote_base_dir The cleanup command builders validated the job directory against the hardcoded default base directory, so any job started with a custom remote_base_dir ran to completion and then failed its cleanup step with "Invalid job directory ... Expected path under '/tmp/airflow-ssh-jobs'". Thread the operator's configured remote_base_dir through build_posix_cleanup_command / build_windows_cleanup_command into _validate_job_dir, falling back to the OS default when unset, so the safety check still constrains the rm target to the directory the job was actually created under. Signed-off-by: Harsh Gupta --- .../providers/ssh/operators/ssh_remote_job.py | 4 ++-- .../airflow/providers/ssh/utils/remote_job.py | 22 ++++++++++++------ .../unit/ssh/operators/test_ssh_remote_job.py | 23 +++++++++++++++++++ .../tests/unit/ssh/utils/test_remote_job.py | 21 +++++++++++++++++ 4 files changed, 61 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..1c73b50840c70 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,9 +452,9 @@ def _cleanup_remote_job(self, job_dir: str, remote_os: str) -> None: """ self.log.info("Cleaning up remote job directory: %s", job_dir) if remote_os == "posix": - cleanup_cmd = build_posix_cleanup_command(job_dir) + cleanup_cmd = build_posix_cleanup_command(job_dir, base_dir=self.remote_base_dir) else: - cleanup_cmd = build_windows_cleanup_command(job_dir) + cleanup_cmd = build_windows_cleanup_command(job_dir, base_dir=self.remote_base_dir) 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..8a3e5fab1270d 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,22 @@ 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"], base_dir: 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 base_dir: Base directory the job was created under (e.g. the operator's + ``remote_base_dir``). Falls back to the OS-specific default base directory. :raises ValueError: If job_dir doesn't start with the expected base path """ if remote_os == "posix": - expected_prefix = POSIX_DEFAULT_BASE_DIR + "/" + expected_prefix = (base_dir or POSIX_DEFAULT_BASE_DIR) + "/" else: - expected_prefix = WINDOWS_DEFAULT_BASE_DIR + "\\" + expected_prefix = (base_dir or WINDOWS_DEFAULT_BASE_DIR) + "\\" if not job_dir.startswith(expected_prefix): raise ValueError( @@ -452,27 +456,31 @@ 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, base_dir: str | None = None) -> str: """ Build a POSIX command to clean up the job directory. :param job_dir: Path to the job directory + :param base_dir: Base directory the job was created under. Defaults to the + OS-specific default base directory. :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", base_dir) return f"rm -rf '{job_dir}'" -def build_windows_cleanup_command(job_dir: str) -> str: +def build_windows_cleanup_command(job_dir: str, base_dir: str | None = None) -> str: """ Build a PowerShell command to clean up the job directory. :param job_dir: Path to the job directory + :param base_dir: Base directory the job was created under. Defaults to the + OS-specific default base directory. :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", base_dir) 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/operators/test_ssh_remote_job.py b/providers/ssh/tests/unit/ssh/operators/test_ssh_remote_job.py index 6ecf80baef28d..f46541cd9be73 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 @@ -460,3 +460,26 @@ 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_with_custom_remote_base_dir(self): + """Cleanup validates the job dir against the operator's remote_base_dir, not the default. + + Regression test for https://github.com/apache/airflow/issues/69813: with a custom + remote_base_dir the job ran fine but cleanup raised ValueError because the job dir + was validated against the default base directory. + """ + self.mock_hook.exec_ssh_client_command.return_value = (0, b"", b"") + + op = SSHRemoteJobOperator( + task_id="test_task", + ssh_conn_id="test_conn", + command="/path/to/script.sh", + remote_base_dir="/tmp-data/airflow-ssh-jobs", + ) + + # Must not raise; before the fix this failed with "Invalid job directory". + op._cleanup_remote_job("/tmp-data/airflow-ssh-jobs/test_job_123", "posix") + + assert self.mock_hook.exec_ssh_client_command.call_count == 1 + cleanup_cmd = self.mock_hook.exec_ssh_client_command.call_args.args[1] + assert "rm -rf '/tmp-data/airflow-ssh-jobs/test_job_123'" in cleanup_cmd 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..a8a39362dd0ce 100644 --- a/providers/ssh/tests/unit/ssh/utils/test_remote_job.py +++ b/providers/ssh/tests/unit/ssh/utils/test_remote_job.py @@ -413,3 +413,24 @@ 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_posix_cleanup_accepts_custom_base_dir(self): + """Test POSIX cleanup accepts job dirs under a custom base directory.""" + cmd = build_posix_cleanup_command("/data/airflow-jobs/job_123", base_dir="/data/airflow-jobs") + assert "rm -rf" in cmd + assert "/data/airflow-jobs/job_123" in cmd + + def test_windows_cleanup_accepts_custom_base_dir(self): + """Test Windows cleanup accepts job dirs under a custom base directory.""" + cmd = build_windows_cleanup_command("D:\\airflow-jobs\\job_123", base_dir="D:\\airflow-jobs") + assert "powershell.exe" in cmd + + def test_posix_cleanup_custom_base_dir_rejects_outside_paths(self): + """With a custom base dir, paths outside it (even the default) are rejected.""" + with pytest.raises(ValueError, match="Invalid job directory"): + build_posix_cleanup_command("/tmp/airflow-ssh-jobs/job_123", base_dir="/data/airflow-jobs") + + def test_windows_cleanup_custom_base_dir_rejects_outside_paths(self): + """With a custom base dir, paths outside it (even the default) are rejected.""" + with pytest.raises(ValueError, match="Invalid job directory"): + build_windows_cleanup_command("$env:TEMP\\airflow-ssh-jobs\\job_123", base_dir="D:\\airflow-jobs")