[AIRFLOW-1779] Add keepalive packets to ssh hook

Make use of paramiko's set_keepalive method to
send keepalive packets every
keepalive_interval seconds.  This will prevent
long running queries with no terminal
output from being termanated as idle, for example
by an intermediate NAT.

Set on by default with a 30 second interval.

Closes apache#2749 from RJKeevil/add-sshhook-keepalive

(cherry picked from commit 1bde783)
Signed-off-by: Bolke de Bruin <[email protected]>

airflow/contrib/hooks/ssh_hook.py

@@ -46,6 +46,8 @@ class SSHHook(BaseHook, LoggingMixin):
     :type key_file: str
     :param timeout: timeout for the attempt to connect to the remote_host.
     :type timeout: int
+    :param keepalive_interval: send a keepalive packet to remote host every keepalive_interval seconds
+    :type keepalive_interval: int
     """
 
     def __init__(self,
@@ -54,7 +56,8 @@ def __init__(self,
                  username=None,
                  password=None,
                  key_file=None,
-                 timeout=10
+                 timeout=10,
+                 keepalive_interval=30
                  ):
         super(SSHHook, self).__init__(ssh_conn_id)
         self.ssh_conn_id = ssh_conn_id
@@ -63,6 +66,7 @@ def __init__(self,
         self.password = password
         self.key_file = key_file
         self.timeout = timeout
+        self.keepalive_interval = keepalive_interval
         # Default values, overridable from Connection
         self.compress = True
         self.no_host_key_check = True
@@ -140,6 +144,9 @@ def get_conn(self):
                                    compress=self.compress,
                                    sock=host_proxy)
 
+                if self.keepalive_interval:
+                    client.get_transport().set_keepalive(self.keepalive_interval)
+
                 self.client = client
             except paramiko.AuthenticationException as auth_error:
                 self.log.error(

tests/contrib/hooks/test_ssh_hook.py

@@ -33,7 +33,7 @@ class SSHHookTest(unittest.TestCase):
     def setUp(self):
         configuration.load_test_config()
         from airflow.contrib.hooks.ssh_hook import SSHHook
-        self.hook = SSHHook(ssh_conn_id='ssh_default')
+        self.hook = SSHHook(ssh_conn_id='ssh_default', keepalive_interval=10)
         self.hook.no_host_key_check = True
 
     def test_ssh_connection(self):