diff --git a/providers/src/airflow/providers/ssh/hooks/ssh.py b/providers/src/airflow/providers/ssh/hooks/ssh.py index e5ce9ec27cfe5..ecd3990d06b9a 100644 --- a/providers/src/airflow/providers/ssh/hooks/ssh.py +++ b/providers/src/airflow/providers/ssh/hooks/ssh.py @@ -74,6 +74,7 @@ class SSHHook(BaseHook): iterable of algorithm identifiers, which will be disabled for the lifetime of the transport :param ciphers: list of ciphers to use in order of preference + :param auth_timeout: timeout (in seconds) for the attempt to authenticate with the remote_host """ # List of classes to try loading private keys as, ordered (roughly) by most common to least common @@ -121,6 +122,7 @@ def __init__( banner_timeout: float = 30.0, disabled_algorithms: dict | None = None, ciphers: list[str] | None = None, + auth_timeout: int | None = None, ) -> None: super().__init__() self.ssh_conn_id = ssh_conn_id @@ -138,6 +140,7 @@ def __init__( self.disabled_algorithms = disabled_algorithms self.ciphers = ciphers self.host_proxy_cmd = None + self.auth_timeout = auth_timeout # Default values, overridable from Connection self.compress = True @@ -332,6 +335,7 @@ def get_conn(self) -> paramiko.SSHClient: "sock": self.host_proxy, "look_for_keys": self.look_for_keys, "banner_timeout": self.banner_timeout, + "auth_timeout": self.auth_timeout, } if self.password: diff --git a/providers/tests/ssh/hooks/test_ssh.py b/providers/tests/ssh/hooks/test_ssh.py index 6b3e5dcdae116..ac40d01c6f68c 100644 --- a/providers/tests/ssh/hooks/test_ssh.py +++ b/providers/tests/ssh/hooks/test_ssh.py @@ -362,6 +362,7 @@ def test_ssh_connection_with_password(self, ssh_mock): port="port", sock=None, look_for_keys=True, + auth_timeout=None, ) @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient") @@ -381,6 +382,7 @@ def test_ssh_connection_without_password(self, ssh_mock): port="port", sock=None, look_for_keys=True, + auth_timeout=None, ) @mock.patch("airflow.providers.ssh.hooks.ssh.SSHTunnelForwarder") @@ -569,6 +571,7 @@ def test_ssh_connection_with_private_key_extra(self, ssh_mock): port="port", sock=None, look_for_keys=True, + auth_timeout=None, ) @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient") @@ -592,6 +595,7 @@ def test_ssh_connection_with_private_key_passphrase_extra(self, ssh_mock): port="port", sock=None, look_for_keys=True, + auth_timeout=None, ) @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient") @@ -654,6 +658,7 @@ def test_ssh_connection_with_conn_timeout(self, ssh_mock): password="password", conn_timeout=20, key_file="fake.file", + auth_timeout=10, ) with hook.get_conn(): @@ -668,6 +673,7 @@ def test_ssh_connection_with_conn_timeout(self, ssh_mock): port="port", sock=None, look_for_keys=True, + auth_timeout=10, ) @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient") @@ -695,6 +701,7 @@ def test_ssh_connection_with_conn_timeout_and_timeout(self, ssh_mock): port="port", sock=None, look_for_keys=True, + auth_timeout=None, ) @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient") @@ -718,6 +725,7 @@ def test_ssh_connection_with_timeout_extra(self, ssh_mock): port="port", sock=None, look_for_keys=True, + auth_timeout=None, ) @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient") @@ -743,6 +751,7 @@ def test_ssh_connection_with_conn_timeout_extra(self, ssh_mock): port="port", sock=None, look_for_keys=True, + auth_timeout=None, ) @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient") @@ -768,6 +777,7 @@ def test_ssh_connection_with_timeout_extra_and_conn_timeout_extra(self, ssh_mock port="port", sock=None, look_for_keys=True, + auth_timeout=None, ) @pytest.mark.parametrize( @@ -837,6 +847,7 @@ def test_ssh_connection_with_all_timeout_param_and_extra_combinations( port="port", sock=None, look_for_keys=True, + auth_timeout=None, ) @pytest.mark.parametrize( @@ -901,6 +912,7 @@ def test_ssh_with_extra_disabled_algorithms(self, ssh_mock): sock=None, look_for_keys=True, disabled_algorithms=TEST_DISABLED_ALGORITHMS, + auth_timeout=None, ) @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient")