SSHOperator - Restore ability to override SSHHook cmd_timeout

airflow/providers/ssh/hooks/ssh.py

@@ -471,18 +471,24 @@ def exec_ssh_client_command(
         command: str,
         get_pty: bool,
         environment: dict | None,
-        timeout: int | None = None,
+        timeout: int | ArgNotSet | None = NOTSET,
     ) -> tuple[int, bytes, bytes]:
         self.log.info("Running command: %s", command)
 
-        if timeout is None:
-            timeout = self.cmd_timeout  # type: ignore[assignment]
+        cmd_timeout: int | None
+        if not isinstance(timeout, ArgNotSet):
+            cmd_timeout = timeout
+        elif not isinstance(self.cmd_timeout, ArgNotSet):
+            cmd_timeout = self.cmd_timeout
+        else:
+            cmd_timeout = CMD_TIMEOUT
+        del timeout  # Too easy to confuse with "timedout" below.
 
         # set timeout taken as params
         stdin, stdout, stderr = ssh_client.exec_command(
             command=command,
             get_pty=get_pty,
-            timeout=timeout,
+            timeout=cmd_timeout,
             environment=environment,
         )
         # get channels
@@ -505,8 +511,8 @@ def exec_ssh_client_command(
 
         # read from both stdout and stderr
         while not channel.closed or channel.recv_ready() or channel.recv_stderr_ready():
-            readq, _, _ = select([channel], [], [], timeout)
-            if timeout is not None:
+            readq, _, _ = select([channel], [], [], cmd_timeout)
+            if cmd_timeout is not None:
                 timedout = len(readq) == 0
             for recv in readq:
                 if recv.recv_ready():

airflow/providers/ssh/operators/ssh.py

@@ -143,7 +143,7 @@ def exec_ssh_client_command(self, ssh_client: SSHClient, command: str):
         )
         assert self.ssh_hook
         return self.ssh_hook.exec_ssh_client_command(
-            ssh_client, command, environment=self.environment, get_pty=self.get_pty
+            ssh_client, command, timeout=self.cmd_timeout, environment=self.environment, get_pty=self.get_pty
         )
 
     def raise_for_status(self, exit_status: int, stderr: bytes, context=None) -> None:
@@ -156,7 +156,7 @@ def raise_for_status(self, exit_status: int, stderr: bytes, context=None) -> Non
     def run_ssh_client_command(self, ssh_client: SSHClient, command: str, context=None) -> bytes:
         assert self.ssh_hook
         exit_status, agg_stdout, agg_stderr = self.ssh_hook.exec_ssh_client_command(
-            ssh_client, command, environment=self.environment, get_pty=self.get_pty
+            ssh_client, command, timeout=self.cmd_timeout, environment=self.environment, get_pty=self.get_pty
         )
         self.raise_for_status(exit_status, agg_stderr, context=context)
         return agg_stdout

tests/providers/ssh/hooks/test_ssh.py

@@ -974,24 +974,6 @@ def test_exec_ssh_client_command(self):
             )
             assert ret == (0, b"airflow\n", b"")
 
-    @pytest.mark.flaky(reruns=5)
-    def test_command_timeout_default(self):
-        hook = SSHHook(
-            ssh_conn_id="ssh_default",
-            conn_timeout=30,
-            banner_timeout=100,
-        )
-
-        with hook.get_conn() as client:
-            with pytest.raises(AirflowException):
-                hook.exec_ssh_client_command(
-                    client,
-                    "sleep 10",
-                    False,
-                    None,
-                    1,
-                )
-
     @pytest.mark.flaky(reruns=5)
     def test_command_timeout_success(self):
         hook = SSHHook(
@@ -1028,6 +1010,24 @@ def test_command_timeout_fail(self):
                     None,
                 )
 
+    def test_command_timeout_not_set(self):
+        hook = SSHHook(
+            ssh_conn_id="ssh_default",
+            conn_timeout=30,
+            cmd_timeout=None,
+            banner_timeout=100,
+        )
+
+        with hook.get_conn() as client:
+            # sleeping for 20 sec which is longer than default timeout of 10 seconds
+            # to validate that no timeout is applied
+            hook.exec_ssh_client_command(
+                client,
+                "sleep 20",
+                environment=False,
+                get_pty=None,
+            )
+
     @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient")
     def test_ssh_connection_with_no_host_key_check_true_and_allow_host_key_changes_true(self, ssh_mock):
         hook = SSHHook(ssh_conn_id=self.CONN_SSH_WITH_NO_HOST_KEY_CHECK_TRUE_AND_ALLOW_HOST_KEY_CHANGES_TRUE)

tests/providers/ssh/operators/test_ssh.py

@@ -28,6 +28,7 @@
 from airflow.providers.ssh.hooks.ssh import SSHHook
 from airflow.providers.ssh.operators.ssh import SSHOperator
 from airflow.utils.timezone import datetime
+from airflow.utils.types import NOTSET
 from tests.test_utils.config import conf_vars
 
 TEST_DAG_ID = "unit_tests_ssh_test_op"
@@ -108,7 +109,7 @@ def test_return_value(self, enable_xcom_pickling, output, expected):
             result = task.execute(None)
             assert result == expected
             self.exec_ssh_client_command.assert_called_with(
-                mock.ANY, COMMAND, environment={"TEST": "value"}, get_pty=False
+                mock.ANY, COMMAND, timeout=NOTSET, environment={"TEST": "value"}, get_pty=False
             )
 
     @mock.patch("os.environ", {"AIRFLOW_CONN_" + TEST_CONN_ID.upper(): "ssh://test_id@localhost"})