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 7c126c7c6e7d2..fec1be1536913 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 d8d7d8ff3fd59..5cc0699322ed0 100644 --- a/providers/ssh/src/airflow/providers/ssh/utils/remote_job.py +++ b/providers/ssh/src/airflow/providers/ssh/utils/remote_job.py @@ -30,23 +30,26 @@ 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: The base directory job_dir is expected to be under. Defaults to the + standard POSIX/Windows base directory, but callers configured with a custom + ``remote_base_dir`` must pass it here so job_dir is validated against the base + directory that was actually used to build it. :raises ValueError: If job_dir doesn't start with the expected base path """ - if remote_os == "posix": - expected_prefix = POSIX_DEFAULT_BASE_DIR + "/" - else: - expected_prefix = WINDOWS_DEFAULT_BASE_DIR + "\\" + if base_dir is None: + base_dir = POSIX_DEFAULT_BASE_DIR if remote_os == "posix" else WINDOWS_DEFAULT_BASE_DIR + expected_prefix = base_dir + ("\\" if remote_os == "windows" else "/") if not job_dir.startswith(expected_prefix): - raise ValueError( - f"Invalid job directory '{job_dir}'. Expected path under '{expected_prefix[:-1]}' for safety." - ) + raise ValueError(f"Invalid job directory '{job_dir}'. Expected path under '{base_dir}' for safety.") def _validate_env_var_name(name: str) -> None: @@ -456,27 +459,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: The configured base directory job_dir is expected to be under + (see :func:`_validate_job_dir`). Defaults to the standard 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=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: The configured base directory job_dir is expected to be under + (see :func:`_validate_job_dir`). Defaults to the standard 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=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 55e144bab2727..ceb5511d30811 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 @@ -393,6 +393,43 @@ 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_and_custom_remote_base_dir(self): + """ + Cleanup must validate job_dir against the operator's configured + remote_base_dir, not the hardcoded default. Regression test for a bug + where a custom remote_base_dir made execution/monitoring succeed but + cleanup always raised ValueError, since it only accepted job_dir under + the default base directory. + """ + op = SSHRemoteJobOperator( + task_id="test_task", + ssh_conn_id="test_conn", + command="/path/to/script.sh", + cleanup="on_success", + remote_base_dir="/opt/custom-airflow-jobs", + ) + op.remote_base_dir = "/opt/custom-airflow-jobs" # post-templating value + + event = { + "done": True, + "status": "success", + "exit_code": 0, + "job_id": "test_job_123", + "job_dir": "/opt/custom-airflow-jobs/test_job_123", + "log_file": "/opt/custom-airflow-jobs/test_job_123/stdout.log", + "exit_code_file": "/opt/custom-airflow-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" in call_args[0][1] + assert "/opt/custom-airflow-jobs/test_job_123" in call_args[0][1] + def test_on_kill(self): """Test on_kill attempts to kill remote process.""" op = SSHRemoteJobOperator( 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 53c1f6f80b455..2dadebaa366dd 100644 --- a/providers/ssh/tests/unit/ssh/utils/test_remote_job.py +++ b/providers/ssh/tests/unit/ssh/utils/test_remote_job.py @@ -453,3 +453,18 @@ 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): + """Cleanup must validate against the base_dir that was actually configured.""" + cmd = build_posix_cleanup_command("/custom/base/job_123", base_dir="/custom/base") + assert "/custom/base/job_123" in cmd + + def test_posix_cleanup_rejects_default_dir_when_custom_base_dir_configured(self): + """A job_dir under the default base must not pass when a custom base_dir was used.""" + with pytest.raises(ValueError, match="Invalid job directory"): + build_posix_cleanup_command("/tmp/airflow-ssh-jobs/job_123", base_dir="/custom/base") + + def test_windows_cleanup_accepts_custom_base_dir(self): + """Cleanup must validate against the base_dir that was actually configured.""" + cmd = build_windows_cleanup_command("C:\\custom\\base\\job_123", base_dir="C:\\custom\\base") + assert "powershell.exe" in cmd