Fix SSHRemoteJobOperator custom remote_base_dir cleanup validation

providers/ssh/src/airflow/providers/ssh/operators/ssh_remote_job.py

@@ -451,10 +451,17 @@ 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)
-        if remote_os == "posix":
-            cleanup_cmd = build_posix_cleanup_command(job_dir)
-        else:
-            cleanup_cmd = build_windows_cleanup_command(job_dir)
+        expected_base = self._paths.base_dir if self._paths else self.remote_base_dir
+        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):

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 = POSIX_DEFAULT_BASE_DIR if expected_base is None else expected_base
+        expected_prefix = expected_base + "/"
     else:
-        expected_prefix = 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):
         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")

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 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."""
         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

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")
@@ -404,12 +412,46 @@ 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"):
             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_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."""
+        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."""
         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",
+            )
+
+    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="")