Add retryer to SFTP hook connection (#13065)

apache · Feb 3, 2021 · 9af452a · 9af452a
1 parent 0e8c77b
commit 9af452a
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/airflow/providers/sftp/hooks/sftp.py b/airflow/providers/sftp/hooks/sftp.py
@@ -21,6 +21,8 @@
 from typing import Dict, List, Optional, Tuple
 
 import pysftp
+import tenacity
+from paramiko import SSHException
 
 from airflow.providers.ssh.hooks.ssh import SSHHook
 
@@ -109,6 +111,12 @@ def __init__(self, ftp_conn_id: str = 'sftp_default', *args, **kwargs) -> None:
                     )
                     self.key_file = extra_options.get('private_key')
 
+    @tenacity.retry(
+        stop=tenacity.stop_after_delay(10),
+        wait=tenacity.wait_exponential(multiplier=1, max=10),
+        retry=tenacity.retry_if_exception_type(SSHException),
+        reraise=True,
+    )
     def get_conn(self) -> pysftp.Connection:
         """Returns an SFTP connection object"""
         if self.conn is None: