Repository: apache/airflow
Pull request: Add a retry with wait interval for SSH operator #14489


closes: #14489

Retry attempts are added if there are any errors connecting or establishing SSH session.
This will NOT retry failures like wrong host, auth errors, socket timeouts.

Problem statement:
Add a retry with wait interval for SSH operator
<!--

Welcome to Apache Airflow!  For a smooth issue process, try to answer the following questions.
Don't worry if they're not all applicable; just try to include what you can :-)

If you need to include code snippets or logs, please put them in fenced code
blocks.  If they're super-long, please use the details tag like
<details><summary>super-long log</summary> lots of stuff </details>

Please delete these comment blocks before submitting the issue.

-->

**Description**
currently SSH operator fails if authentication fails without retying.  We can add two more parameters as mentioned below which will retry in case of Authentication failure after waiting for configured time. 

- max_retry - maximum time SSH operator should retry in case of exception 
-  wait - how many seconds it should wait before next retry 


<!-- A short description of your feature -->

**Use case / motivation**
We are using SSH operator heavily in our production jobs, And what I have noticed is sometimes SSH operator fails to authenticate, however open re-running jobs it run successfully. And  this happens ofthen. We have ended up writing our own custom operator for this. However, if we can implement this, this could help others as well. 

<!-- What do you want to happen?

Rather than telling us how you might implement this solution, try to take a
step back and describe what you are trying to achieve.

-->

Implement the suggested feature for ssh operator. 


**Are you willing to submit a PR?**
I will submit the PR if feature gets approval. 

<!--- We accept contributions! -->

**Related Issues**
N/A
<!-- Is there currently another issue associated with this? -->
No

Hints from issue discussion:
SGTM. GoogleCompute retry authorizations, so generic 
hook can also do it. 
@mik-laj  can you elaborate on the generic hook thing? I was thinking to modify the existing hook and add an option for this. 
We have 2 hooks that allows us to connect to SSH server. Second hook have retry mechanism. See: https://github.com/apache/airflow/blob/3f7ebfdfe2a1fa90b0854028a5db057adacd46c1/airflow/providers/google/cloud/hooks/compute_ssh.py#L251
@mik-laj  perfect. we can similar implementation for this as well, then 
https://github.com/apache/airflow/blob/3f7ebfdfe2a1fa90b0854028a5db057adacd46c1/airflow/providers/ssh/hooks/ssh.py
@eladkal Can I take this up?
assigned it to you @nshetty15 
@nshetty15 are you still working on it?
@eladkal  I had some issues setting up airflow. Its good now. I will submit a PR soon for this one.

Commit message: Add a retry with wait interval for SSH operator #14489

Patch to review:
diff --git a/airflow/providers/ssh/CHANGELOG.rst b/airflow/providers/ssh/CHANGELOG.rst
index 8a6182cb598c3..cfdf47bba0815 100644
--- a/airflow/providers/ssh/CHANGELOG.rst
+++ b/airflow/providers/ssh/CHANGELOG.rst
@@ -19,6 +19,17 @@
 Changelog
 ---------
 
+2.4.0
+.....
+
+Features
+~~~~~~~~
+
+* ``Add a retry with wait interval for SSH operator (#14489)``
+
+.. Below changes are excluded from the changelog. Move them to
+   appropriate section above if needed. Do not delete the lines(!):
+
 2.3.0
 .....
 
diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index cc4ef1ca48f1f..7620457aa38d1 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -17,6 +17,8 @@
 # under the License.
 """Hook for SSH connections."""
 import os
+import random
+import time
 import warnings
 from base64 import decodebytes
 from io import StringIO
@@ -257,58 +259,74 @@ def get_conn(self) -> paramiko.SSHClient:
         :rtype: paramiko.client.SSHClient
         """
         self.log.debug('Creating SSH client for conn_id: %s', self.ssh_conn_id)
-        client = paramiko.SSHClient()
 
-        if not self.allow_host_key_change:
-            self.log.warning(
-                "Remote Identification Change is not verified. "
-                "This won't protect against Man-In-The-Middle attacks"
-            )
-            client.load_system_host_keys()
+        max_time_to_wait = 10
+        for time_to_wait in self._exponential_sleep_generator(initial=1, maximum=max_time_to_wait):
+            try:
+                client = paramiko.SSHClient()
 
-        if self.no_host_key_check:
-            self.log.warning("No Host Key Verification. This won't protect against Man-In-The-Middle attacks")
-            # Default is RejectPolicy
-            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-        else:
-            if self.host_key is not None:
-                client_host_keys = client.get_host_keys()
-                if self.port == SSH_PORT:
-                    client_host_keys.add(self.remote_host, self.host_key.get_name(), self.host_key)
-                else:
-                    client_host_keys.add(
-                        f"[{self.remote_host}]:{self.port}", self.host_key.get_name(), self.host_key
+                if not self.allow_host_key_change:
+                    self.log.warning(
+                        "Remote Identification Change is not verified. "
+                        "This won't protect against Man-In-The-Middle attacks"
                     )
-            else:
-                pass  # will fallback to system host keys if none explicitly specified in conn extra
-
-        connect_kwargs = dict(
-            hostname=self.remote_host,
-            username=self.username,
-            timeout=self.conn_timeout,
-            compress=self.compress,
-            port=self.port,
-            sock=self.host_proxy,
-            look_for_keys=self.look_for_keys,
-        )
-
-        if self.password:
-            password = self.password.strip()
-            connect_kwargs.update(password=password)
-
-        if self.pkey:
-            connect_kwargs.update(pkey=self.pkey)
-
-        if self.key_file:
-            connect_kwargs.update(key_filename=self.key_file)
-
-        client.connect(**connect_kwargs)
+                    client.load_system_host_keys()
 
-        if self.keepalive_interval:
-            client.get_transport().set_keepalive(self.keepalive_interval)
-
-        self.client = client
-        return client
+                if self.no_host_key_check:
+                    self.log.warning(
+                        "No Host Key Verification. This won't protect against Man-In-The-Middle attacks"
+                    )
+                    # Default is RejectPolicy
+                    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+                else:
+                    if self.host_key is not None:
+                        client_host_keys = client.get_host_keys()
+                        if self.port == SSH_PORT:
+                            client_host_keys.add(self.remote_host, self.host_key.get_name(), self.host_key)
+                        else:
+                            client_host_keys.add(
+                                f"[{self.remote_host}]:{self.port}", self.host_key.get_name(), self.host_key
+                            )
+                    else:
+                        pass  # will fallback to system host keys if none explicitly specified in conn extra
+
+                connect_kwargs = dict(
+                    hostname=self.remote_host,
+                    username=self.username,
+                    timeout=self.conn_timeout,
+                    compress=self.compress,
+                    port=self.port,
+                    sock=self.host_proxy,
+                    look_for_keys=self.look_for_keys,
+                )
+                self.log.warning(connect_kwargs)
+
+                if self.password:
+                    password = self.password.strip()
+                    connect_kwargs.update(password=password)
+
+                if self.pkey:
+                    connect_kwargs.update(pkey=self.pkey)
+
+                if self.key_file:
+                    connect_kwargs.update(key_filename=self.key_file)
+
+                client.connect(**connect_kwargs)
+
+                if self.keepalive_interval:
+                    client.get_transport().set_keepalive(self.keepalive_interval)
+
+                self.client = client
+                return client
+
+            except paramiko.SSHException:
+                # exponential_sleep_generator is an infinite generator, so we need to
+                # check the end condition.
+                if time_to_wait == max_time_to_wait:
+                    raise
+            self.log.info("Failed to connect. Waiting %ds to retry", time_to_wait)
+            time.sleep(time_to_wait)
+        raise AirflowException("SSH connection failed")
 
     def __enter__(self) -> 'SSHHook':
         warnings.warn(
@@ -411,3 +429,23 @@ def _pkey_from_private_key(self, private_key: str, passphrase: Optional[str] = N
             'Ensure key provided is valid for one of the following'
             'key formats: RSA, DSS, ECDSA, or Ed25519'
         )
+
+    def _exponential_sleep_generator(self, initial, maximum, multiplier=2.0):
+        """
+        Generates sleep intervals based on the exponential back-off algorithm.
+
+        :param initial: The minimum amount of time to delay. This must be greater than 0.
+        :type initial: float
+        :param maximum: The maximum amount of time to delay.
+        :type maximum: float
+        :param multiplier: The multiplier applied to the delay.
+        :type multiplier: float
+
+        :return: successive sleep intervals.
+        """
+        delay = initial
+        while True:
+            # Introduce jitter by yielding a delay that is uniformly distributed
+            # to average out to the delay time.
+            yield min(random.uniform(0.0, delay * 2.0), maximum)
+            delay = delay * multiplier
diff --git a/airflow/providers/ssh/provider.yaml b/airflow/providers/ssh/provider.yaml
index 7e9b1e66506a3..601bcb304a328 100644
--- a/airflow/providers/ssh/provider.yaml
+++ b/airflow/providers/ssh/provider.yaml
@@ -22,6 +22,7 @@ description: |
     `Secure Shell (SSH) <https://tools.ietf.org/html/rfc4251>`__
 
 versions:
+  - 2.4.0
   - 2.3.0
   - 2.2.0
   - 2.1.1


Code review comments:
diff_hunk: @@ -257,58 +259,74 @@ def get_conn(self) -> paramiko.SSHClient:
         :rtype: paramiko.client.SSHClient
         """
         self.log.debug('Creating SSH client for conn_id: %s', self.ssh_conn_id)
-        client = paramiko.SSHClient()
 
-        if not self.allow_host_key_change:
-            self.log.warning(
-                "Remote Identification Change is not verified. "
-                "This won't protect against Man-In-The-Middle attacks"
-            )
-            client.load_system_host_keys()
+        max_time_to_wait = 10
+        for time_to_wait in self._exponential_sleep_generator(initial=1, maximum=max_time_to_wait):
path: airflow/providers/ssh/hooks/ssh.py
text: @user1:
We should use tenacity here (we already use it elsewhere).

@author:
okay.. let me change this to use tenacity.

diff_hunk: @@ -257,58 +259,74 @@ def get_conn(self) -> paramiko.SSHClient:
         :rtype: paramiko.client.SSHClient
         """
         self.log.debug('Creating SSH client for conn_id: %s', self.ssh_conn_id)
-        client = paramiko.SSHClient()
 
-        if not self.allow_host_key_change:
-            self.log.warning(
-                "Remote Identification Change is not verified. "
-                "This won't protect against Man-In-The-Middle attacks"
-            )
-            client.load_system_host_keys()
+        max_time_to_wait = 10
+        for time_to_wait in self._exponential_sleep_generator(initial=1, maximum=max_time_to_wait):
+            try:
+                client = paramiko.SSHClient()
 
-        if self.no_host_key_check:
-            self.log.warning("No Host Key Verification. This won't protect against Man-In-The-Middle attacks")
-            # Default is RejectPolicy
-            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-        else:
-            if self.host_key is not None:
-                client_host_keys = client.get_host_keys()
-                if self.port == SSH_PORT:
-                    client_host_keys.add(self.remote_host, self.host_key.get_name(), self.host_key)
-                else:
-                    client_host_keys.add(
-                        f"[{self.remote_host}]:{self.port}", self.host_key.get_name(), self.host_key
+                if not self.allow_host_key_change:
path: airflow/providers/ssh/hooks/ssh.py
text: @user1:
I think we shoudl move it up to before the loop. Otherwise it will keep on repeating in the logs when ssh fails.