Repository: apache/airflow
Pull request: [Airflow 16364] Add conn_timeout and cmd_timeout params to SSHOperator; add conn_timeout param to SSHHook
closes #16364 (Timeout is ambiguous in SSHHook and SSHOperator)

Add conn_timeout and cmd_timeout params to SSHOperator. Add conn_timeout param to SSHHook. The legacy timeout param remains but is marked as deprecated. The new params take precedence over the legacy timeout one if both/all are provided.

---
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/CONTRIBUTING.rst#pull-request-guidelines)** for more information.
In case of fundamental code change, Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvements+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in [UPDATING.md](https://github.com/apache/airflow/blob/main/UPDATING.md).


Problem statement:
Timeout is ambiguous in SSHHook and SSHOperator
In SSHHook the timeout argument of the constructor is used to set a connection timeout. This is fine.

But in SSHOperator the timeout argument of the constructor is used for *both* the timeout of the SSHHook *and* the timeout of the command itself (see paramiko's ssh client exec_command use of the timeout parameter). This ambiguous use of the same parameter is very dirty.

I see two ways to clean the behaviour: 

1. Let the SSHHook constructor be the only way to handle the connection timeout (thus, if one wants a specific timeout they should explicitely build a hook to be passed to the operator using the operator's constructor).
2. Split the timeout argument in SSHOperator into two arguments conn_timeout and cmd_timeout for example.

The choice between 1 and 2 depends on how frequently people are supposed to want to change the connection timeout. If it is something very frequent. then go for 2. if not go for 1.

BR and thanks for the code!

Hints from issue discussion:
I feel Option 2 is clearer regardless, since it is still quite confusing if `SSHHook(timeout=...)` and `SSHOperator(timeout=...)` mean two different things. It’s probably better to be explicit about what the timeout is for:

1. Create `conn_timeout` and pass it to `SSHCHook(timeout)`.
2. Create `cmd_timeout` for the command line timeout.
3. Keep `SSHOperator(timeout)`; if it’s passed, set `conn_timeout` and `cmd_timeout` to `timeout` if they are `None`. `TypeError` if `timeout` is all three are passed.
4. (Optional) Deprecate `SSHOperator(timeout)` in favour of `conn_timeout` and `cmd_timeout`.
Agree and I think point 4. should not be optional. We should deprecate it. 

One more thing - I noticed that "timeout" (among others) options for sshHook which also should be updated:

```
                if "timeout" in extra_options:
                    self.timeout = int(extra_options["timeout"], 10)
```

I though also that we could convert the SSHHook to be configurable via UI changing the `extra__ssh__` names similarly to #15159 (and follow up fix #16388). However, unfortunately this is a bit more complex, there is a small caveat with this - we have SFTPHook that derives from SSHHook and this means that we could not immediately reuse the logic of the fields defined there (`extras__ssh__cmd_timeout` vs. `extras__sftp__cmd_timeout`) .

I think it's not worth for this change, but we should probably reimplement the way the extra UI fields are defined to support inheritance better.
@ashb See ^^ I think you mentioned you had some idea that we should get rid of the `__extra__CONN_TYPE__ prefix` ? correct?  It indeed makes things complex for inheritance case and we have about half-half connections now - the UI ones are implementing __extra__CONN_TYPE__ prefix but there are a number of them that have "bare" extras. 
> @ashb See ^^ I think you mentioned you had some idea that we should get rid of the `__extra__CONN_TYPE__ prefix` ? correct? It indeed makes things complex for inheritance case and we have about half-half connections now - the UI ones are implementing **extra__CONN_TYPE** prefix but there are a number of them that have "bare" extras.

@potiuk Yeah, I always just found it annoying to have the prefix on it, mostly when setting the connection via a URI. Are you saying this issue means we should keep the prefix, or you would like to remove it too?
> @potiuk Yeah, I always just found it annoying to have the prefix on it, mostly when setting the connection via a URI. Are you saying this issue means we should keep the prefix, or you would like to remove it too?

SSH/SFTP does not have the prefix (it does not support UI fields for extras). Since long term remove the prefix (I agree it's annoying)- we should not touch it here then and keep it without the prefix and implement removal of the prefix later (and then we can add UI support for SSH/SFTP). 

So summarizing @NBardelot : just don't forget to add the "cmd_timeout" and "conn_timeout" to extras (without prefix).
Do we need `cmd_timeout` on the Hook? Isn't that more a property that only makes sense for the operator because it changes from task-to-task, and isn't anything to do with the Connection.
> Do we need `cmd_timeout` on the Hook? Isn't that more a property that only makes sense for the operator because it changes from task-to-task, and isn't anything to do with the Connection.

I think both should be in Hook/connection - acting as default value for all SSH commands. It could be done via dag defaults as well, but both "cmd_timeout" and "timeout" are a bit too generic IMHO for Dag defaults, so it would have to be "ssh_command_timeout" or something like tha. In this case, it rather belongs to connection.
@potiuk or @ashb - Can someone please assign to me if no one else is already working on it? Thanks!
Please!

Just to be clear on the preferred usage going forward:
- SSHOperator uses two different timeouts:
   - A connection timeout passed into SSHHook (see below)
   - A command timeout passed into paramiko.SSHClient.exec_command() which sets the command’s channel timeout
- SSHHook uses a single timeout:
  - A connection timeout passed into paramiko.SSHClient.connect() which is a timeout for the TCP connect.

It doesn't seem to make sense to accept a cmd_timeout in the hook, so I am planning to do the following:
- Accept both conn_timeout and cmd_timeout in SSHOperator.
- Only accept conn_timeout in SSHHook.
- Add a deprecation warning to SSHHook for the use of timeout in extra_options; favor conn_timeout if provided.

Commit message: Merge branch 'apache:main' into airflow-16364

Patch to review:
diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index a8994bdece1b7..5f9c637b9d18b 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -56,8 +56,10 @@ class SSHHook(BaseHook):
     :type key_file: str
     :param port: port of remote host to connect (Default is paramiko SSH_PORT)
     :type port: int
-    :param timeout: timeout for the attempt to connect to the remote_host.
+    :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host.
     :type timeout: int
+    :param conn_timeout: timeout for the attempt to connect to the remote_host.
+    :type conn_timeout: int
     :param keepalive_interval: send a keepalive packet to remote host every
         keepalive_interval seconds
     :type keepalive_interval: int
@@ -101,7 +103,8 @@ def __init__(
         password: Optional[str] = None,
         key_file: Optional[str] = None,
         port: Optional[int] = None,
-        timeout: int = 10,
+        timeout: Optional[int] = 10,
+        conn_timeout: Optional[int] = None,
         keepalive_interval: int = 30,
     ) -> None:
         super().__init__()
@@ -113,6 +116,7 @@ def __init__(
         self.pkey = None
         self.port = port
         self.timeout = timeout
+        self.conn_timeout = conn_timeout
         self.keepalive_interval = keepalive_interval
 
         # Default values, overridable from Connection
@@ -137,6 +141,7 @@ def __init__(
                 self.remote_host = conn.host
             if self.port is None:
                 self.port = conn.port
+
             if conn.extra is not None:
                 extra_options = conn.extra_dejson
                 if "key_file" in extra_options and self.key_file is None:
@@ -148,7 +153,17 @@ def __init__(
                     self.pkey = self._pkey_from_private_key(private_key, passphrase=private_key_passphrase)
 
                 if "timeout" in extra_options:
-                    self.timeout = int(extra_options["timeout"], 10)
+                    warnings.warn(
+                        'Extra option `timeout` is deprecated.'
+                        'Please use `conn_timeout` instead.'
+                        'The old option `timeout` will be removed in a future version.',
+                        DeprecationWarning,
+                        stacklevel=2,
+                    )
+                    self.timeout = int(extra_options.get('timeout', 10))
+
+                if "conn_timeout" in extra_options:
+                    self.conn_timeout = int(extra_options.get('conn_timeout'))
 
                 if "compress" in extra_options and str(extra_options["compress"]).lower() == 'false':
                     self.compress = False
@@ -185,6 +200,18 @@ def __init__(
                     self.host_key = key_constructor(data=decoded_host_key)
                     self.no_host_key_check = False
 
+        if self.timeout:
+            warnings.warn(
+                'Parameter `timeout` is deprecated.'
+                'Please use `conn_timeout` instead.'
+                'The old option `timeout` will be removed in a future version.',
+                DeprecationWarning,
+                stacklevel=1,
+            )
+
+        if self.conn_timeout is None:
+            self.conn_timeout = self.timeout
+
         if self.pkey and self.key_file:
             raise AirflowException(
                 "Params key_file and private_key both provided.  Must provide no more than one."
@@ -253,7 +280,7 @@ def get_conn(self) -> paramiko.SSHClient:
         connect_kwargs = dict(
             hostname=self.remote_host,
             username=self.username,
-            timeout=self.timeout,
+            timeout=self.conn_timeout,
             compress=self.compress,
             port=self.port,
             sock=self.host_proxy,
diff --git a/airflow/providers/ssh/operators/ssh.py b/airflow/providers/ssh/operators/ssh.py
index c82e0e12c30a4..27380977a2abb 100644
--- a/airflow/providers/ssh/operators/ssh.py
+++ b/airflow/providers/ssh/operators/ssh.py
@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import warnings
 from base64 import b64encode
 from select import select
 from typing import Optional, Union
@@ -43,8 +44,12 @@ class SSHOperator(BaseOperator):
     :type remote_host: str
     :param command: command to execute on remote host. (templated)
     :type command: str
-    :param timeout: timeout (in seconds) for executing the command. The default is 10 seconds.
+    :param timeout: (deprecated) timeout (in seconds) for executing the command. The default is 10 seconds.
     :type timeout: int
+    :param conn_timeout: timeout (in seconds) for maintaining the connection. The default is 10 seconds.
+    :type conn_timeout: int
+    :param cmd_timeout: timeout (in seconds) for executing the command. The default is 10 seconds.
+    :type cmd_timeout: int
     :param environment: a dict of shell environment variables. Note that the
         server will reject them silently if `AcceptEnv` is not set in SSH config.
     :type environment: dict
@@ -66,7 +71,9 @@ def __init__(
         ssh_conn_id: Optional[str] = None,
         remote_host: Optional[str] = None,
         command: Optional[str] = None,
-        timeout: int = 10,
+        timeout: Optional[int] = 10,
+        conn_timeout: Optional[int] = None,
+        cmd_timeout: Optional[int] = None,
         environment: Optional[dict] = None,
         get_pty: bool = False,
         **kwargs,
@@ -77,9 +84,24 @@ def __init__(
         self.remote_host = remote_host
         self.command = command
         self.timeout = timeout
+        self.conn_timeout = conn_timeout
+        self.cmd_timeout = cmd_timeout
+        if self.conn_timeout is None:
+            self.conn_timeout = self.timeout
+        if self.cmd_timeout is None:
+            self.cmd_timeout = self.timeout
         self.environment = environment
         self.get_pty = (self.command.startswith('sudo') or get_pty) if self.command else get_pty
 
+        if self.timeout:
+            warnings.warn(
+                'Parameter `timeout` is deprecated.'
+                'Please use `conn_timeout` and `cmd_timeout` instead.'
+                'The old option `timeout` will be removed in a future version.',
+                DeprecationWarning,
+                stacklevel=1,
+            )
+
     def execute(self, context) -> Union[bytes, str, bool]:
         try:
             if self.ssh_conn_id:
@@ -89,7 +111,7 @@ def execute(self, context) -> Union[bytes, str, bool]:
                     self.log.info(
                         "ssh_hook is not provided or invalid. Trying ssh_conn_id to create SSHHook."
                     )
-                    self.ssh_hook = SSHHook(ssh_conn_id=self.ssh_conn_id, timeout=self.timeout)
+                    self.ssh_hook = SSHHook(ssh_conn_id=self.ssh_conn_id, conn_timeout=self.conn_timeout)
 
             if not self.ssh_hook:
                 raise AirflowException("Cannot operate without ssh_hook or ssh_conn_id.")
@@ -112,7 +134,7 @@ def execute(self, context) -> Union[bytes, str, bool]:
                 stdin, stdout, stderr = ssh_client.exec_command(
                     command=self.command,
                     get_pty=self.get_pty,
-                    timeout=self.timeout,
+                    timeout=self.cmd_timeout,
                     environment=self.environment,
                 )
                 # get channels
@@ -133,7 +155,7 @@ def execute(self, context) -> Union[bytes, str, bool]:
 
                 # read from both stdout and stderr
                 while not channel.closed or channel.recv_ready() or channel.recv_stderr_ready():
-                    readq, _, _ = select([channel], [], [], self.timeout)
+                    readq, _, _ = select([channel], [], [], self.cmd_timeout)
                     for recv in readq:
                         if recv.recv_ready():
                             line = stdout.channel.recv(len(recv.in_buffer))
diff --git a/tests/providers/ssh/hooks/test_ssh.py b/tests/providers/ssh/hooks/test_ssh.py
index 70ba4eb0ce642..4afa6fb12e060 100644
--- a/tests/providers/ssh/hooks/test_ssh.py
+++ b/tests/providers/ssh/hooks/test_ssh.py
@@ -68,6 +68,9 @@ def generate_host_key(pkey: paramiko.PKey):
 TEST_PKEY_ECDSA = paramiko.ECDSAKey.generate()
 TEST_PRIVATE_KEY_ECDSA = generate_key_string(pkey=TEST_PKEY_ECDSA)
 
+TEST_TIMEOUT = 20
+TEST_CONN_TIMEOUT = 30
+
 PASSPHRASE = ''.join(random.choice(string.ascii_letters) for i in range(10))
 TEST_ENCRYPTED_PRIVATE_KEY = generate_key_string(pkey=TEST_PKEY, passphrase=PASSPHRASE)
 
@@ -76,6 +79,8 @@ class TestSSHHook(unittest.TestCase):
     CONN_SSH_WITH_PRIVATE_KEY_EXTRA = 'ssh_with_private_key_extra'
     CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA = 'ssh_with_private_key_ecdsa_extra'
     CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA = 'ssh_with_private_key_passphrase_extra'
+    CONN_SSH_WITH_TIMEOUT_EXTRA = 'ssh_with_timeout_extra'
+    CONN_SSH_WITH_CONN_TIMEOUT_EXTRA = 'ssh_with_conn_timeout_extra'
     CONN_SSH_WITH_EXTRA = 'ssh_with_extra'
     CONN_SSH_WITH_EXTRA_FALSE_LOOK_FOR_KEYS = 'ssh_with_extra_false_look_for_keys'
     CONN_SSH_WITH_HOST_KEY_EXTRA = 'ssh_with_host_key_extra'
@@ -91,6 +96,8 @@ def tearDownClass(cls) -> None:
                 cls.CONN_SSH_WITH_PRIVATE_KEY_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA,
+                cls.CONN_SSH_WITH_TIMEOUT_EXTRA,
+                cls.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA,
                 cls.CONN_SSH_WITH_EXTRA,
                 cls.CONN_SSH_WITH_HOST_KEY_EXTRA,
                 cls.CONN_SSH_WITH_HOST_KEY_EXTRA_WITH_TYPE,
@@ -126,11 +133,7 @@ def setUpClass(cls) -> None:
                 conn_id=cls.CONN_SSH_WITH_PRIVATE_KEY_EXTRA,
                 host='localhost',
                 conn_type='ssh',
-                extra=json.dumps(
-                    {
-                        "private_key": TEST_PRIVATE_KEY,
-                    }
-                ),
+                extra=json.dumps({"private_key": TEST_PRIVATE_KEY}),
             )
         )
         db.merge_conn(
@@ -151,6 +154,22 @@ def setUpClass(cls) -> None:
                 extra=json.dumps({"private_key": TEST_PRIVATE_KEY_ECDSA}),
             )
         )
+        db.merge_conn(
+            Connection(
+                conn_id=cls.CONN_SSH_WITH_TIMEOUT_EXTRA,
+                host='localhost',
+                conn_type='ssh',
+                extra=json.dumps({"timeout": TEST_TIMEOUT}),
+            )
+        )
+        db.merge_conn(
+            Connection(
+                conn_id=cls.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA,
+                host='localhost',
+                conn_type='ssh',
+                extra=json.dumps({"conn_timeout": TEST_CONN_TIMEOUT}),
+            )
+        )
         db.merge_conn(
             Connection(
                 conn_id=cls.CONN_SSH_WITH_HOST_KEY_EXTRA,
@@ -475,6 +494,98 @@ def test_ssh_connection_with_no_host_key_where_no_host_key_check_is_false(self,
             assert ssh_client.return_value.connect.called is True
             assert ssh_client.return_value.get_host_keys.return_value.add.called is False
 
+    @mock.patch('airflow.providers.ssh.hooks.ssh.paramiko.SSHClient')
+    def test_ssh_connection_with_conn_timeout(self, ssh_mock):
+        hook = SSHHook(
+            remote_host='remote_host',
+            port='port',
+            username='username',
+            password='password',
+            conn_timeout=20,
+            key_file='fake.file',
+        )
+
+        with hook.get_conn():
+            ssh_mock.return_value.connect.assert_called_once_with(
+                hostname='remote_host',
+                username='username',
+                password='password',
+                key_filename='fake.file',
+                timeout=20,
+                compress=True,
+                port='port',
+                sock=None,
+                look_for_keys=True,
+            )
+
+    @mock.patch('airflow.providers.ssh.hooks.ssh.paramiko.SSHClient')
+    def test_ssh_connection_with_conn_timeout_and_timeout(self, ssh_mock):
+        hook = SSHHook(
+            remote_host='remote_host',
+            port='port',
+            username='username',
+            password='password',
+            timeout=10,
+            conn_timeout=20,
+            key_file='fake.file',
+        )
+
+        with hook.get_conn():
+            ssh_mock.return_value.connect.assert_called_once_with(
+                hostname='remote_host',
+                username='username',
+                password='password',
+                key_filename='fake.file',
+                timeout=20,
+                compress=True,
+                port='port',
+                sock=None,
+                look_for_keys=True,
+            )
+
+    @mock.patch('airflow.providers.ssh.hooks.ssh.paramiko.SSHClient')
+    def test_ssh_connection_with_timeout_extra(self, ssh_mock):
+        hook = SSHHook(
+            ssh_conn_id=self.CONN_SSH_WITH_TIMEOUT_EXTRA,
+            remote_host='remote_host',
+            port='port',
+            username='username',
+            timeout=10,
+        )
+
+        with hook.get_conn():
+            ssh_mock.return_value.connect.assert_called_once_with(
+                hostname='remote_host',
+                username='username',
+                timeout=20,
+                compress=True,
+                port='port',
+                sock=None,
+                look_for_keys=True,
+            )
+
+    @mock.patch('airflow.providers.ssh.hooks.ssh.paramiko.SSHClient')
+    def test_ssh_connection_with_conn_timeout_extra(self, ssh_mock):
+        hook = SSHHook(
+            ssh_conn_id=self.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA,
+            remote_host='remote_host',
+            port='port',
+            username='username',
+            timeout=10,
+            conn_timeout=15,
+        )
+
+        with hook.get_conn():
+            ssh_mock.return_value.connect.assert_called_once_with(
+                hostname='remote_host',
+                username='username',
+                timeout=30,
+                compress=True,
+                port='port',
+                sock=None,
+                look_for_keys=True,
+            )
+
     def test_openssh_private_key(self):
         # Paramiko behaves differently with OpenSSH generated keys to paramiko
         # generated keys, so we need a test one.
diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index e39b8fbb8c156..a73c2087c61fb 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -31,7 +31,9 @@
 
 TEST_DAG_ID = 'unit_tests_ssh_test_op'
 TEST_CONN_ID = "conn_id_for_testing"
-TIMEOUT = 5
+CONN_TIMEOUT = 5
+CMD_TIMEOUT = 7
+TIMEOUT = 12
 DEFAULT_DATE = datetime(2017, 1, 1)
 COMMAND = "echo -n airflow"
 COMMAND_WITH_SUDO = "sudo " + COMMAND
@@ -52,7 +54,7 @@ def setUp(self):
         self.hook = hook
         self.dag = dag
 
-    def test_hook_created_correctly(self):
+    def test_hook_created_correctly_with_timeout(self):
         timeout = 20
         ssh_id = "ssh_default"
         task = SSHOperator(
@@ -62,7 +64,26 @@ def test_hook_created_correctly(self):
 
         task.execute(None)
 
-        assert timeout == task.ssh_hook.timeout
+        assert timeout == task.ssh_hook.conn_timeout
+        assert ssh_id == task.ssh_hook.ssh_conn_id
+
+    def test_hook_created_correctly(self):
+        conn_timeout = 20
+        cmd_timeout = 45
+        ssh_id = "ssh_default"
+        task = SSHOperator(
+            task_id="test",
+            command=COMMAND,
+            dag=self.dag,
+            conn_timeout=conn_timeout,
+            cmd_timeout=cmd_timeout,
+            ssh_conn_id="ssh_default",
+        )
+        assert task is not None
+
+        task.execute(None)
+
+        assert conn_timeout == task.ssh_hook.conn_timeout
         assert ssh_id == task.ssh_hook.ssh_conn_id
 
     @conf_vars({('core', 'enable_xcom_pickling'): 'False'})
@@ -138,7 +159,9 @@ def test_no_output_command(self):
     def test_arg_checking(self):
         # Exception should be raised if neither ssh_hook nor ssh_conn_id is provided
         with pytest.raises(AirflowException, match="Cannot operate without ssh_hook or ssh_conn_id."):
-            task_0 = SSHOperator(task_id="test", command=COMMAND, timeout=TIMEOUT, dag=self.dag)
+            task_0 = SSHOperator(
+                task_id="test", command=COMMAND, conn_timeout=TIMEOUT, cmd_timeout=TIMEOUT, dag=self.dag
+            )
             task_0.execute(None)
 
         # if ssh_hook is invalid/not provided, use ssh_conn_id to create SSHHook
@@ -147,7 +170,8 @@ def test_arg_checking(self):
             ssh_hook="string_rather_than_SSHHook",  # invalid ssh_hook
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -160,7 +184,8 @@ def test_arg_checking(self):
             task_id="test_2",
             ssh_conn_id=TEST_CONN_ID,  # no ssh_hook provided
             command=COMMAND,
-            timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -175,7 +200,8 @@ def test_arg_checking(self):
             ssh_hook=self.hook,
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -184,6 +210,39 @@ def test_arg_checking(self):
             pass
         assert task_3.ssh_hook.ssh_conn_id == self.hook.ssh_conn_id
 
+        # if valid conn_timeout, cmd_timeout, and timeout are provided, ignore timeout
+        task_4 = SSHOperator(
+            task_id="test_4",
+            ssh_conn_id=TEST_CONN_ID,
+            command=COMMAND,
+            timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
+            dag=self.dag,
+        )
+        try:
+            task_4.execute(None)
+        except Exception:
+            pass
+        assert task_4.conn_timeout == CONN_TIMEOUT
+        assert task_4.cmd_timeout == CMD_TIMEOUT
+        assert task_4.ssh_hook.conn_timeout == CONN_TIMEOUT
+
+        # if valid only timeout is provided, use for both conn_timeout and cmd_timeout
+        task_5 = SSHOperator(
+            task_id="test_5",
+            ssh_conn_id=TEST_CONN_ID,
+            command=COMMAND,
+            timeout=TIMEOUT,
+            dag=self.dag,
+        )
+        try:
+            task_5.execute(None)
+        except Exception:
+            pass
+        assert task_5.cmd_timeout == TIMEOUT
+        assert task_5.ssh_hook.conn_timeout == TIMEOUT
+
     @parameterized.expand(
         [
             (COMMAND, False, False),
@@ -198,7 +257,9 @@ def test_get_pyt_set_correctly(self, command, get_pty_in, get_pty_out):
             task_id="test",
             ssh_hook=self.hook,
             command=command,
-            timeout=TIMEOUT,
+            # timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             get_pty=get_pty_in,
             dag=self.dag,
         )


Code review comments:
diff_hunk: @@ -56,8 +56,10 @@ class SSHHook(BaseHook):
     :type key_file: str
     :param port: port of remote host to connect (Default is paramiko SSH_PORT)
     :type port: int
-    :param timeout: timeout for the attempt to connect to the remote_host.
+    :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host.
     :type timeout: int
+    :param conn_timeout: timeout for the attempt to connect to the remote_host.
+    :type conn_timeout: int
path: airflow/providers/ssh/hooks/ssh.py
text: @user1:
We should put this above the deprecated `timeout`

diff_hunk: @@ -101,7 +103,8 @@ def __init__(
         password: Optional[str] = None,
         key_file: Optional[str] = None,
         port: Optional[int] = None,
-        timeout: int = 10,
+        timeout: Optional[int] = 10,
path: airflow/providers/ssh/hooks/ssh.py
text: @user1:
Is this change needed/correct? I feel we want to either change the default to `None` (so we can better detect whether `timeout` is passed in by the user), or the type should continue to be `int` (since the user shouldn’t do `timeout=None` in any situation).

@author:
Yeah, looking at it again you're right. It should default to None otherwise the user will need to purposefully pass timeout = None to turn off the deprecation warning.

In order to maintain backwards compatibility, since the current versions of the hook/operator allow the user to omit timeout, I'll need to handle a situation where no timeout is passed and set the default later in the code.