jedcunningham commented on a change in pull request #17236:
URL: https://github.com/apache/airflow/pull/17236#discussion_r683689193
##########
File path: tests/providers/ssh/hooks/test_ssh.py
##########
@@ -475,6 +504,122 @@ 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,
+ )
+
+ # conn_timeout parameter wins over extra options
+ with hook.get_conn():
+ ssh_mock.return_value.connect.assert_called_once_with(
+ hostname='remote_host',
+ username='username',
+ timeout=15,
+ 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_and_conn_timeout_extra(self,
ssh_mock):
+ hook = SSHHook(
+ ssh_conn_id=self.CONN_SSH_WITH_TIMEOUT_AND_CONN_TIMEOUT_EXTRA,
+ remote_host='remote_host',
+ port='port',
+ username='username',
+ timeout=10,
+ conn_timeout=15,
+ )
+
+ # conn_timeout parameter wins over extra options
+ with hook.get_conn():
+ ssh_mock.return_value.connect.assert_called_once_with(
+ hostname='remote_host',
+ username='username',
+ timeout=15,
+ compress=True,
+ port='port',
+ sock=None,
+ look_for_keys=True,
+ )
Review comment:
Aren't we missing a test where neither `timeout` and `conn_timeout` are
passed and we use the value from the connection extras? Also missing a test
where the default is used if the connection doesn't have one set?
##########
File path: airflow/providers/ssh/operators/ssh.py
##########
@@ -43,7 +46,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 conn_timeout: timeout (in seconds) for maintaining the connection.
The default is 10 seconds.
Review comment:
Maybe the same language here as is used on the `remote_host`? "If
provided, it will replace..."
##########
File path: airflow/providers/ssh/operators/ssh.py
##########
@@ -77,9 +87,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.timeout else
TIMEOUT_DEFAULT
Review comment:
I think by setting a default here, we effectively break allowing the
connections `conn_timeout` to be used over the default. Maybe this would be
better? Btw, I think we need better test coverage on default / extra handling
like the hook too.
```suggestion
if self.conn_timeout is None and self.timeout:
self.conn_timeout = self.timeout
```
This'll let `None` be passed into the hook, which will let the connection
extras be used and we avoid having the "default" set in 2 places.
##########
File path: docs/apache-airflow-providers-ssh/connections/ssh.rst
##########
@@ -47,7 +47,8 @@ Extra (optional)
* ``key_file`` - Full Path of the private SSH Key file that will be used
to connect to the remote_host.
* ``private_key`` - Content of the private key used to connect to the
remote_host.
* ``private_key_passphrase`` - Content of the private key passphrase used
to decrypt the private key.
- * ``timeout`` - An optional timeout (in seconds) for the TCP connect.
Default is ``10``.
+ * ``conn_timeout`` - An optional timeout (in seconds) for the TCP connect.
Default is ``10``. The paramater to
:class:`~airflow.providers.ssh.hooks.ssh.SSHHook` conn_timeout takes precedence.
Review comment:
Maybe we should document on the other side instead of here, like
`remote_host`? Thoughts?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
