From a0a97208d8aa1b7fcb3823d7164cd7aa49055d16 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Sun, 25 Jul 2021 19:43:36 -0400
Subject: [PATCH 01/42] Issue 16363 - add conn_timeout and cmd_timeout params
 to SSHOperator; add conn_timoue param to SSHHook

---
 airflow/providers/ssh/hooks/ssh.py        |  40 +++++++-
 airflow/providers/ssh/operators/ssh.py    |  32 +++++-
 tests/providers/ssh/hooks/test_ssh.py     | 117 +++++++++++++++++++++-
 tests/providers/ssh/operators/test_ssh.py |  35 +++++--
 4 files changed, 207 insertions(+), 17 deletions(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index a8994bdece1b7..b4f524515e5e3 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -56,8 +56,10 @@ class SSHHook(BaseHook):
     :type key_file: str
     :param port: port of remote host to connect (Default is paramiko SSH_PORT)
     :type port: int
-    :param timeout: timeout for the attempt to connect to the remote_host.
+    :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host.
     :type timeout: int
+    :param conn_timeout: timeout for the attempt to connect to the remote_host.
+    :type conn_timeout: int
     :param keepalive_interval: send a keepalive packet to remote host every
         keepalive_interval seconds
     :type keepalive_interval: int
@@ -101,7 +103,8 @@ def __init__(
         password: Optional[str] = None,
         key_file: Optional[str] = None,
         port: Optional[int] = None,
-        timeout: int = 10,
+        timeout: Optional[int] = 10,
+        conn_timeout: Optional[int] = None,
         keepalive_interval: int = 30,
     ) -> None:
         super().__init__()
@@ -113,6 +116,7 @@ def __init__(
         self.pkey = None
         self.port = port
         self.timeout = timeout
+        self.conn_timeout = conn_timeout
         self.keepalive_interval = keepalive_interval
 
         # Default values, overridable from Connection
@@ -137,6 +141,7 @@ def __init__(
                 self.remote_host = conn.host
             if self.port is None:
                 self.port = conn.port
+
             if conn.extra is not None:
                 extra_options = conn.extra_dejson
                 if "key_file" in extra_options and self.key_file is None:
@@ -148,7 +153,22 @@ def __init__(
                     self.pkey = self._pkey_from_private_key(private_key, passphrase=private_key_passphrase)
 
                 if "timeout" in extra_options:
-                    self.timeout = int(extra_options["timeout"], 10)
+                    warnings.warn(
+                        'Extra option `timeout` is deprecated.'
+                        'Please use `conn_timeout` instead.'
+                        'The old option `timeout` will be removed in a future version.',
+                        DeprecationWarning,
+                        stacklevel=2,
+                    )
+                    self.timeout = int(extra_options.get('timeout', 10))
+
+                if "conn_timeout" in extra_options:
+                    self.conn_timeout = int(extra_options.get('conn_timeout'))
+
+                # TODO Rework this - I don't like that just having conn.extra set will override 
+                # self.conn_timeout = int(extra_options.get('conn_timeout',
+                    # extra_options.get('timeout', 10))
+                #)
 
                 if "compress" in extra_options and str(extra_options["compress"]).lower() == 'false':
                     self.compress = False
@@ -185,6 +205,18 @@ def __init__(
                     self.host_key = key_constructor(data=decoded_host_key)
                     self.no_host_key_check = False
 
+        if self.timeout:
+            warnings.warn(
+                'Parameter `timeout` is deprecated.'
+                'Please use `conn_timeout` instead.'
+                'The old option `timeout` will be removed in a future version.',
+                DeprecationWarning,
+                stacklevel=1,
+            )
+
+        if self.conn_timeout is None:
+            self.conn_timeout = self.timeout
+
         if self.pkey and self.key_file:
             raise AirflowException(
                 "Params key_file and private_key both provided.  Must provide no more than one."
@@ -253,7 +285,7 @@ def get_conn(self) -> paramiko.SSHClient:
         connect_kwargs = dict(
             hostname=self.remote_host,
             username=self.username,
-            timeout=self.timeout,
+            timeout=self.conn_timeout,
             compress=self.compress,
             port=self.port,
             sock=self.host_proxy,
diff --git a/airflow/providers/ssh/operators/ssh.py b/airflow/providers/ssh/operators/ssh.py
index c82e0e12c30a4..27380977a2abb 100644
--- a/airflow/providers/ssh/operators/ssh.py
+++ b/airflow/providers/ssh/operators/ssh.py
@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import warnings
 from base64 import b64encode
 from select import select
 from typing import Optional, Union
@@ -43,8 +44,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 timeout: (deprecated) timeout (in seconds) for executing the command. The default is 10 seconds.
     :type timeout: int
+    :param conn_timeout: timeout (in seconds) for maintaining the connection. The default is 10 seconds.
+    :type conn_timeout: int
+    :param cmd_timeout: timeout (in seconds) for executing the command. The default is 10 seconds.
+    :type cmd_timeout: int
     :param environment: a dict of shell environment variables. Note that the
         server will reject them silently if `AcceptEnv` is not set in SSH config.
     :type environment: dict
@@ -66,7 +71,9 @@ def __init__(
         ssh_conn_id: Optional[str] = None,
         remote_host: Optional[str] = None,
         command: Optional[str] = None,
-        timeout: int = 10,
+        timeout: Optional[int] = 10,
+        conn_timeout: Optional[int] = None,
+        cmd_timeout: Optional[int] = None,
         environment: Optional[dict] = None,
         get_pty: bool = False,
         **kwargs,
@@ -77,9 +84,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.cmd_timeout is None:
+            self.cmd_timeout = self.timeout
         self.environment = environment
         self.get_pty = (self.command.startswith('sudo') or get_pty) if self.command else get_pty
 
+        if self.timeout:
+            warnings.warn(
+                'Parameter `timeout` is deprecated.'
+                'Please use `conn_timeout` and `cmd_timeout` instead.'
+                'The old option `timeout` will be removed in a future version.',
+                DeprecationWarning,
+                stacklevel=1,
+            )
+
     def execute(self, context) -> Union[bytes, str, bool]:
         try:
             if self.ssh_conn_id:
@@ -89,7 +111,7 @@ def execute(self, context) -> Union[bytes, str, bool]:
                     self.log.info(
                         "ssh_hook is not provided or invalid. Trying ssh_conn_id to create SSHHook."
                     )
-                    self.ssh_hook = SSHHook(ssh_conn_id=self.ssh_conn_id, timeout=self.timeout)
+                    self.ssh_hook = SSHHook(ssh_conn_id=self.ssh_conn_id, conn_timeout=self.conn_timeout)
 
             if not self.ssh_hook:
                 raise AirflowException("Cannot operate without ssh_hook or ssh_conn_id.")
@@ -112,7 +134,7 @@ def execute(self, context) -> Union[bytes, str, bool]:
                 stdin, stdout, stderr = ssh_client.exec_command(
                     command=self.command,
                     get_pty=self.get_pty,
-                    timeout=self.timeout,
+                    timeout=self.cmd_timeout,
                     environment=self.environment,
                 )
                 # get channels
@@ -133,7 +155,7 @@ def execute(self, context) -> Union[bytes, str, bool]:
 
                 # read from both stdout and stderr
                 while not channel.closed or channel.recv_ready() or channel.recv_stderr_ready():
-                    readq, _, _ = select([channel], [], [], self.timeout)
+                    readq, _, _ = select([channel], [], [], self.cmd_timeout)
                     for recv in readq:
                         if recv.recv_ready():
                             line = stdout.channel.recv(len(recv.in_buffer))
diff --git a/tests/providers/ssh/hooks/test_ssh.py b/tests/providers/ssh/hooks/test_ssh.py
index 70ba4eb0ce642..18389be798ba7 100644
--- a/tests/providers/ssh/hooks/test_ssh.py
+++ b/tests/providers/ssh/hooks/test_ssh.py
@@ -68,6 +68,9 @@ def generate_host_key(pkey: paramiko.PKey):
 TEST_PKEY_ECDSA = paramiko.ECDSAKey.generate()
 TEST_PRIVATE_KEY_ECDSA = generate_key_string(pkey=TEST_PKEY_ECDSA)
 
+TEST_TIMEOUT = 20
+TEST_CONN_TIMEOUT = 30
+
 PASSPHRASE = ''.join(random.choice(string.ascii_letters) for i in range(10))
 TEST_ENCRYPTED_PRIVATE_KEY = generate_key_string(pkey=TEST_PKEY, passphrase=PASSPHRASE)
 
@@ -76,6 +79,8 @@ class TestSSHHook(unittest.TestCase):
     CONN_SSH_WITH_PRIVATE_KEY_EXTRA = 'ssh_with_private_key_extra'
     CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA = 'ssh_with_private_key_ecdsa_extra'
     CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA = 'ssh_with_private_key_passphrase_extra'
+    CONN_SSH_WITH_TIMEOUT_EXTRA = 'ssh_with_timeout_extra'
+    CONN_SSH_WITH_CONN_TIMEOUT_EXTRA = 'ssh_with_conn_timeout_extra'
     CONN_SSH_WITH_EXTRA = 'ssh_with_extra'
     CONN_SSH_WITH_EXTRA_FALSE_LOOK_FOR_KEYS = 'ssh_with_extra_false_look_for_keys'
     CONN_SSH_WITH_HOST_KEY_EXTRA = 'ssh_with_host_key_extra'
@@ -91,6 +96,8 @@ def tearDownClass(cls) -> None:
                 cls.CONN_SSH_WITH_PRIVATE_KEY_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA,
+                cls.CONN_SSH_WITH_TIMEOUT_EXTRA,
+                cls.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA,
                 cls.CONN_SSH_WITH_EXTRA,
                 cls.CONN_SSH_WITH_HOST_KEY_EXTRA,
                 cls.CONN_SSH_WITH_HOST_KEY_EXTRA_WITH_TYPE,
@@ -128,7 +135,7 @@ def setUpClass(cls) -> None:
                 conn_type='ssh',
                 extra=json.dumps(
                     {
-                        "private_key": TEST_PRIVATE_KEY,
+                        "private_key": TEST_PRIVATE_KEY
                     }
                 ),
             )
@@ -151,6 +158,22 @@ def setUpClass(cls) -> None:
                 extra=json.dumps({"private_key": TEST_PRIVATE_KEY_ECDSA}),
             )
         )
+        db.merge_conn(
+            Connection(
+                conn_id=cls.CONN_SSH_WITH_TIMEOUT_EXTRA,
+                host='localhost',
+                conn_type='ssh',
+                extra=json.dumps({"timeout": TEST_TIMEOUT}),
+            )
+        )
+        db.merge_conn(
+            Connection(
+                conn_id=cls.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA,
+                host='localhost',
+                conn_type='ssh',
+                extra=json.dumps({"conn_timeout": TEST_CONN_TIMEOUT}),
+            )
+        )
         db.merge_conn(
             Connection(
                 conn_id=cls.CONN_SSH_WITH_HOST_KEY_EXTRA,
@@ -475,6 +498,98 @@ 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
+        )
+
+        with hook.get_conn():
+            ssh_mock.return_value.connect.assert_called_once_with(
+                hostname='remote_host',
+                username='username',
+                timeout=30,
+                compress=True,
+                port='port',
+                sock=None,
+                look_for_keys=True,
+            )
+
     def test_openssh_private_key(self):
         # Paramiko behaves differently with OpenSSH generated keys to paramiko
         # generated keys, so we need a test one.
diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index e39b8fbb8c156..f918278768133 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -52,7 +52,7 @@ def setUp(self):
         self.hook = hook
         self.dag = dag
 
-    def test_hook_created_correctly(self):
+    def test_hook_created_correctly_with_timeout(self):
         timeout = 20
         ssh_id = "ssh_default"
         task = SSHOperator(
@@ -61,8 +61,21 @@ def test_hook_created_correctly(self):
         assert task is not None
 
         task.execute(None)
+        
+        assert timeout == task.ssh_hook.conn_timeout
+        assert ssh_id == task.ssh_hook.ssh_conn_id
 
-        assert timeout == task.ssh_hook.timeout
+    def test_hook_created_correctly(self):
+        conn_timeout = 20
+        ssh_id = "ssh_default"
+        task = SSHOperator(
+            task_id="test", command=COMMAND, dag=self.dag, conn_timeout=conn_timeout, ssh_conn_id="ssh_default"
+        )
+        assert task is not None
+
+        task.execute(None)
+        
+        assert conn_timeout == task.ssh_hook.conn_timeout
         assert ssh_id == task.ssh_hook.ssh_conn_id
 
     @conf_vars({('core', 'enable_xcom_pickling'): 'False'})
@@ -138,7 +151,8 @@ def test_no_output_command(self):
     def test_arg_checking(self):
         # Exception should be raised if neither ssh_hook nor ssh_conn_id is provided
         with pytest.raises(AirflowException, match="Cannot operate without ssh_hook or ssh_conn_id."):
-            task_0 = SSHOperator(task_id="test", command=COMMAND, timeout=TIMEOUT, dag=self.dag)
+            #  task_0 = SSHOperator(task_id="test", command=COMMAND, timeout=TIMEOUT, dag=self.dag)
+            task_0 = SSHOperator(task_id="test", command=COMMAND, conn_timeout=TIMEOUT, cmd_timeout=TIMEOUT, dag=self.dag)
             task_0.execute(None)
 
         # if ssh_hook is invalid/not provided, use ssh_conn_id to create SSHHook
@@ -147,7 +161,8 @@ def test_arg_checking(self):
             ssh_hook="string_rather_than_SSHHook",  # invalid ssh_hook
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -160,7 +175,9 @@ def test_arg_checking(self):
             task_id="test_2",
             ssh_conn_id=TEST_CONN_ID,  # no ssh_hook provided
             command=COMMAND,
-            timeout=TIMEOUT,
+            # timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -175,7 +192,9 @@ def test_arg_checking(self):
             ssh_hook=self.hook,
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            timeout=TIMEOUT,
+            # timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -198,7 +217,9 @@ def test_get_pyt_set_correctly(self, command, get_pty_in, get_pty_out):
             task_id="test",
             ssh_hook=self.hook,
             command=command,
-            timeout=TIMEOUT,
+            # timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             get_pty=get_pty_in,
             dag=self.dag,
         )

From d3f0052b0017f1b1c9e23403486bf41c8c249e7d Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Sun, 25 Jul 2021 20:24:56 -0400
Subject: [PATCH 02/42] Update SSHOperator tests for new conn_timeout and
 cmd_timeout params

---
 tests/providers/ssh/operators/test_ssh.py | 55 ++++++++++++++++++-----
 1 file changed, 44 insertions(+), 11 deletions(-)

diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index f918278768133..0d8fbcae493fa 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -31,7 +31,9 @@
 
 TEST_DAG_ID = 'unit_tests_ssh_test_op'
 TEST_CONN_ID = "conn_id_for_testing"
-TIMEOUT = 5
+CONN_TIMEOUT = 5
+CMD_TIMEOUT = 7
+TIMEOUT = 12
 DEFAULT_DATE = datetime(2017, 1, 1)
 COMMAND = "echo -n airflow"
 COMMAND_WITH_SUDO = "sudo " + COMMAND
@@ -67,9 +69,10 @@ def test_hook_created_correctly_with_timeout(self):
 
     def test_hook_created_correctly(self):
         conn_timeout = 20
+        cmd_timeout = 45
         ssh_id = "ssh_default"
         task = SSHOperator(
-            task_id="test", command=COMMAND, dag=self.dag, conn_timeout=conn_timeout, ssh_conn_id="ssh_default"
+            task_id="test", command=COMMAND, dag=self.dag, conn_timeout=conn_timeout, cmd_timeout=cmd_timeout, ssh_conn_id="ssh_default"
         )
         assert task is not None
 
@@ -151,7 +154,6 @@ def test_no_output_command(self):
     def test_arg_checking(self):
         # Exception should be raised if neither ssh_hook nor ssh_conn_id is provided
         with pytest.raises(AirflowException, match="Cannot operate without ssh_hook or ssh_conn_id."):
-            #  task_0 = SSHOperator(task_id="test", command=COMMAND, timeout=TIMEOUT, dag=self.dag)
             task_0 = SSHOperator(task_id="test", command=COMMAND, conn_timeout=TIMEOUT, cmd_timeout=TIMEOUT, dag=self.dag)
             task_0.execute(None)
 
@@ -161,8 +163,8 @@ def test_arg_checking(self):
             ssh_hook="string_rather_than_SSHHook",  # invalid ssh_hook
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            conn_timeout=TIMEOUT,
-            cmd_timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -175,9 +177,8 @@ def test_arg_checking(self):
             task_id="test_2",
             ssh_conn_id=TEST_CONN_ID,  # no ssh_hook provided
             command=COMMAND,
-            # timeout=TIMEOUT,
-            conn_timeout=TIMEOUT,
-            cmd_timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -192,9 +193,8 @@ def test_arg_checking(self):
             ssh_hook=self.hook,
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            # timeout=TIMEOUT,
-            conn_timeout=TIMEOUT,
-            cmd_timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -203,6 +203,39 @@ def test_arg_checking(self):
             pass
         assert task_3.ssh_hook.ssh_conn_id == self.hook.ssh_conn_id
 
+        # if valid conn_timeout, cmd_timeout, and timeout are provided, ignore timeout
+        task_4 = SSHOperator(
+            task_id="test_4",
+            ssh_conn_id=TEST_CONN_ID,
+            command=COMMAND,
+            timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
+            dag=self.dag,
+        )
+        try:
+            task_4.execute(None)
+        except Exception:
+            pass
+        assert task_4.conn_timeout == CONN_TIMEOUT
+        assert task_4.cmd_timeout == CMD_TIMEOUT
+        assert task_4.ssh_hook.conn_timeout == CONN_TIMEOUT
+
+        # if valid only timeout is provided, use for both conn_timeout and cmd_timeout
+        task_5 = SSHOperator(
+            task_id="test_5",
+            ssh_conn_id=TEST_CONN_ID,
+            command=COMMAND,
+            timeout=TIMEOUT,
+            dag=self.dag,
+        )
+        try:
+            task_5.execute(None)
+        except Exception:
+            pass
+        assert task_5.cmd_timeout == TIMEOUT
+        assert task_5.ssh_hook.conn_timeout == TIMEOUT
+
     @parameterized.expand(
         [
             (COMMAND, False, False),

From 9784e56a9da47d8e7e6fd0acbb617f1b01580a99 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Sun, 25 Jul 2021 19:43:36 -0400
Subject: [PATCH 03/42] Issue 16363 - add conn_timeout and cmd_timeout params
 to SSHOperator; add conn_timoue param to SSHHook

---
 airflow/providers/ssh/hooks/ssh.py        |  40 +++++++-
 airflow/providers/ssh/operators/ssh.py    |  32 +++++-
 tests/providers/ssh/hooks/test_ssh.py     | 117 +++++++++++++++++++++-
 tests/providers/ssh/operators/test_ssh.py |  35 +++++--
 4 files changed, 207 insertions(+), 17 deletions(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index a8994bdece1b7..b4f524515e5e3 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -56,8 +56,10 @@ class SSHHook(BaseHook):
     :type key_file: str
     :param port: port of remote host to connect (Default is paramiko SSH_PORT)
     :type port: int
-    :param timeout: timeout for the attempt to connect to the remote_host.
+    :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host.
     :type timeout: int
+    :param conn_timeout: timeout for the attempt to connect to the remote_host.
+    :type conn_timeout: int
     :param keepalive_interval: send a keepalive packet to remote host every
         keepalive_interval seconds
     :type keepalive_interval: int
@@ -101,7 +103,8 @@ def __init__(
         password: Optional[str] = None,
         key_file: Optional[str] = None,
         port: Optional[int] = None,
-        timeout: int = 10,
+        timeout: Optional[int] = 10,
+        conn_timeout: Optional[int] = None,
         keepalive_interval: int = 30,
     ) -> None:
         super().__init__()
@@ -113,6 +116,7 @@ def __init__(
         self.pkey = None
         self.port = port
         self.timeout = timeout
+        self.conn_timeout = conn_timeout
         self.keepalive_interval = keepalive_interval
 
         # Default values, overridable from Connection
@@ -137,6 +141,7 @@ def __init__(
                 self.remote_host = conn.host
             if self.port is None:
                 self.port = conn.port
+
             if conn.extra is not None:
                 extra_options = conn.extra_dejson
                 if "key_file" in extra_options and self.key_file is None:
@@ -148,7 +153,22 @@ def __init__(
                     self.pkey = self._pkey_from_private_key(private_key, passphrase=private_key_passphrase)
 
                 if "timeout" in extra_options:
-                    self.timeout = int(extra_options["timeout"], 10)
+                    warnings.warn(
+                        'Extra option `timeout` is deprecated.'
+                        'Please use `conn_timeout` instead.'
+                        'The old option `timeout` will be removed in a future version.',
+                        DeprecationWarning,
+                        stacklevel=2,
+                    )
+                    self.timeout = int(extra_options.get('timeout', 10))
+
+                if "conn_timeout" in extra_options:
+                    self.conn_timeout = int(extra_options.get('conn_timeout'))
+
+                # TODO Rework this - I don't like that just having conn.extra set will override 
+                # self.conn_timeout = int(extra_options.get('conn_timeout',
+                    # extra_options.get('timeout', 10))
+                #)
 
                 if "compress" in extra_options and str(extra_options["compress"]).lower() == 'false':
                     self.compress = False
@@ -185,6 +205,18 @@ def __init__(
                     self.host_key = key_constructor(data=decoded_host_key)
                     self.no_host_key_check = False
 
+        if self.timeout:
+            warnings.warn(
+                'Parameter `timeout` is deprecated.'
+                'Please use `conn_timeout` instead.'
+                'The old option `timeout` will be removed in a future version.',
+                DeprecationWarning,
+                stacklevel=1,
+            )
+
+        if self.conn_timeout is None:
+            self.conn_timeout = self.timeout
+
         if self.pkey and self.key_file:
             raise AirflowException(
                 "Params key_file and private_key both provided.  Must provide no more than one."
@@ -253,7 +285,7 @@ def get_conn(self) -> paramiko.SSHClient:
         connect_kwargs = dict(
             hostname=self.remote_host,
             username=self.username,
-            timeout=self.timeout,
+            timeout=self.conn_timeout,
             compress=self.compress,
             port=self.port,
             sock=self.host_proxy,
diff --git a/airflow/providers/ssh/operators/ssh.py b/airflow/providers/ssh/operators/ssh.py
index c82e0e12c30a4..27380977a2abb 100644
--- a/airflow/providers/ssh/operators/ssh.py
+++ b/airflow/providers/ssh/operators/ssh.py
@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import warnings
 from base64 import b64encode
 from select import select
 from typing import Optional, Union
@@ -43,8 +44,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 timeout: (deprecated) timeout (in seconds) for executing the command. The default is 10 seconds.
     :type timeout: int
+    :param conn_timeout: timeout (in seconds) for maintaining the connection. The default is 10 seconds.
+    :type conn_timeout: int
+    :param cmd_timeout: timeout (in seconds) for executing the command. The default is 10 seconds.
+    :type cmd_timeout: int
     :param environment: a dict of shell environment variables. Note that the
         server will reject them silently if `AcceptEnv` is not set in SSH config.
     :type environment: dict
@@ -66,7 +71,9 @@ def __init__(
         ssh_conn_id: Optional[str] = None,
         remote_host: Optional[str] = None,
         command: Optional[str] = None,
-        timeout: int = 10,
+        timeout: Optional[int] = 10,
+        conn_timeout: Optional[int] = None,
+        cmd_timeout: Optional[int] = None,
         environment: Optional[dict] = None,
         get_pty: bool = False,
         **kwargs,
@@ -77,9 +84,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.cmd_timeout is None:
+            self.cmd_timeout = self.timeout
         self.environment = environment
         self.get_pty = (self.command.startswith('sudo') or get_pty) if self.command else get_pty
 
+        if self.timeout:
+            warnings.warn(
+                'Parameter `timeout` is deprecated.'
+                'Please use `conn_timeout` and `cmd_timeout` instead.'
+                'The old option `timeout` will be removed in a future version.',
+                DeprecationWarning,
+                stacklevel=1,
+            )
+
     def execute(self, context) -> Union[bytes, str, bool]:
         try:
             if self.ssh_conn_id:
@@ -89,7 +111,7 @@ def execute(self, context) -> Union[bytes, str, bool]:
                     self.log.info(
                         "ssh_hook is not provided or invalid. Trying ssh_conn_id to create SSHHook."
                     )
-                    self.ssh_hook = SSHHook(ssh_conn_id=self.ssh_conn_id, timeout=self.timeout)
+                    self.ssh_hook = SSHHook(ssh_conn_id=self.ssh_conn_id, conn_timeout=self.conn_timeout)
 
             if not self.ssh_hook:
                 raise AirflowException("Cannot operate without ssh_hook or ssh_conn_id.")
@@ -112,7 +134,7 @@ def execute(self, context) -> Union[bytes, str, bool]:
                 stdin, stdout, stderr = ssh_client.exec_command(
                     command=self.command,
                     get_pty=self.get_pty,
-                    timeout=self.timeout,
+                    timeout=self.cmd_timeout,
                     environment=self.environment,
                 )
                 # get channels
@@ -133,7 +155,7 @@ def execute(self, context) -> Union[bytes, str, bool]:
 
                 # read from both stdout and stderr
                 while not channel.closed or channel.recv_ready() or channel.recv_stderr_ready():
-                    readq, _, _ = select([channel], [], [], self.timeout)
+                    readq, _, _ = select([channel], [], [], self.cmd_timeout)
                     for recv in readq:
                         if recv.recv_ready():
                             line = stdout.channel.recv(len(recv.in_buffer))
diff --git a/tests/providers/ssh/hooks/test_ssh.py b/tests/providers/ssh/hooks/test_ssh.py
index 70ba4eb0ce642..18389be798ba7 100644
--- a/tests/providers/ssh/hooks/test_ssh.py
+++ b/tests/providers/ssh/hooks/test_ssh.py
@@ -68,6 +68,9 @@ def generate_host_key(pkey: paramiko.PKey):
 TEST_PKEY_ECDSA = paramiko.ECDSAKey.generate()
 TEST_PRIVATE_KEY_ECDSA = generate_key_string(pkey=TEST_PKEY_ECDSA)
 
+TEST_TIMEOUT = 20
+TEST_CONN_TIMEOUT = 30
+
 PASSPHRASE = ''.join(random.choice(string.ascii_letters) for i in range(10))
 TEST_ENCRYPTED_PRIVATE_KEY = generate_key_string(pkey=TEST_PKEY, passphrase=PASSPHRASE)
 
@@ -76,6 +79,8 @@ class TestSSHHook(unittest.TestCase):
     CONN_SSH_WITH_PRIVATE_KEY_EXTRA = 'ssh_with_private_key_extra'
     CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA = 'ssh_with_private_key_ecdsa_extra'
     CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA = 'ssh_with_private_key_passphrase_extra'
+    CONN_SSH_WITH_TIMEOUT_EXTRA = 'ssh_with_timeout_extra'
+    CONN_SSH_WITH_CONN_TIMEOUT_EXTRA = 'ssh_with_conn_timeout_extra'
     CONN_SSH_WITH_EXTRA = 'ssh_with_extra'
     CONN_SSH_WITH_EXTRA_FALSE_LOOK_FOR_KEYS = 'ssh_with_extra_false_look_for_keys'
     CONN_SSH_WITH_HOST_KEY_EXTRA = 'ssh_with_host_key_extra'
@@ -91,6 +96,8 @@ def tearDownClass(cls) -> None:
                 cls.CONN_SSH_WITH_PRIVATE_KEY_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA,
+                cls.CONN_SSH_WITH_TIMEOUT_EXTRA,
+                cls.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA,
                 cls.CONN_SSH_WITH_EXTRA,
                 cls.CONN_SSH_WITH_HOST_KEY_EXTRA,
                 cls.CONN_SSH_WITH_HOST_KEY_EXTRA_WITH_TYPE,
@@ -128,7 +135,7 @@ def setUpClass(cls) -> None:
                 conn_type='ssh',
                 extra=json.dumps(
                     {
-                        "private_key": TEST_PRIVATE_KEY,
+                        "private_key": TEST_PRIVATE_KEY
                     }
                 ),
             )
@@ -151,6 +158,22 @@ def setUpClass(cls) -> None:
                 extra=json.dumps({"private_key": TEST_PRIVATE_KEY_ECDSA}),
             )
         )
+        db.merge_conn(
+            Connection(
+                conn_id=cls.CONN_SSH_WITH_TIMEOUT_EXTRA,
+                host='localhost',
+                conn_type='ssh',
+                extra=json.dumps({"timeout": TEST_TIMEOUT}),
+            )
+        )
+        db.merge_conn(
+            Connection(
+                conn_id=cls.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA,
+                host='localhost',
+                conn_type='ssh',
+                extra=json.dumps({"conn_timeout": TEST_CONN_TIMEOUT}),
+            )
+        )
         db.merge_conn(
             Connection(
                 conn_id=cls.CONN_SSH_WITH_HOST_KEY_EXTRA,
@@ -475,6 +498,98 @@ 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
+        )
+
+        with hook.get_conn():
+            ssh_mock.return_value.connect.assert_called_once_with(
+                hostname='remote_host',
+                username='username',
+                timeout=30,
+                compress=True,
+                port='port',
+                sock=None,
+                look_for_keys=True,
+            )
+
     def test_openssh_private_key(self):
         # Paramiko behaves differently with OpenSSH generated keys to paramiko
         # generated keys, so we need a test one.
diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index e39b8fbb8c156..f918278768133 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -52,7 +52,7 @@ def setUp(self):
         self.hook = hook
         self.dag = dag
 
-    def test_hook_created_correctly(self):
+    def test_hook_created_correctly_with_timeout(self):
         timeout = 20
         ssh_id = "ssh_default"
         task = SSHOperator(
@@ -61,8 +61,21 @@ def test_hook_created_correctly(self):
         assert task is not None
 
         task.execute(None)
+        
+        assert timeout == task.ssh_hook.conn_timeout
+        assert ssh_id == task.ssh_hook.ssh_conn_id
 
-        assert timeout == task.ssh_hook.timeout
+    def test_hook_created_correctly(self):
+        conn_timeout = 20
+        ssh_id = "ssh_default"
+        task = SSHOperator(
+            task_id="test", command=COMMAND, dag=self.dag, conn_timeout=conn_timeout, ssh_conn_id="ssh_default"
+        )
+        assert task is not None
+
+        task.execute(None)
+        
+        assert conn_timeout == task.ssh_hook.conn_timeout
         assert ssh_id == task.ssh_hook.ssh_conn_id
 
     @conf_vars({('core', 'enable_xcom_pickling'): 'False'})
@@ -138,7 +151,8 @@ def test_no_output_command(self):
     def test_arg_checking(self):
         # Exception should be raised if neither ssh_hook nor ssh_conn_id is provided
         with pytest.raises(AirflowException, match="Cannot operate without ssh_hook or ssh_conn_id."):
-            task_0 = SSHOperator(task_id="test", command=COMMAND, timeout=TIMEOUT, dag=self.dag)
+            #  task_0 = SSHOperator(task_id="test", command=COMMAND, timeout=TIMEOUT, dag=self.dag)
+            task_0 = SSHOperator(task_id="test", command=COMMAND, conn_timeout=TIMEOUT, cmd_timeout=TIMEOUT, dag=self.dag)
             task_0.execute(None)
 
         # if ssh_hook is invalid/not provided, use ssh_conn_id to create SSHHook
@@ -147,7 +161,8 @@ def test_arg_checking(self):
             ssh_hook="string_rather_than_SSHHook",  # invalid ssh_hook
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -160,7 +175,9 @@ def test_arg_checking(self):
             task_id="test_2",
             ssh_conn_id=TEST_CONN_ID,  # no ssh_hook provided
             command=COMMAND,
-            timeout=TIMEOUT,
+            # timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -175,7 +192,9 @@ def test_arg_checking(self):
             ssh_hook=self.hook,
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            timeout=TIMEOUT,
+            # timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -198,7 +217,9 @@ def test_get_pyt_set_correctly(self, command, get_pty_in, get_pty_out):
             task_id="test",
             ssh_hook=self.hook,
             command=command,
-            timeout=TIMEOUT,
+            # timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             get_pty=get_pty_in,
             dag=self.dag,
         )

From dc2d15e83420d0ba6643dc855d3ea0fc5848b982 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Sun, 25 Jul 2021 20:24:56 -0400
Subject: [PATCH 04/42] Update SSHOperator tests for new conn_timeout and
 cmd_timeout params

---
 tests/providers/ssh/operators/test_ssh.py | 55 ++++++++++++++++++-----
 1 file changed, 44 insertions(+), 11 deletions(-)

diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index f918278768133..0d8fbcae493fa 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -31,7 +31,9 @@
 
 TEST_DAG_ID = 'unit_tests_ssh_test_op'
 TEST_CONN_ID = "conn_id_for_testing"
-TIMEOUT = 5
+CONN_TIMEOUT = 5
+CMD_TIMEOUT = 7
+TIMEOUT = 12
 DEFAULT_DATE = datetime(2017, 1, 1)
 COMMAND = "echo -n airflow"
 COMMAND_WITH_SUDO = "sudo " + COMMAND
@@ -67,9 +69,10 @@ def test_hook_created_correctly_with_timeout(self):
 
     def test_hook_created_correctly(self):
         conn_timeout = 20
+        cmd_timeout = 45
         ssh_id = "ssh_default"
         task = SSHOperator(
-            task_id="test", command=COMMAND, dag=self.dag, conn_timeout=conn_timeout, ssh_conn_id="ssh_default"
+            task_id="test", command=COMMAND, dag=self.dag, conn_timeout=conn_timeout, cmd_timeout=cmd_timeout, ssh_conn_id="ssh_default"
         )
         assert task is not None
 
@@ -151,7 +154,6 @@ def test_no_output_command(self):
     def test_arg_checking(self):
         # Exception should be raised if neither ssh_hook nor ssh_conn_id is provided
         with pytest.raises(AirflowException, match="Cannot operate without ssh_hook or ssh_conn_id."):
-            #  task_0 = SSHOperator(task_id="test", command=COMMAND, timeout=TIMEOUT, dag=self.dag)
             task_0 = SSHOperator(task_id="test", command=COMMAND, conn_timeout=TIMEOUT, cmd_timeout=TIMEOUT, dag=self.dag)
             task_0.execute(None)
 
@@ -161,8 +163,8 @@ def test_arg_checking(self):
             ssh_hook="string_rather_than_SSHHook",  # invalid ssh_hook
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            conn_timeout=TIMEOUT,
-            cmd_timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -175,9 +177,8 @@ def test_arg_checking(self):
             task_id="test_2",
             ssh_conn_id=TEST_CONN_ID,  # no ssh_hook provided
             command=COMMAND,
-            # timeout=TIMEOUT,
-            conn_timeout=TIMEOUT,
-            cmd_timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -192,9 +193,8 @@ def test_arg_checking(self):
             ssh_hook=self.hook,
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            # timeout=TIMEOUT,
-            conn_timeout=TIMEOUT,
-            cmd_timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -203,6 +203,39 @@ def test_arg_checking(self):
             pass
         assert task_3.ssh_hook.ssh_conn_id == self.hook.ssh_conn_id
 
+        # if valid conn_timeout, cmd_timeout, and timeout are provided, ignore timeout
+        task_4 = SSHOperator(
+            task_id="test_4",
+            ssh_conn_id=TEST_CONN_ID,
+            command=COMMAND,
+            timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
+            dag=self.dag,
+        )
+        try:
+            task_4.execute(None)
+        except Exception:
+            pass
+        assert task_4.conn_timeout == CONN_TIMEOUT
+        assert task_4.cmd_timeout == CMD_TIMEOUT
+        assert task_4.ssh_hook.conn_timeout == CONN_TIMEOUT
+
+        # if valid only timeout is provided, use for both conn_timeout and cmd_timeout
+        task_5 = SSHOperator(
+            task_id="test_5",
+            ssh_conn_id=TEST_CONN_ID,
+            command=COMMAND,
+            timeout=TIMEOUT,
+            dag=self.dag,
+        )
+        try:
+            task_5.execute(None)
+        except Exception:
+            pass
+        assert task_5.cmd_timeout == TIMEOUT
+        assert task_5.ssh_hook.conn_timeout == TIMEOUT
+
     @parameterized.expand(
         [
             (COMMAND, False, False),

From e2b67e6a571eed8e7b97aff8eb0e7fd2669af8f4 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Mon, 26 Jul 2021 13:19:47 -0400
Subject: [PATCH 05/42] Removed commented TODO for SSHHook changes

---
 airflow/providers/ssh/hooks/ssh.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index b4f524515e5e3..5f9c637b9d18b 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -165,11 +165,6 @@ def __init__(
                 if "conn_timeout" in extra_options:
                     self.conn_timeout = int(extra_options.get('conn_timeout'))
 
-                # TODO Rework this - I don't like that just having conn.extra set will override 
-                # self.conn_timeout = int(extra_options.get('conn_timeout',
-                    # extra_options.get('timeout', 10))
-                #)
-
                 if "compress" in extra_options and str(extra_options["compress"]).lower() == 'false':
                     self.compress = False
 

From b72d9be655155a655b45c7c69b8235d838936a02 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Mon, 26 Jul 2021 16:37:29 -0400
Subject: [PATCH 06/42] Fix static checks for operators/test_ssh.py and
 hooks/test_ssh.py

---
 tests/providers/ssh/hooks/test_ssh.py     |  8 ++------
 tests/providers/ssh/operators/test_ssh.py | 15 +++++++++++----
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/tests/providers/ssh/hooks/test_ssh.py b/tests/providers/ssh/hooks/test_ssh.py
index 18389be798ba7..4afa6fb12e060 100644
--- a/tests/providers/ssh/hooks/test_ssh.py
+++ b/tests/providers/ssh/hooks/test_ssh.py
@@ -133,11 +133,7 @@ def setUpClass(cls) -> None:
                 conn_id=cls.CONN_SSH_WITH_PRIVATE_KEY_EXTRA,
                 host='localhost',
                 conn_type='ssh',
-                extra=json.dumps(
-                    {
-                        "private_key": TEST_PRIVATE_KEY
-                    }
-                ),
+                extra=json.dumps({"private_key": TEST_PRIVATE_KEY}),
             )
         )
         db.merge_conn(
@@ -576,7 +572,7 @@ def test_ssh_connection_with_conn_timeout_extra(self, ssh_mock):
             port='port',
             username='username',
             timeout=10,
-            conn_timeout=15
+            conn_timeout=15,
         )
 
         with hook.get_conn():
diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index 0d8fbcae493fa..a73c2087c61fb 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -63,7 +63,7 @@ def test_hook_created_correctly_with_timeout(self):
         assert task is not None
 
         task.execute(None)
-        
+
         assert timeout == task.ssh_hook.conn_timeout
         assert ssh_id == task.ssh_hook.ssh_conn_id
 
@@ -72,12 +72,17 @@ def test_hook_created_correctly(self):
         cmd_timeout = 45
         ssh_id = "ssh_default"
         task = SSHOperator(
-            task_id="test", command=COMMAND, dag=self.dag, conn_timeout=conn_timeout, cmd_timeout=cmd_timeout, ssh_conn_id="ssh_default"
+            task_id="test",
+            command=COMMAND,
+            dag=self.dag,
+            conn_timeout=conn_timeout,
+            cmd_timeout=cmd_timeout,
+            ssh_conn_id="ssh_default",
         )
         assert task is not None
 
         task.execute(None)
-        
+
         assert conn_timeout == task.ssh_hook.conn_timeout
         assert ssh_id == task.ssh_hook.ssh_conn_id
 
@@ -154,7 +159,9 @@ def test_no_output_command(self):
     def test_arg_checking(self):
         # Exception should be raised if neither ssh_hook nor ssh_conn_id is provided
         with pytest.raises(AirflowException, match="Cannot operate without ssh_hook or ssh_conn_id."):
-            task_0 = SSHOperator(task_id="test", command=COMMAND, conn_timeout=TIMEOUT, cmd_timeout=TIMEOUT, dag=self.dag)
+            task_0 = SSHOperator(
+                task_id="test", command=COMMAND, conn_timeout=TIMEOUT, cmd_timeout=TIMEOUT, dag=self.dag
+            )
             task_0.execute(None)
 
         # if ssh_hook is invalid/not provided, use ssh_conn_id to create SSHHook

From af3e425a66c0b4d83ae741c00261876e6826f30d Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Sun, 25 Jul 2021 19:43:36 -0400
Subject: [PATCH 07/42] Issue 16363 - add conn_timeout and cmd_timeout params
 to SSHOperator; add conn_timoue param to SSHHook

---
 airflow/providers/ssh/hooks/ssh.py        |  40 +++++++-
 airflow/providers/ssh/operators/ssh.py    |  32 +++++-
 tests/providers/ssh/hooks/test_ssh.py     | 117 +++++++++++++++++++++-
 tests/providers/ssh/operators/test_ssh.py |  35 +++++--
 4 files changed, 207 insertions(+), 17 deletions(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index a8994bdece1b7..b4f524515e5e3 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -56,8 +56,10 @@ class SSHHook(BaseHook):
     :type key_file: str
     :param port: port of remote host to connect (Default is paramiko SSH_PORT)
     :type port: int
-    :param timeout: timeout for the attempt to connect to the remote_host.
+    :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host.
     :type timeout: int
+    :param conn_timeout: timeout for the attempt to connect to the remote_host.
+    :type conn_timeout: int
     :param keepalive_interval: send a keepalive packet to remote host every
         keepalive_interval seconds
     :type keepalive_interval: int
@@ -101,7 +103,8 @@ def __init__(
         password: Optional[str] = None,
         key_file: Optional[str] = None,
         port: Optional[int] = None,
-        timeout: int = 10,
+        timeout: Optional[int] = 10,
+        conn_timeout: Optional[int] = None,
         keepalive_interval: int = 30,
     ) -> None:
         super().__init__()
@@ -113,6 +116,7 @@ def __init__(
         self.pkey = None
         self.port = port
         self.timeout = timeout
+        self.conn_timeout = conn_timeout
         self.keepalive_interval = keepalive_interval
 
         # Default values, overridable from Connection
@@ -137,6 +141,7 @@ def __init__(
                 self.remote_host = conn.host
             if self.port is None:
                 self.port = conn.port
+
             if conn.extra is not None:
                 extra_options = conn.extra_dejson
                 if "key_file" in extra_options and self.key_file is None:
@@ -148,7 +153,22 @@ def __init__(
                     self.pkey = self._pkey_from_private_key(private_key, passphrase=private_key_passphrase)
 
                 if "timeout" in extra_options:
-                    self.timeout = int(extra_options["timeout"], 10)
+                    warnings.warn(
+                        'Extra option `timeout` is deprecated.'
+                        'Please use `conn_timeout` instead.'
+                        'The old option `timeout` will be removed in a future version.',
+                        DeprecationWarning,
+                        stacklevel=2,
+                    )
+                    self.timeout = int(extra_options.get('timeout', 10))
+
+                if "conn_timeout" in extra_options:
+                    self.conn_timeout = int(extra_options.get('conn_timeout'))
+
+                # TODO Rework this - I don't like that just having conn.extra set will override 
+                # self.conn_timeout = int(extra_options.get('conn_timeout',
+                    # extra_options.get('timeout', 10))
+                #)
 
                 if "compress" in extra_options and str(extra_options["compress"]).lower() == 'false':
                     self.compress = False
@@ -185,6 +205,18 @@ def __init__(
                     self.host_key = key_constructor(data=decoded_host_key)
                     self.no_host_key_check = False
 
+        if self.timeout:
+            warnings.warn(
+                'Parameter `timeout` is deprecated.'
+                'Please use `conn_timeout` instead.'
+                'The old option `timeout` will be removed in a future version.',
+                DeprecationWarning,
+                stacklevel=1,
+            )
+
+        if self.conn_timeout is None:
+            self.conn_timeout = self.timeout
+
         if self.pkey and self.key_file:
             raise AirflowException(
                 "Params key_file and private_key both provided.  Must provide no more than one."
@@ -253,7 +285,7 @@ def get_conn(self) -> paramiko.SSHClient:
         connect_kwargs = dict(
             hostname=self.remote_host,
             username=self.username,
-            timeout=self.timeout,
+            timeout=self.conn_timeout,
             compress=self.compress,
             port=self.port,
             sock=self.host_proxy,
diff --git a/airflow/providers/ssh/operators/ssh.py b/airflow/providers/ssh/operators/ssh.py
index c82e0e12c30a4..27380977a2abb 100644
--- a/airflow/providers/ssh/operators/ssh.py
+++ b/airflow/providers/ssh/operators/ssh.py
@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import warnings
 from base64 import b64encode
 from select import select
 from typing import Optional, Union
@@ -43,8 +44,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 timeout: (deprecated) timeout (in seconds) for executing the command. The default is 10 seconds.
     :type timeout: int
+    :param conn_timeout: timeout (in seconds) for maintaining the connection. The default is 10 seconds.
+    :type conn_timeout: int
+    :param cmd_timeout: timeout (in seconds) for executing the command. The default is 10 seconds.
+    :type cmd_timeout: int
     :param environment: a dict of shell environment variables. Note that the
         server will reject them silently if `AcceptEnv` is not set in SSH config.
     :type environment: dict
@@ -66,7 +71,9 @@ def __init__(
         ssh_conn_id: Optional[str] = None,
         remote_host: Optional[str] = None,
         command: Optional[str] = None,
-        timeout: int = 10,
+        timeout: Optional[int] = 10,
+        conn_timeout: Optional[int] = None,
+        cmd_timeout: Optional[int] = None,
         environment: Optional[dict] = None,
         get_pty: bool = False,
         **kwargs,
@@ -77,9 +84,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.cmd_timeout is None:
+            self.cmd_timeout = self.timeout
         self.environment = environment
         self.get_pty = (self.command.startswith('sudo') or get_pty) if self.command else get_pty
 
+        if self.timeout:
+            warnings.warn(
+                'Parameter `timeout` is deprecated.'
+                'Please use `conn_timeout` and `cmd_timeout` instead.'
+                'The old option `timeout` will be removed in a future version.',
+                DeprecationWarning,
+                stacklevel=1,
+            )
+
     def execute(self, context) -> Union[bytes, str, bool]:
         try:
             if self.ssh_conn_id:
@@ -89,7 +111,7 @@ def execute(self, context) -> Union[bytes, str, bool]:
                     self.log.info(
                         "ssh_hook is not provided or invalid. Trying ssh_conn_id to create SSHHook."
                     )
-                    self.ssh_hook = SSHHook(ssh_conn_id=self.ssh_conn_id, timeout=self.timeout)
+                    self.ssh_hook = SSHHook(ssh_conn_id=self.ssh_conn_id, conn_timeout=self.conn_timeout)
 
             if not self.ssh_hook:
                 raise AirflowException("Cannot operate without ssh_hook or ssh_conn_id.")
@@ -112,7 +134,7 @@ def execute(self, context) -> Union[bytes, str, bool]:
                 stdin, stdout, stderr = ssh_client.exec_command(
                     command=self.command,
                     get_pty=self.get_pty,
-                    timeout=self.timeout,
+                    timeout=self.cmd_timeout,
                     environment=self.environment,
                 )
                 # get channels
@@ -133,7 +155,7 @@ def execute(self, context) -> Union[bytes, str, bool]:
 
                 # read from both stdout and stderr
                 while not channel.closed or channel.recv_ready() or channel.recv_stderr_ready():
-                    readq, _, _ = select([channel], [], [], self.timeout)
+                    readq, _, _ = select([channel], [], [], self.cmd_timeout)
                     for recv in readq:
                         if recv.recv_ready():
                             line = stdout.channel.recv(len(recv.in_buffer))
diff --git a/tests/providers/ssh/hooks/test_ssh.py b/tests/providers/ssh/hooks/test_ssh.py
index 70ba4eb0ce642..18389be798ba7 100644
--- a/tests/providers/ssh/hooks/test_ssh.py
+++ b/tests/providers/ssh/hooks/test_ssh.py
@@ -68,6 +68,9 @@ def generate_host_key(pkey: paramiko.PKey):
 TEST_PKEY_ECDSA = paramiko.ECDSAKey.generate()
 TEST_PRIVATE_KEY_ECDSA = generate_key_string(pkey=TEST_PKEY_ECDSA)
 
+TEST_TIMEOUT = 20
+TEST_CONN_TIMEOUT = 30
+
 PASSPHRASE = ''.join(random.choice(string.ascii_letters) for i in range(10))
 TEST_ENCRYPTED_PRIVATE_KEY = generate_key_string(pkey=TEST_PKEY, passphrase=PASSPHRASE)
 
@@ -76,6 +79,8 @@ class TestSSHHook(unittest.TestCase):
     CONN_SSH_WITH_PRIVATE_KEY_EXTRA = 'ssh_with_private_key_extra'
     CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA = 'ssh_with_private_key_ecdsa_extra'
     CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA = 'ssh_with_private_key_passphrase_extra'
+    CONN_SSH_WITH_TIMEOUT_EXTRA = 'ssh_with_timeout_extra'
+    CONN_SSH_WITH_CONN_TIMEOUT_EXTRA = 'ssh_with_conn_timeout_extra'
     CONN_SSH_WITH_EXTRA = 'ssh_with_extra'
     CONN_SSH_WITH_EXTRA_FALSE_LOOK_FOR_KEYS = 'ssh_with_extra_false_look_for_keys'
     CONN_SSH_WITH_HOST_KEY_EXTRA = 'ssh_with_host_key_extra'
@@ -91,6 +96,8 @@ def tearDownClass(cls) -> None:
                 cls.CONN_SSH_WITH_PRIVATE_KEY_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA,
+                cls.CONN_SSH_WITH_TIMEOUT_EXTRA,
+                cls.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA,
                 cls.CONN_SSH_WITH_EXTRA,
                 cls.CONN_SSH_WITH_HOST_KEY_EXTRA,
                 cls.CONN_SSH_WITH_HOST_KEY_EXTRA_WITH_TYPE,
@@ -128,7 +135,7 @@ def setUpClass(cls) -> None:
                 conn_type='ssh',
                 extra=json.dumps(
                     {
-                        "private_key": TEST_PRIVATE_KEY,
+                        "private_key": TEST_PRIVATE_KEY
                     }
                 ),
             )
@@ -151,6 +158,22 @@ def setUpClass(cls) -> None:
                 extra=json.dumps({"private_key": TEST_PRIVATE_KEY_ECDSA}),
             )
         )
+        db.merge_conn(
+            Connection(
+                conn_id=cls.CONN_SSH_WITH_TIMEOUT_EXTRA,
+                host='localhost',
+                conn_type='ssh',
+                extra=json.dumps({"timeout": TEST_TIMEOUT}),
+            )
+        )
+        db.merge_conn(
+            Connection(
+                conn_id=cls.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA,
+                host='localhost',
+                conn_type='ssh',
+                extra=json.dumps({"conn_timeout": TEST_CONN_TIMEOUT}),
+            )
+        )
         db.merge_conn(
             Connection(
                 conn_id=cls.CONN_SSH_WITH_HOST_KEY_EXTRA,
@@ -475,6 +498,98 @@ 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
+        )
+
+        with hook.get_conn():
+            ssh_mock.return_value.connect.assert_called_once_with(
+                hostname='remote_host',
+                username='username',
+                timeout=30,
+                compress=True,
+                port='port',
+                sock=None,
+                look_for_keys=True,
+            )
+
     def test_openssh_private_key(self):
         # Paramiko behaves differently with OpenSSH generated keys to paramiko
         # generated keys, so we need a test one.
diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index e39b8fbb8c156..f918278768133 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -52,7 +52,7 @@ def setUp(self):
         self.hook = hook
         self.dag = dag
 
-    def test_hook_created_correctly(self):
+    def test_hook_created_correctly_with_timeout(self):
         timeout = 20
         ssh_id = "ssh_default"
         task = SSHOperator(
@@ -61,8 +61,21 @@ def test_hook_created_correctly(self):
         assert task is not None
 
         task.execute(None)
+        
+        assert timeout == task.ssh_hook.conn_timeout
+        assert ssh_id == task.ssh_hook.ssh_conn_id
 
-        assert timeout == task.ssh_hook.timeout
+    def test_hook_created_correctly(self):
+        conn_timeout = 20
+        ssh_id = "ssh_default"
+        task = SSHOperator(
+            task_id="test", command=COMMAND, dag=self.dag, conn_timeout=conn_timeout, ssh_conn_id="ssh_default"
+        )
+        assert task is not None
+
+        task.execute(None)
+        
+        assert conn_timeout == task.ssh_hook.conn_timeout
         assert ssh_id == task.ssh_hook.ssh_conn_id
 
     @conf_vars({('core', 'enable_xcom_pickling'): 'False'})
@@ -138,7 +151,8 @@ def test_no_output_command(self):
     def test_arg_checking(self):
         # Exception should be raised if neither ssh_hook nor ssh_conn_id is provided
         with pytest.raises(AirflowException, match="Cannot operate without ssh_hook or ssh_conn_id."):
-            task_0 = SSHOperator(task_id="test", command=COMMAND, timeout=TIMEOUT, dag=self.dag)
+            #  task_0 = SSHOperator(task_id="test", command=COMMAND, timeout=TIMEOUT, dag=self.dag)
+            task_0 = SSHOperator(task_id="test", command=COMMAND, conn_timeout=TIMEOUT, cmd_timeout=TIMEOUT, dag=self.dag)
             task_0.execute(None)
 
         # if ssh_hook is invalid/not provided, use ssh_conn_id to create SSHHook
@@ -147,7 +161,8 @@ def test_arg_checking(self):
             ssh_hook="string_rather_than_SSHHook",  # invalid ssh_hook
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -160,7 +175,9 @@ def test_arg_checking(self):
             task_id="test_2",
             ssh_conn_id=TEST_CONN_ID,  # no ssh_hook provided
             command=COMMAND,
-            timeout=TIMEOUT,
+            # timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -175,7 +192,9 @@ def test_arg_checking(self):
             ssh_hook=self.hook,
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            timeout=TIMEOUT,
+            # timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -198,7 +217,9 @@ def test_get_pyt_set_correctly(self, command, get_pty_in, get_pty_out):
             task_id="test",
             ssh_hook=self.hook,
             command=command,
-            timeout=TIMEOUT,
+            # timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             get_pty=get_pty_in,
             dag=self.dag,
         )

From 99041154b3094e88264c9b5f9026e6cbcab70c9e Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Sun, 25 Jul 2021 20:24:56 -0400
Subject: [PATCH 08/42] Update SSHOperator tests for new conn_timeout and
 cmd_timeout params

---
 tests/providers/ssh/operators/test_ssh.py | 55 ++++++++++++++++++-----
 1 file changed, 44 insertions(+), 11 deletions(-)

diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index f918278768133..0d8fbcae493fa 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -31,7 +31,9 @@
 
 TEST_DAG_ID = 'unit_tests_ssh_test_op'
 TEST_CONN_ID = "conn_id_for_testing"
-TIMEOUT = 5
+CONN_TIMEOUT = 5
+CMD_TIMEOUT = 7
+TIMEOUT = 12
 DEFAULT_DATE = datetime(2017, 1, 1)
 COMMAND = "echo -n airflow"
 COMMAND_WITH_SUDO = "sudo " + COMMAND
@@ -67,9 +69,10 @@ def test_hook_created_correctly_with_timeout(self):
 
     def test_hook_created_correctly(self):
         conn_timeout = 20
+        cmd_timeout = 45
         ssh_id = "ssh_default"
         task = SSHOperator(
-            task_id="test", command=COMMAND, dag=self.dag, conn_timeout=conn_timeout, ssh_conn_id="ssh_default"
+            task_id="test", command=COMMAND, dag=self.dag, conn_timeout=conn_timeout, cmd_timeout=cmd_timeout, ssh_conn_id="ssh_default"
         )
         assert task is not None
 
@@ -151,7 +154,6 @@ def test_no_output_command(self):
     def test_arg_checking(self):
         # Exception should be raised if neither ssh_hook nor ssh_conn_id is provided
         with pytest.raises(AirflowException, match="Cannot operate without ssh_hook or ssh_conn_id."):
-            #  task_0 = SSHOperator(task_id="test", command=COMMAND, timeout=TIMEOUT, dag=self.dag)
             task_0 = SSHOperator(task_id="test", command=COMMAND, conn_timeout=TIMEOUT, cmd_timeout=TIMEOUT, dag=self.dag)
             task_0.execute(None)
 
@@ -161,8 +163,8 @@ def test_arg_checking(self):
             ssh_hook="string_rather_than_SSHHook",  # invalid ssh_hook
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            conn_timeout=TIMEOUT,
-            cmd_timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -175,9 +177,8 @@ def test_arg_checking(self):
             task_id="test_2",
             ssh_conn_id=TEST_CONN_ID,  # no ssh_hook provided
             command=COMMAND,
-            # timeout=TIMEOUT,
-            conn_timeout=TIMEOUT,
-            cmd_timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -192,9 +193,8 @@ def test_arg_checking(self):
             ssh_hook=self.hook,
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            # timeout=TIMEOUT,
-            conn_timeout=TIMEOUT,
-            cmd_timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -203,6 +203,39 @@ def test_arg_checking(self):
             pass
         assert task_3.ssh_hook.ssh_conn_id == self.hook.ssh_conn_id
 
+        # if valid conn_timeout, cmd_timeout, and timeout are provided, ignore timeout
+        task_4 = SSHOperator(
+            task_id="test_4",
+            ssh_conn_id=TEST_CONN_ID,
+            command=COMMAND,
+            timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
+            dag=self.dag,
+        )
+        try:
+            task_4.execute(None)
+        except Exception:
+            pass
+        assert task_4.conn_timeout == CONN_TIMEOUT
+        assert task_4.cmd_timeout == CMD_TIMEOUT
+        assert task_4.ssh_hook.conn_timeout == CONN_TIMEOUT
+
+        # if valid only timeout is provided, use for both conn_timeout and cmd_timeout
+        task_5 = SSHOperator(
+            task_id="test_5",
+            ssh_conn_id=TEST_CONN_ID,
+            command=COMMAND,
+            timeout=TIMEOUT,
+            dag=self.dag,
+        )
+        try:
+            task_5.execute(None)
+        except Exception:
+            pass
+        assert task_5.cmd_timeout == TIMEOUT
+        assert task_5.ssh_hook.conn_timeout == TIMEOUT
+
     @parameterized.expand(
         [
             (COMMAND, False, False),

From 5231613b9bb854a95c5836efbc7f86cefc8a036a Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Mon, 26 Jul 2021 13:19:47 -0400
Subject: [PATCH 09/42] Removed commented TODO for SSHHook changes

---
 airflow/providers/ssh/hooks/ssh.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index b4f524515e5e3..5f9c637b9d18b 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -165,11 +165,6 @@ def __init__(
                 if "conn_timeout" in extra_options:
                     self.conn_timeout = int(extra_options.get('conn_timeout'))
 
-                # TODO Rework this - I don't like that just having conn.extra set will override 
-                # self.conn_timeout = int(extra_options.get('conn_timeout',
-                    # extra_options.get('timeout', 10))
-                #)
-
                 if "compress" in extra_options and str(extra_options["compress"]).lower() == 'false':
                     self.compress = False
 

From f3395821fd4d967e041ca91d4d3411928ad6f236 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Mon, 26 Jul 2021 16:37:29 -0400
Subject: [PATCH 10/42] Fix static checks for operators/test_ssh.py and
 hooks/test_ssh.py

---
 tests/providers/ssh/hooks/test_ssh.py     |  8 ++------
 tests/providers/ssh/operators/test_ssh.py | 15 +++++++++++----
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/tests/providers/ssh/hooks/test_ssh.py b/tests/providers/ssh/hooks/test_ssh.py
index 18389be798ba7..4afa6fb12e060 100644
--- a/tests/providers/ssh/hooks/test_ssh.py
+++ b/tests/providers/ssh/hooks/test_ssh.py
@@ -133,11 +133,7 @@ def setUpClass(cls) -> None:
                 conn_id=cls.CONN_SSH_WITH_PRIVATE_KEY_EXTRA,
                 host='localhost',
                 conn_type='ssh',
-                extra=json.dumps(
-                    {
-                        "private_key": TEST_PRIVATE_KEY
-                    }
-                ),
+                extra=json.dumps({"private_key": TEST_PRIVATE_KEY}),
             )
         )
         db.merge_conn(
@@ -576,7 +572,7 @@ def test_ssh_connection_with_conn_timeout_extra(self, ssh_mock):
             port='port',
             username='username',
             timeout=10,
-            conn_timeout=15
+            conn_timeout=15,
         )
 
         with hook.get_conn():
diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index 0d8fbcae493fa..a73c2087c61fb 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -63,7 +63,7 @@ def test_hook_created_correctly_with_timeout(self):
         assert task is not None
 
         task.execute(None)
-        
+
         assert timeout == task.ssh_hook.conn_timeout
         assert ssh_id == task.ssh_hook.ssh_conn_id
 
@@ -72,12 +72,17 @@ def test_hook_created_correctly(self):
         cmd_timeout = 45
         ssh_id = "ssh_default"
         task = SSHOperator(
-            task_id="test", command=COMMAND, dag=self.dag, conn_timeout=conn_timeout, cmd_timeout=cmd_timeout, ssh_conn_id="ssh_default"
+            task_id="test",
+            command=COMMAND,
+            dag=self.dag,
+            conn_timeout=conn_timeout,
+            cmd_timeout=cmd_timeout,
+            ssh_conn_id="ssh_default",
         )
         assert task is not None
 
         task.execute(None)
-        
+
         assert conn_timeout == task.ssh_hook.conn_timeout
         assert ssh_id == task.ssh_hook.ssh_conn_id
 
@@ -154,7 +159,9 @@ def test_no_output_command(self):
     def test_arg_checking(self):
         # Exception should be raised if neither ssh_hook nor ssh_conn_id is provided
         with pytest.raises(AirflowException, match="Cannot operate without ssh_hook or ssh_conn_id."):
-            task_0 = SSHOperator(task_id="test", command=COMMAND, conn_timeout=TIMEOUT, cmd_timeout=TIMEOUT, dag=self.dag)
+            task_0 = SSHOperator(
+                task_id="test", command=COMMAND, conn_timeout=TIMEOUT, cmd_timeout=TIMEOUT, dag=self.dag
+            )
             task_0.execute(None)
 
         # if ssh_hook is invalid/not provided, use ssh_conn_id to create SSHHook

From 98dc572974369a01543c4214b6566a8d577b9083 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Wed, 28 Jul 2021 20:10:01 -0400
Subject: [PATCH 11/42] Set deprecated timeout param default value to None; add
 more tests

---
 airflow/providers/ssh/hooks/ssh.py        | 10 ++++---
 airflow/providers/ssh/operators/ssh.py    | 12 +++++----
 tests/providers/ssh/hooks/test_ssh.py     | 32 +++++++++++++++++++++++
 tests/providers/ssh/operators/test_ssh.py |  2 +-
 4 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index 5f9c637b9d18b..3ddae25df8f9c 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -34,6 +34,8 @@
 except ImportError:
     from getpass import getuser
 
+TIMEOUT_DEFAULT = 10
+
 
 class SSHHook(BaseHook):
     """
@@ -56,10 +58,10 @@ class SSHHook(BaseHook):
     :type key_file: str
     :param port: port of remote host to connect (Default is paramiko SSH_PORT)
     :type port: int
-    :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host.
-    :type timeout: int
     :param conn_timeout: timeout for the attempt to connect to the remote_host.
     :type conn_timeout: int
+    :param timeout: (Deprecated). 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
@@ -103,7 +105,7 @@ def __init__(
         password: Optional[str] = None,
         key_file: Optional[str] = None,
         port: Optional[int] = None,
-        timeout: Optional[int] = 10,
+        timeout: Optional[int] = None,
         conn_timeout: Optional[int] = None,
         keepalive_interval: int = 30,
     ) -> None:
@@ -210,7 +212,7 @@ def __init__(
             )
 
         if self.conn_timeout is None:
-            self.conn_timeout = self.timeout
+            self.conn_timeout = self.timeout if self.timeout else TIMEOUT_DEFAULT
 
         if self.pkey and self.key_file:
             raise AirflowException(
diff --git a/airflow/providers/ssh/operators/ssh.py b/airflow/providers/ssh/operators/ssh.py
index 27380977a2abb..b99a6ba97a510 100644
--- a/airflow/providers/ssh/operators/ssh.py
+++ b/airflow/providers/ssh/operators/ssh.py
@@ -26,6 +26,8 @@
 from airflow.models import BaseOperator
 from airflow.providers.ssh.hooks.ssh import SSHHook
 
+TIMEOUT_DEFAULT = 10
+
 
 class SSHOperator(BaseOperator):
     """
@@ -44,12 +46,12 @@ class SSHOperator(BaseOperator):
     :type remote_host: str
     :param command: command to execute on remote host. (templated)
     :type command: str
-    :param timeout: (deprecated) timeout (in seconds) for executing the command. The default is 10 seconds.
-    :type timeout: int
     :param conn_timeout: timeout (in seconds) for maintaining the connection. The default is 10 seconds.
     :type conn_timeout: int
     :param cmd_timeout: timeout (in seconds) for executing the command. The default is 10 seconds.
     :type cmd_timeout: int
+    :param timeout: (deprecated) timeout (in seconds) for executing the command. The default is 10 seconds.
+    :type timeout: int
     :param environment: a dict of shell environment variables. Note that the
         server will reject them silently if `AcceptEnv` is not set in SSH config.
     :type environment: dict
@@ -71,7 +73,7 @@ def __init__(
         ssh_conn_id: Optional[str] = None,
         remote_host: Optional[str] = None,
         command: Optional[str] = None,
-        timeout: Optional[int] = 10,
+        timeout: Optional[int] = None,
         conn_timeout: Optional[int] = None,
         cmd_timeout: Optional[int] = None,
         environment: Optional[dict] = None,
@@ -87,9 +89,9 @@ def __init__(
         self.conn_timeout = conn_timeout
         self.cmd_timeout = cmd_timeout
         if self.conn_timeout is None:
-            self.conn_timeout = self.timeout
+            self.conn_timeout = self.timeout if self.timeout else TIMEOUT_DEFAULT
         if self.cmd_timeout is None:
-            self.cmd_timeout = self.timeout
+            self.cmd_timeout = self.timeout if self.timeout else TIMEOUT_DEFAULT
         self.environment = environment
         self.get_pty = (self.command.startswith('sudo') or get_pty) if self.command else get_pty
 
diff --git a/tests/providers/ssh/hooks/test_ssh.py b/tests/providers/ssh/hooks/test_ssh.py
index 4afa6fb12e060..139c9d0c18435 100644
--- a/tests/providers/ssh/hooks/test_ssh.py
+++ b/tests/providers/ssh/hooks/test_ssh.py
@@ -81,6 +81,7 @@ class TestSSHHook(unittest.TestCase):
     CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA = 'ssh_with_private_key_passphrase_extra'
     CONN_SSH_WITH_TIMEOUT_EXTRA = 'ssh_with_timeout_extra'
     CONN_SSH_WITH_CONN_TIMEOUT_EXTRA = 'ssh_with_conn_timeout_extra'
+    CONN_SSH_WITH_TIMEOUT_AND_CONN_TIMEOUT_EXTRA = 'ssh_with_timeout_and_conn_timeout_extra'
     CONN_SSH_WITH_EXTRA = 'ssh_with_extra'
     CONN_SSH_WITH_EXTRA_FALSE_LOOK_FOR_KEYS = 'ssh_with_extra_false_look_for_keys'
     CONN_SSH_WITH_HOST_KEY_EXTRA = 'ssh_with_host_key_extra'
@@ -98,6 +99,7 @@ def tearDownClass(cls) -> None:
                 cls.CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA,
                 cls.CONN_SSH_WITH_TIMEOUT_EXTRA,
                 cls.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA,
+                cls.CONN_SSH_WITH_TIMEOUT_AND_CONN_TIMEOUT_EXTRA,
                 cls.CONN_SSH_WITH_EXTRA,
                 cls.CONN_SSH_WITH_HOST_KEY_EXTRA,
                 cls.CONN_SSH_WITH_HOST_KEY_EXTRA_WITH_TYPE,
@@ -170,6 +172,14 @@ def setUpClass(cls) -> None:
                 extra=json.dumps({"conn_timeout": TEST_CONN_TIMEOUT}),
             )
         )
+        db.merge_conn(
+            Connection(
+                conn_id=cls.CONN_SSH_WITH_TIMEOUT_AND_CONN_TIMEOUT_EXTRA,
+                host='localhost',
+                conn_type='ssh',
+                extra=json.dumps({"conn_timeout": TEST_CONN_TIMEOUT, 'timeout': TEST_TIMEOUT}),
+            )
+        )
         db.merge_conn(
             Connection(
                 conn_id=cls.CONN_SSH_WITH_HOST_KEY_EXTRA,
@@ -586,6 +596,28 @@ def test_ssh_connection_with_conn_timeout_extra(self, ssh_mock):
                 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,
+        )
+
+        with hook.get_conn():
+            ssh_mock.return_value.connect.assert_called_once_with(
+                hostname='remote_host',
+                username='username',
+                timeout=30,
+                compress=True,
+                port='port',
+                sock=None,
+                look_for_keys=True,
+            )
+
     def test_openssh_private_key(self):
         # Paramiko behaves differently with OpenSSH generated keys to paramiko
         # generated keys, so we need a test one.
diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index a73c2087c61fb..4d93a100745c4 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -228,7 +228,7 @@ def test_arg_checking(self):
         assert task_4.cmd_timeout == CMD_TIMEOUT
         assert task_4.ssh_hook.conn_timeout == CONN_TIMEOUT
 
-        # if valid only timeout is provided, use for both conn_timeout and cmd_timeout
+        # if only timeout is provided, use for both conn_timeout and cmd_timeout
         task_5 = SSHOperator(
             task_id="test_5",
             ssh_conn_id=TEST_CONN_ID,

From 9553020cd0a4601d08272b98287ece5cdfb3b465 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Wed, 28 Jul 2021 21:20:53 -0400
Subject: [PATCH 12/42] Remove unnecessary extra_options.get() calls for
 timeout and conn_timeout

---
 airflow/providers/ssh/hooks/ssh.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index 3ddae25df8f9c..e399fd8baaf75 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -162,10 +162,10 @@ def __init__(
                         DeprecationWarning,
                         stacklevel=2,
                     )
-                    self.timeout = int(extra_options.get('timeout', 10))
+                    self.timeout = int(extra_options['timeout'])
 
                 if "conn_timeout" in extra_options:
-                    self.conn_timeout = int(extra_options.get('conn_timeout'))
+                    self.conn_timeout = int(extra_options['conn_timeout'])
 
                 if "compress" in extra_options and str(extra_options["compress"]).lower() == 'false':
                     self.compress = False

From 40e7f43b1d7017f3cd8b44bf4b20ff8a11cead49 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Sun, 1 Aug 2021 18:14:21 -0400
Subject: [PATCH 13/42] Add reference to new timeout parameters in docstrings
 for deprecated timeout

---
 airflow/providers/ssh/hooks/ssh.py     | 2 +-
 airflow/providers/ssh/operators/ssh.py | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index e399fd8baaf75..024a66a46f170 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -60,7 +60,7 @@ class SSHHook(BaseHook):
     :type port: int
     :param conn_timeout: timeout for the attempt to connect to the remote_host.
     :type conn_timeout: int
-    :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host.
+    :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host. Use conn_timeout instead.
     :type timeout: int
     :param keepalive_interval: send a keepalive packet to remote host every
         keepalive_interval seconds
diff --git a/airflow/providers/ssh/operators/ssh.py b/airflow/providers/ssh/operators/ssh.py
index b99a6ba97a510..ebec70bf2d220 100644
--- a/airflow/providers/ssh/operators/ssh.py
+++ b/airflow/providers/ssh/operators/ssh.py
@@ -51,6 +51,7 @@ class SSHOperator(BaseOperator):
     :param cmd_timeout: timeout (in seconds) for executing the command. The default is 10 seconds.
     :type cmd_timeout: int
     :param timeout: (deprecated) timeout (in seconds) for executing the command. The default is 10 seconds.
+        Use conn_timeout and cmd_timeout parameters instead.
     :type timeout: int
     :param environment: a dict of shell environment variables. Note that the
         server will reject them silently if `AcceptEnv` is not set in SSH config.

From 86c731a9cd270ad8979b96bf4104ad3720b9b38f Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Sun, 1 Aug 2021 20:21:21 -0400
Subject: [PATCH 14/42] Move timeout docstring text to 2nd line to keep line
 within 110 characters

---
 airflow/providers/ssh/hooks/ssh.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index 024a66a46f170..332de6eff59cf 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -60,7 +60,8 @@ class SSHHook(BaseHook):
     :type port: int
     :param conn_timeout: timeout for the attempt to connect to the remote_host.
     :type conn_timeout: int
-    :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host. Use conn_timeout instead.
+    :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host.
+        Use conn_timeout instead.
     :type timeout: int
     :param keepalive_interval: send a keepalive packet to remote host every
         keepalive_interval seconds

From 62ef8c548750b69822b7817c735a39fbbc5aae5c Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Sun, 1 Aug 2021 20:40:02 -0400
Subject: [PATCH 15/42] Favor conn_timeout parameter over extra options

---
 airflow/providers/ssh/hooks/ssh.py    | 2 +-
 tests/providers/ssh/hooks/test_ssh.py | 6 ++++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index 332de6eff59cf..9ea2a86cc4515 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -165,7 +165,7 @@ def __init__(
                     )
                     self.timeout = int(extra_options['timeout'])
 
-                if "conn_timeout" in extra_options:
+                if "conn_timeout" in extra_options and self.conn_timeout is None:
                     self.conn_timeout = int(extra_options['conn_timeout'])
 
                 if "compress" in extra_options and str(extra_options["compress"]).lower() == 'false':
diff --git a/tests/providers/ssh/hooks/test_ssh.py b/tests/providers/ssh/hooks/test_ssh.py
index 139c9d0c18435..7f3c9c2167291 100644
--- a/tests/providers/ssh/hooks/test_ssh.py
+++ b/tests/providers/ssh/hooks/test_ssh.py
@@ -585,11 +585,12 @@ def test_ssh_connection_with_conn_timeout_extra(self, ssh_mock):
             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=30,
+                timeout=15,
                 compress=True,
                 port='port',
                 sock=None,
@@ -607,11 +608,12 @@ def test_ssh_connection_with_timeout_extra_and_conn_timeout_extra(self, ssh_mock
             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=30,
+                timeout=15,
                 compress=True,
                 port='port',
                 sock=None,

From 87bbbaabb642b67093ee816929d285c08f3bf409 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Wed, 4 Aug 2021 22:03:13 -0400
Subject: [PATCH 16/42] Add conn_timeout to SSH Connection docs and mark
 timeout as deprecated

---
 docs/apache-airflow-providers-ssh/connections/ssh.rst | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/docs/apache-airflow-providers-ssh/connections/ssh.rst b/docs/apache-airflow-providers-ssh/connections/ssh.rst
index ec9e6d165539e..bcfb01e128e18 100644
--- a/docs/apache-airflow-providers-ssh/connections/ssh.rst
+++ b/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.
+    * ``timeout`` - Deprecated - use conn_timeout instead.
     * ``compress`` - ``true`` to ask the remote client/server to compress traffic; ``false`` to refuse compression. Default is ``true``.
     * ``no_host_key_check`` - Set to ``false`` to restrict connecting to hosts with no entries in ``~/.ssh/known_hosts`` (Hosts file). This provides maximum protection against trojan horse attacks, but can be troublesome when the ``/etc/ssh/ssh_known_hosts`` file is poorly maintained or connections to new hosts are frequently made. This option forces the user to manually add all new hosts. Default is ``true``, ssh will automatically add new host keys to the user known hosts files.
     * ``allow_host_key_change`` - Set to ``true`` if you want to allow connecting to hosts that has host key changed or when you get 'REMOTE HOST IDENTIFICATION HAS CHANGED' error.  This wont protect against Man-In-The-Middle attacks. Other possible solution is to remove the host entry from ``~/.ssh/known_hosts`` file. Default is ``false``.
@@ -60,7 +61,7 @@ Extra (optional)
 
        {
           "key_file": "/home/airflow/.ssh/id_rsa",
-          "timeout": "10",
+          "conn_timeout": "10",
           "compress": "false",
           "look_for_keys": "false",
           "allow_host_key_change": "false",
@@ -75,7 +76,7 @@ Extra (optional)
 
     .. code-block:: bash
 
-        export AIRFLOW_CONN_MAIN_SERVER='ssh://user:pass@localhost:22?timeout=10&compress=false&no_host_key_check=false&allow_host_key_change=true&key_file=%2Fhome%2Fairflow%2F.ssh%2Fid_rsa'
+        export AIRFLOW_CONN_MAIN_SERVER='ssh://user:pass@localhost:22?conn_timeout=10&compress=false&no_host_key_check=false&allow_host_key_change=true&key_file=%2Fhome%2Fairflow%2F.ssh%2Fid_rsa'
 
     Example connection string with ``private_key`` (actual private key provided in connection):
 

From df6aef87f24f4ad1b59eb3b9fa2edccc2861d052 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Mon, 23 Aug 2021 23:09:19 -0400
Subject: [PATCH 17/42] Fix issue with defaults for timeouts in
 SSHHook/SSHOperator; add parametrized unit tests for SSHHook

---
 airflow/providers/ssh/operators/ssh.py    | 10 ++--
 tests/providers/ssh/hooks/test_ssh.py     | 66 +++++++++++++++++++++++
 tests/providers/ssh/operators/test_ssh.py | 19 +++++++
 3 files changed, 91 insertions(+), 4 deletions(-)

diff --git a/airflow/providers/ssh/operators/ssh.py b/airflow/providers/ssh/operators/ssh.py
index b0dc19eb20277..7f97b031b414d 100644
--- a/airflow/providers/ssh/operators/ssh.py
+++ b/airflow/providers/ssh/operators/ssh.py
@@ -26,7 +26,7 @@
 from airflow.models import BaseOperator
 from airflow.providers.ssh.hooks.ssh import SSHHook
 
-TIMEOUT_DEFAULT = 10
+CMD_TIMEOUT = 10
 
 
 class SSHOperator(BaseOperator):
@@ -47,6 +47,8 @@ class SSHOperator(BaseOperator):
     :param command: command to execute on remote host. (templated)
     :type command: str
     :param conn_timeout: timeout (in seconds) for maintaining the connection. The default is 10 seconds.
+        Nullable. If provided, it will replace the `conn_timeout` which was
+        predefined in the connection of `ssh_conn_id`.
     :type conn_timeout: int
     :param cmd_timeout: timeout (in seconds) for executing the command. The default is 10 seconds.
     :type cmd_timeout: int
@@ -89,10 +91,10 @@ def __init__(
         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
+        if self.conn_timeout is None and self.timeout:
+            self.conn_timeout = self.timeout
         if self.cmd_timeout is None:
-            self.cmd_timeout = self.timeout if self.timeout else TIMEOUT_DEFAULT
+            self.cmd_timeout = self.timeout if self.timeout else CMD_TIMEOUT
         self.environment = environment
         self.get_pty = (self.command.startswith('sudo') or get_pty) if self.command else get_pty
 
diff --git a/tests/providers/ssh/hooks/test_ssh.py b/tests/providers/ssh/hooks/test_ssh.py
index 7f3c9c2167291..be63093ccfbff 100644
--- a/tests/providers/ssh/hooks/test_ssh.py
+++ b/tests/providers/ssh/hooks/test_ssh.py
@@ -25,6 +25,7 @@
 from unittest import mock
 
 import paramiko
+from parameterized import parameterized
 
 from airflow import settings
 from airflow.models import Connection
@@ -76,6 +77,7 @@ def generate_host_key(pkey: paramiko.PKey):
 
 
 class TestSSHHook(unittest.TestCase):
+    CONN_SSH_WITH_NO_EXTRA = 'ssh_with_no_extra'
     CONN_SSH_WITH_PRIVATE_KEY_EXTRA = 'ssh_with_private_key_extra'
     CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA = 'ssh_with_private_key_ecdsa_extra'
     CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA = 'ssh_with_private_key_passphrase_extra'
@@ -94,6 +96,7 @@ class TestSSHHook(unittest.TestCase):
     def tearDownClass(cls) -> None:
         with create_session() as session:
             conns_to_reset = [
+                cls.CONN_SSH_WITH_NO_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA,
@@ -113,6 +116,14 @@ def tearDownClass(cls) -> None:
 
     @classmethod
     def setUpClass(cls) -> None:
+        db.merge_conn(
+            Connection(
+                conn_id=cls.CONN_SSH_WITH_NO_EXTRA,
+                host='localhost',
+                conn_type='ssh',
+                extra=None,
+            )
+        )
         db.merge_conn(
             Connection(
                 conn_id=cls.CONN_SSH_WITH_EXTRA,
@@ -620,6 +631,61 @@ def test_ssh_connection_with_timeout_extra_and_conn_timeout_extra(self, ssh_mock
                 look_for_keys=True,
             )
 
+    @parameterized.expand(
+        [
+            (TEST_TIMEOUT, TEST_CONN_TIMEOUT, True, True, TEST_CONN_TIMEOUT),
+            (TEST_TIMEOUT, TEST_CONN_TIMEOUT, True, False, TEST_CONN_TIMEOUT),
+            (TEST_TIMEOUT, TEST_CONN_TIMEOUT, False, True, TEST_CONN_TIMEOUT),
+            (TEST_TIMEOUT, TEST_CONN_TIMEOUT, False, False, TEST_CONN_TIMEOUT),
+            (TEST_TIMEOUT, None, True, True, TEST_CONN_TIMEOUT),
+            (TEST_TIMEOUT, None, True, False, TEST_TIMEOUT),
+            (TEST_TIMEOUT, None, False, True, TEST_CONN_TIMEOUT),
+            (TEST_TIMEOUT, None, False, False, TEST_TIMEOUT),
+            (None, TEST_CONN_TIMEOUT, True, True, TEST_CONN_TIMEOUT),
+            (None, TEST_CONN_TIMEOUT, True, False, TEST_CONN_TIMEOUT),
+            (None, TEST_CONN_TIMEOUT, False, True, TEST_CONN_TIMEOUT),
+            (None, TEST_CONN_TIMEOUT, False, False, TEST_CONN_TIMEOUT),
+            (None, None, True, True, TEST_CONN_TIMEOUT),
+            (None, None, True, False, TEST_TIMEOUT),
+            (None, None, False, True, TEST_CONN_TIMEOUT),
+            (None, None, False, False, 10),
+        ]
+    )
+    @mock.patch('airflow.providers.ssh.hooks.ssh.paramiko.SSHClient')
+    def test_ssh_connection_with_all_timeout_param_and_extra_combinations(
+        self, timeout, conn_timeout, timeoutextra, conn_timeoutextra, expected_value, ssh_mock
+    ):
+
+        if timeoutextra and conn_timeoutextra:
+            ssh_conn_id = self.CONN_SSH_WITH_TIMEOUT_AND_CONN_TIMEOUT_EXTRA
+        elif timeoutextra and not conn_timeoutextra:
+            ssh_conn_id = self.CONN_SSH_WITH_TIMEOUT_EXTRA
+        elif not timeoutextra and conn_timeoutextra:
+            ssh_conn_id = self.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA
+        else:
+            ssh_conn_id = self.CONN_SSH_WITH_NO_EXTRA
+
+        hook = SSHHook(
+            ssh_conn_id=ssh_conn_id,
+            remote_host='remote_host',
+            port='port',
+            username='username',
+            timeout=timeout,
+            conn_timeout=conn_timeout,
+        )
+
+        # 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=expected_value,
+                compress=True,
+                port='port',
+                sock=None,
+                look_for_keys=True,
+            )
+
     def test_openssh_private_key(self):
         # Paramiko behaves differently with OpenSSH generated keys to paramiko
         # generated keys, so we need a test one.
diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index 4d93a100745c4..5771843e3eca8 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -31,6 +31,7 @@
 
 TEST_DAG_ID = 'unit_tests_ssh_test_op'
 TEST_CONN_ID = "conn_id_for_testing"
+DEFAULT_TIMEOUT = 10
 CONN_TIMEOUT = 5
 CMD_TIMEOUT = 7
 TIMEOUT = 12
@@ -243,6 +244,24 @@ def test_arg_checking(self):
         assert task_5.cmd_timeout == TIMEOUT
         assert task_5.ssh_hook.conn_timeout == TIMEOUT
 
+        # if no timeouts are provided (conn_timeout, cmd_timeout, timeout), use the hook default
+        # for the conn_timeout and the operator default for the cmd_timeout
+        task_6 = SSHOperator(
+            task_id="task_6",
+            ssh_conn_id=TEST_CONN_ID,
+            command=COMMAND,
+            conn_timeout=None,
+            timeout=None,
+            cmd_timeout=None,
+            dag=self.dag,
+        )
+        try:
+            task_6.execute(None)
+        except Exception:
+            pass
+        assert task_6.cmd_timeout == DEFAULT_TIMEOUT
+        assert task_6.ssh_hook.conn_timeout == DEFAULT_TIMEOUT
+
     @parameterized.expand(
         [
             (COMMAND, False, False),

From d92bf2c9d5fcd2f49388e533368dd62c438a84d5 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Tue, 24 Aug 2021 07:58:46 -0400
Subject: [PATCH 18/42] Fix spelling mistake in ssh.rst

---
 docs/apache-airflow-providers-ssh/connections/ssh.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/apache-airflow-providers-ssh/connections/ssh.rst b/docs/apache-airflow-providers-ssh/connections/ssh.rst
index bcfb01e128e18..59c8324280eeb 100644
--- a/docs/apache-airflow-providers-ssh/connections/ssh.rst
+++ b/docs/apache-airflow-providers-ssh/connections/ssh.rst
@@ -47,7 +47,7 @@ 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.
-    * ``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.
+    * ``conn_timeout`` - An optional timeout (in seconds) for the TCP connect. Default is ``10``. The parameter to :class:`~airflow.providers.ssh.hooks.ssh.SSHHook` conn_timeout takes precedence.
     * ``timeout`` - Deprecated - use conn_timeout instead.
     * ``compress`` - ``true`` to ask the remote client/server to compress traffic; ``false`` to refuse compression. Default is ``true``.
     * ``no_host_key_check`` - Set to ``false`` to restrict connecting to hosts with no entries in ``~/.ssh/known_hosts`` (Hosts file). This provides maximum protection against trojan horse attacks, but can be troublesome when the ``/etc/ssh/ssh_known_hosts`` file is poorly maintained or connections to new hosts are frequently made. This option forces the user to manually add all new hosts. Default is ``true``, ssh will automatically add new host keys to the user known hosts files.

From a5ad8c7a25e4991ba380575c4323aa7d5e88e195 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Wed, 1 Sep 2021 17:23:10 -0400
Subject: [PATCH 19/42] Move details for conn_timeout to SSHHook docstring
 instead of ssh.rst

---
 airflow/providers/ssh/hooks/ssh.py                    | 4 +++-
 docs/apache-airflow-providers-ssh/connections/ssh.rst | 2 +-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index 9ea2a86cc4515..eda5990288290 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -58,7 +58,9 @@ class SSHHook(BaseHook):
     :type key_file: str
     :param port: port of remote host to connect (Default is paramiko SSH_PORT)
     :type port: int
-    :param conn_timeout: timeout for the attempt to connect to the remote_host.
+    :param conn_timeout: timeout (in seconds) for the attempt to connect to the remote_host.
+        The default is 10 seconds. If provided, it will replace the `conn_timeout` which was
+        predefined in the connection of `ssh_conn_id`.
     :type conn_timeout: int
     :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host.
         Use conn_timeout instead.
diff --git a/docs/apache-airflow-providers-ssh/connections/ssh.rst b/docs/apache-airflow-providers-ssh/connections/ssh.rst
index 59c8324280eeb..c9c109ae8720b 100644
--- a/docs/apache-airflow-providers-ssh/connections/ssh.rst
+++ b/docs/apache-airflow-providers-ssh/connections/ssh.rst
@@ -47,7 +47,7 @@ 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.
-    * ``conn_timeout`` - An optional timeout (in seconds) for the TCP connect. Default is ``10``. The parameter to :class:`~airflow.providers.ssh.hooks.ssh.SSHHook` conn_timeout takes precedence.
+    * ``conn_timeout`` - An optional timeout (in seconds) for the TCP connect. Default is ``10``.
     * ``timeout`` - Deprecated - use conn_timeout instead.
     * ``compress`` - ``true`` to ask the remote client/server to compress traffic; ``false`` to refuse compression. Default is ``true``.
     * ``no_host_key_check`` - Set to ``false`` to restrict connecting to hosts with no entries in ``~/.ssh/known_hosts`` (Hosts file). This provides maximum protection against trojan horse attacks, but can be troublesome when the ``/etc/ssh/ssh_known_hosts`` file is poorly maintained or connections to new hosts are frequently made. This option forces the user to manually add all new hosts. Default is ``true``, ssh will automatically add new host keys to the user known hosts files.

From 55cbfbc630bb42ac864e6eda1644243c99319a3f Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Thu, 2 Sep 2021 22:50:40 -0400
Subject: [PATCH 20/42] Remove commented out line

---
 tests/providers/ssh/operators/test_ssh.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index 9be7d0b31ff7f..4c7da74caba9e 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -276,7 +276,6 @@ def test_get_pyt_set_correctly(self, command, get_pty_in, get_pty_out):
             task_id="test",
             ssh_hook=self.hook,
             command=command,
-            # timeout=TIMEOUT,
             conn_timeout=TIMEOUT,
             cmd_timeout=TIMEOUT,
             get_pty=get_pty_in,

From 60e74b1c7145d36c7cc5fad9f272fcd75954f581 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Sun, 25 Jul 2021 19:43:36 -0400
Subject: [PATCH 21/42] Issue 16363 - add conn_timeout and cmd_timeout params
 to SSHOperator; add conn_timoue param to SSHHook

---
 airflow/providers/ssh/hooks/ssh.py        |  40 +++++++-
 airflow/providers/ssh/operators/ssh.py    |  32 +++++-
 tests/providers/ssh/hooks/test_ssh.py     | 117 +++++++++++++++++++++-
 tests/providers/ssh/operators/test_ssh.py |  35 +++++--
 4 files changed, 207 insertions(+), 17 deletions(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index a8994bdece1b7..b4f524515e5e3 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -56,8 +56,10 @@ class SSHHook(BaseHook):
     :type key_file: str
     :param port: port of remote host to connect (Default is paramiko SSH_PORT)
     :type port: int
-    :param timeout: timeout for the attempt to connect to the remote_host.
+    :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host.
     :type timeout: int
+    :param conn_timeout: timeout for the attempt to connect to the remote_host.
+    :type conn_timeout: int
     :param keepalive_interval: send a keepalive packet to remote host every
         keepalive_interval seconds
     :type keepalive_interval: int
@@ -101,7 +103,8 @@ def __init__(
         password: Optional[str] = None,
         key_file: Optional[str] = None,
         port: Optional[int] = None,
-        timeout: int = 10,
+        timeout: Optional[int] = 10,
+        conn_timeout: Optional[int] = None,
         keepalive_interval: int = 30,
     ) -> None:
         super().__init__()
@@ -113,6 +116,7 @@ def __init__(
         self.pkey = None
         self.port = port
         self.timeout = timeout
+        self.conn_timeout = conn_timeout
         self.keepalive_interval = keepalive_interval
 
         # Default values, overridable from Connection
@@ -137,6 +141,7 @@ def __init__(
                 self.remote_host = conn.host
             if self.port is None:
                 self.port = conn.port
+
             if conn.extra is not None:
                 extra_options = conn.extra_dejson
                 if "key_file" in extra_options and self.key_file is None:
@@ -148,7 +153,22 @@ def __init__(
                     self.pkey = self._pkey_from_private_key(private_key, passphrase=private_key_passphrase)
 
                 if "timeout" in extra_options:
-                    self.timeout = int(extra_options["timeout"], 10)
+                    warnings.warn(
+                        'Extra option `timeout` is deprecated.'
+                        'Please use `conn_timeout` instead.'
+                        'The old option `timeout` will be removed in a future version.',
+                        DeprecationWarning,
+                        stacklevel=2,
+                    )
+                    self.timeout = int(extra_options.get('timeout', 10))
+
+                if "conn_timeout" in extra_options:
+                    self.conn_timeout = int(extra_options.get('conn_timeout'))
+
+                # TODO Rework this - I don't like that just having conn.extra set will override 
+                # self.conn_timeout = int(extra_options.get('conn_timeout',
+                    # extra_options.get('timeout', 10))
+                #)
 
                 if "compress" in extra_options and str(extra_options["compress"]).lower() == 'false':
                     self.compress = False
@@ -185,6 +205,18 @@ def __init__(
                     self.host_key = key_constructor(data=decoded_host_key)
                     self.no_host_key_check = False
 
+        if self.timeout:
+            warnings.warn(
+                'Parameter `timeout` is deprecated.'
+                'Please use `conn_timeout` instead.'
+                'The old option `timeout` will be removed in a future version.',
+                DeprecationWarning,
+                stacklevel=1,
+            )
+
+        if self.conn_timeout is None:
+            self.conn_timeout = self.timeout
+
         if self.pkey and self.key_file:
             raise AirflowException(
                 "Params key_file and private_key both provided.  Must provide no more than one."
@@ -253,7 +285,7 @@ def get_conn(self) -> paramiko.SSHClient:
         connect_kwargs = dict(
             hostname=self.remote_host,
             username=self.username,
-            timeout=self.timeout,
+            timeout=self.conn_timeout,
             compress=self.compress,
             port=self.port,
             sock=self.host_proxy,
diff --git a/airflow/providers/ssh/operators/ssh.py b/airflow/providers/ssh/operators/ssh.py
index 75c61fcf7caa2..1b9097cecdc33 100644
--- a/airflow/providers/ssh/operators/ssh.py
+++ b/airflow/providers/ssh/operators/ssh.py
@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import warnings
 from base64 import b64encode
 from select import select
 from typing import Optional, Union
@@ -43,8 +44,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 timeout: (deprecated) timeout (in seconds) for executing the command. The default is 10 seconds.
     :type timeout: int
+    :param conn_timeout: timeout (in seconds) for maintaining the connection. The default is 10 seconds.
+    :type conn_timeout: int
+    :param cmd_timeout: timeout (in seconds) for executing the command. The default is 10 seconds.
+    :type cmd_timeout: int
     :param environment: a dict of shell environment variables. Note that the
         server will reject them silently if `AcceptEnv` is not set in SSH config.
     :type environment: dict
@@ -66,7 +71,9 @@ def __init__(
         ssh_conn_id: Optional[str] = None,
         remote_host: Optional[str] = None,
         command: Optional[str] = None,
-        timeout: int = 10,
+        timeout: Optional[int] = 10,
+        conn_timeout: Optional[int] = None,
+        cmd_timeout: Optional[int] = None,
         environment: Optional[dict] = None,
         get_pty: bool = False,
         **kwargs,
@@ -77,9 +84,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.cmd_timeout is None:
+            self.cmd_timeout = self.timeout
         self.environment = environment
         self.get_pty = (self.command.startswith('sudo') or get_pty) if self.command else get_pty
 
+        if self.timeout:
+            warnings.warn(
+                'Parameter `timeout` is deprecated.'
+                'Please use `conn_timeout` and `cmd_timeout` instead.'
+                'The old option `timeout` will be removed in a future version.',
+                DeprecationWarning,
+                stacklevel=1,
+            )
+
     def execute(self, context) -> Union[bytes, str, bool]:
         try:
             if self.ssh_conn_id:
@@ -89,7 +111,7 @@ def execute(self, context) -> Union[bytes, str, bool]:
                     self.log.info(
                         "ssh_hook is not provided or invalid. Trying ssh_conn_id to create SSHHook."
                     )
-                    self.ssh_hook = SSHHook(ssh_conn_id=self.ssh_conn_id, timeout=self.timeout)
+                    self.ssh_hook = SSHHook(ssh_conn_id=self.ssh_conn_id, conn_timeout=self.conn_timeout)
 
             if not self.ssh_hook:
                 raise AirflowException("Cannot operate without ssh_hook or ssh_conn_id.")
@@ -112,7 +134,7 @@ def execute(self, context) -> Union[bytes, str, bool]:
                 stdin, stdout, stderr = ssh_client.exec_command(
                     command=self.command,
                     get_pty=self.get_pty,
-                    timeout=self.timeout,
+                    timeout=self.cmd_timeout,
                     environment=self.environment,
                 )
                 # get channels
@@ -133,7 +155,7 @@ def execute(self, context) -> Union[bytes, str, bool]:
 
                 # read from both stdout and stderr
                 while not channel.closed or channel.recv_ready() or channel.recv_stderr_ready():
-                    readq, _, _ = select([channel], [], [], self.timeout)
+                    readq, _, _ = select([channel], [], [], self.cmd_timeout)
                     for recv in readq:
                         if recv.recv_ready():
                             line = stdout.channel.recv(len(recv.in_buffer))
diff --git a/tests/providers/ssh/hooks/test_ssh.py b/tests/providers/ssh/hooks/test_ssh.py
index 70ba4eb0ce642..18389be798ba7 100644
--- a/tests/providers/ssh/hooks/test_ssh.py
+++ b/tests/providers/ssh/hooks/test_ssh.py
@@ -68,6 +68,9 @@ def generate_host_key(pkey: paramiko.PKey):
 TEST_PKEY_ECDSA = paramiko.ECDSAKey.generate()
 TEST_PRIVATE_KEY_ECDSA = generate_key_string(pkey=TEST_PKEY_ECDSA)
 
+TEST_TIMEOUT = 20
+TEST_CONN_TIMEOUT = 30
+
 PASSPHRASE = ''.join(random.choice(string.ascii_letters) for i in range(10))
 TEST_ENCRYPTED_PRIVATE_KEY = generate_key_string(pkey=TEST_PKEY, passphrase=PASSPHRASE)
 
@@ -76,6 +79,8 @@ class TestSSHHook(unittest.TestCase):
     CONN_SSH_WITH_PRIVATE_KEY_EXTRA = 'ssh_with_private_key_extra'
     CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA = 'ssh_with_private_key_ecdsa_extra'
     CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA = 'ssh_with_private_key_passphrase_extra'
+    CONN_SSH_WITH_TIMEOUT_EXTRA = 'ssh_with_timeout_extra'
+    CONN_SSH_WITH_CONN_TIMEOUT_EXTRA = 'ssh_with_conn_timeout_extra'
     CONN_SSH_WITH_EXTRA = 'ssh_with_extra'
     CONN_SSH_WITH_EXTRA_FALSE_LOOK_FOR_KEYS = 'ssh_with_extra_false_look_for_keys'
     CONN_SSH_WITH_HOST_KEY_EXTRA = 'ssh_with_host_key_extra'
@@ -91,6 +96,8 @@ def tearDownClass(cls) -> None:
                 cls.CONN_SSH_WITH_PRIVATE_KEY_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA,
+                cls.CONN_SSH_WITH_TIMEOUT_EXTRA,
+                cls.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA,
                 cls.CONN_SSH_WITH_EXTRA,
                 cls.CONN_SSH_WITH_HOST_KEY_EXTRA,
                 cls.CONN_SSH_WITH_HOST_KEY_EXTRA_WITH_TYPE,
@@ -128,7 +135,7 @@ def setUpClass(cls) -> None:
                 conn_type='ssh',
                 extra=json.dumps(
                     {
-                        "private_key": TEST_PRIVATE_KEY,
+                        "private_key": TEST_PRIVATE_KEY
                     }
                 ),
             )
@@ -151,6 +158,22 @@ def setUpClass(cls) -> None:
                 extra=json.dumps({"private_key": TEST_PRIVATE_KEY_ECDSA}),
             )
         )
+        db.merge_conn(
+            Connection(
+                conn_id=cls.CONN_SSH_WITH_TIMEOUT_EXTRA,
+                host='localhost',
+                conn_type='ssh',
+                extra=json.dumps({"timeout": TEST_TIMEOUT}),
+            )
+        )
+        db.merge_conn(
+            Connection(
+                conn_id=cls.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA,
+                host='localhost',
+                conn_type='ssh',
+                extra=json.dumps({"conn_timeout": TEST_CONN_TIMEOUT}),
+            )
+        )
         db.merge_conn(
             Connection(
                 conn_id=cls.CONN_SSH_WITH_HOST_KEY_EXTRA,
@@ -475,6 +498,98 @@ 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
+        )
+
+        with hook.get_conn():
+            ssh_mock.return_value.connect.assert_called_once_with(
+                hostname='remote_host',
+                username='username',
+                timeout=30,
+                compress=True,
+                port='port',
+                sock=None,
+                look_for_keys=True,
+            )
+
     def test_openssh_private_key(self):
         # Paramiko behaves differently with OpenSSH generated keys to paramiko
         # generated keys, so we need a test one.
diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index 7083354301b70..3fb3d6fb4ca22 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -52,7 +52,7 @@ def setUp(self):
         self.hook = hook
         self.dag = dag
 
-    def test_hook_created_correctly(self):
+    def test_hook_created_correctly_with_timeout(self):
         timeout = 20
         ssh_id = "ssh_default"
         task = SSHOperator(
@@ -61,8 +61,21 @@ def test_hook_created_correctly(self):
         assert task is not None
 
         task.execute(None)
+        
+        assert timeout == task.ssh_hook.conn_timeout
+        assert ssh_id == task.ssh_hook.ssh_conn_id
 
-        assert timeout == task.ssh_hook.timeout
+    def test_hook_created_correctly(self):
+        conn_timeout = 20
+        ssh_id = "ssh_default"
+        task = SSHOperator(
+            task_id="test", command=COMMAND, dag=self.dag, conn_timeout=conn_timeout, ssh_conn_id="ssh_default"
+        )
+        assert task is not None
+
+        task.execute(None)
+        
+        assert conn_timeout == task.ssh_hook.conn_timeout
         assert ssh_id == task.ssh_hook.ssh_conn_id
 
     @conf_vars({('core', 'enable_xcom_pickling'): 'False'})
@@ -138,7 +151,8 @@ def test_no_output_command(self):
     def test_arg_checking(self):
         # Exception should be raised if neither ssh_hook nor ssh_conn_id is provided
         with pytest.raises(AirflowException, match="Cannot operate without ssh_hook or ssh_conn_id."):
-            task_0 = SSHOperator(task_id="test", command=COMMAND, timeout=TIMEOUT, dag=self.dag)
+            #  task_0 = SSHOperator(task_id="test", command=COMMAND, timeout=TIMEOUT, dag=self.dag)
+            task_0 = SSHOperator(task_id="test", command=COMMAND, conn_timeout=TIMEOUT, cmd_timeout=TIMEOUT, dag=self.dag)
             task_0.execute(None)
 
         # if ssh_hook is invalid/not provided, use ssh_conn_id to create SSHHook
@@ -147,7 +161,8 @@ def test_arg_checking(self):
             ssh_hook="string_rather_than_SSHHook",  # invalid ssh_hook
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -160,7 +175,9 @@ def test_arg_checking(self):
             task_id="test_2",
             ssh_conn_id=TEST_CONN_ID,  # no ssh_hook provided
             command=COMMAND,
-            timeout=TIMEOUT,
+            # timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -175,7 +192,9 @@ def test_arg_checking(self):
             ssh_hook=self.hook,
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            timeout=TIMEOUT,
+            # timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -198,7 +217,9 @@ def test_get_pyt_set_correctly(self, command, get_pty_in, get_pty_out):
             task_id="test",
             ssh_hook=self.hook,
             command=command,
-            timeout=TIMEOUT,
+            # timeout=TIMEOUT,
+            conn_timeout=TIMEOUT,
+            cmd_timeout=TIMEOUT,
             get_pty=get_pty_in,
             dag=self.dag,
         )

From b3f73e19dd64a45298d1b1e3816bace685a2fb49 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Sun, 25 Jul 2021 20:24:56 -0400
Subject: [PATCH 22/42] Update SSHOperator tests for new conn_timeout and
 cmd_timeout params

---
 tests/providers/ssh/operators/test_ssh.py | 55 ++++++++++++++++++-----
 1 file changed, 44 insertions(+), 11 deletions(-)

diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index 3fb3d6fb4ca22..437058d699bbc 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -31,7 +31,9 @@
 
 TEST_DAG_ID = 'unit_tests_ssh_test_op'
 TEST_CONN_ID = "conn_id_for_testing"
-TIMEOUT = 5
+CONN_TIMEOUT = 5
+CMD_TIMEOUT = 7
+TIMEOUT = 12
 DEFAULT_DATE = datetime(2017, 1, 1)
 COMMAND = "echo -n airflow"
 COMMAND_WITH_SUDO = "sudo " + COMMAND
@@ -67,9 +69,10 @@ def test_hook_created_correctly_with_timeout(self):
 
     def test_hook_created_correctly(self):
         conn_timeout = 20
+        cmd_timeout = 45
         ssh_id = "ssh_default"
         task = SSHOperator(
-            task_id="test", command=COMMAND, dag=self.dag, conn_timeout=conn_timeout, ssh_conn_id="ssh_default"
+            task_id="test", command=COMMAND, dag=self.dag, conn_timeout=conn_timeout, cmd_timeout=cmd_timeout, ssh_conn_id="ssh_default"
         )
         assert task is not None
 
@@ -151,7 +154,6 @@ def test_no_output_command(self):
     def test_arg_checking(self):
         # Exception should be raised if neither ssh_hook nor ssh_conn_id is provided
         with pytest.raises(AirflowException, match="Cannot operate without ssh_hook or ssh_conn_id."):
-            #  task_0 = SSHOperator(task_id="test", command=COMMAND, timeout=TIMEOUT, dag=self.dag)
             task_0 = SSHOperator(task_id="test", command=COMMAND, conn_timeout=TIMEOUT, cmd_timeout=TIMEOUT, dag=self.dag)
             task_0.execute(None)
 
@@ -161,8 +163,8 @@ def test_arg_checking(self):
             ssh_hook="string_rather_than_SSHHook",  # invalid ssh_hook
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            conn_timeout=TIMEOUT,
-            cmd_timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -175,9 +177,8 @@ def test_arg_checking(self):
             task_id="test_2",
             ssh_conn_id=TEST_CONN_ID,  # no ssh_hook provided
             command=COMMAND,
-            # timeout=TIMEOUT,
-            conn_timeout=TIMEOUT,
-            cmd_timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -192,9 +193,8 @@ def test_arg_checking(self):
             ssh_hook=self.hook,
             ssh_conn_id=TEST_CONN_ID,
             command=COMMAND,
-            # timeout=TIMEOUT,
-            conn_timeout=TIMEOUT,
-            cmd_timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
             dag=self.dag,
         )
         try:
@@ -203,6 +203,39 @@ def test_arg_checking(self):
             pass
         assert task_3.ssh_hook.ssh_conn_id == self.hook.ssh_conn_id
 
+        # if valid conn_timeout, cmd_timeout, and timeout are provided, ignore timeout
+        task_4 = SSHOperator(
+            task_id="test_4",
+            ssh_conn_id=TEST_CONN_ID,
+            command=COMMAND,
+            timeout=TIMEOUT,
+            conn_timeout=CONN_TIMEOUT,
+            cmd_timeout=CMD_TIMEOUT,
+            dag=self.dag,
+        )
+        try:
+            task_4.execute(None)
+        except Exception:
+            pass
+        assert task_4.conn_timeout == CONN_TIMEOUT
+        assert task_4.cmd_timeout == CMD_TIMEOUT
+        assert task_4.ssh_hook.conn_timeout == CONN_TIMEOUT
+
+        # if valid only timeout is provided, use for both conn_timeout and cmd_timeout
+        task_5 = SSHOperator(
+            task_id="test_5",
+            ssh_conn_id=TEST_CONN_ID,
+            command=COMMAND,
+            timeout=TIMEOUT,
+            dag=self.dag,
+        )
+        try:
+            task_5.execute(None)
+        except Exception:
+            pass
+        assert task_5.cmd_timeout == TIMEOUT
+        assert task_5.ssh_hook.conn_timeout == TIMEOUT
+
     @parameterized.expand(
         [
             (COMMAND, False, False),

From 0c1c1efebe11c562eca6e71b6fc12fd0f97486b8 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Mon, 26 Jul 2021 13:19:47 -0400
Subject: [PATCH 23/42] Removed commented TODO for SSHHook changes

---
 airflow/providers/ssh/hooks/ssh.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index b4f524515e5e3..5f9c637b9d18b 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -165,11 +165,6 @@ def __init__(
                 if "conn_timeout" in extra_options:
                     self.conn_timeout = int(extra_options.get('conn_timeout'))
 
-                # TODO Rework this - I don't like that just having conn.extra set will override 
-                # self.conn_timeout = int(extra_options.get('conn_timeout',
-                    # extra_options.get('timeout', 10))
-                #)
-
                 if "compress" in extra_options and str(extra_options["compress"]).lower() == 'false':
                     self.compress = False
 

From e51910f9524a11308d7228b513cd0bca8f89b3d2 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Mon, 26 Jul 2021 16:37:29 -0400
Subject: [PATCH 24/42] Fix static checks for operators/test_ssh.py and
 hooks/test_ssh.py

---
 tests/providers/ssh/hooks/test_ssh.py     |  8 ++------
 tests/providers/ssh/operators/test_ssh.py | 15 +++++++++++----
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/tests/providers/ssh/hooks/test_ssh.py b/tests/providers/ssh/hooks/test_ssh.py
index 18389be798ba7..4afa6fb12e060 100644
--- a/tests/providers/ssh/hooks/test_ssh.py
+++ b/tests/providers/ssh/hooks/test_ssh.py
@@ -133,11 +133,7 @@ def setUpClass(cls) -> None:
                 conn_id=cls.CONN_SSH_WITH_PRIVATE_KEY_EXTRA,
                 host='localhost',
                 conn_type='ssh',
-                extra=json.dumps(
-                    {
-                        "private_key": TEST_PRIVATE_KEY
-                    }
-                ),
+                extra=json.dumps({"private_key": TEST_PRIVATE_KEY}),
             )
         )
         db.merge_conn(
@@ -576,7 +572,7 @@ def test_ssh_connection_with_conn_timeout_extra(self, ssh_mock):
             port='port',
             username='username',
             timeout=10,
-            conn_timeout=15
+            conn_timeout=15,
         )
 
         with hook.get_conn():
diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index 437058d699bbc..fb7d819b057d6 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -63,7 +63,7 @@ def test_hook_created_correctly_with_timeout(self):
         assert task is not None
 
         task.execute(None)
-        
+
         assert timeout == task.ssh_hook.conn_timeout
         assert ssh_id == task.ssh_hook.ssh_conn_id
 
@@ -72,12 +72,17 @@ def test_hook_created_correctly(self):
         cmd_timeout = 45
         ssh_id = "ssh_default"
         task = SSHOperator(
-            task_id="test", command=COMMAND, dag=self.dag, conn_timeout=conn_timeout, cmd_timeout=cmd_timeout, ssh_conn_id="ssh_default"
+            task_id="test",
+            command=COMMAND,
+            dag=self.dag,
+            conn_timeout=conn_timeout,
+            cmd_timeout=cmd_timeout,
+            ssh_conn_id="ssh_default",
         )
         assert task is not None
 
         task.execute(None)
-        
+
         assert conn_timeout == task.ssh_hook.conn_timeout
         assert ssh_id == task.ssh_hook.ssh_conn_id
 
@@ -154,7 +159,9 @@ def test_no_output_command(self):
     def test_arg_checking(self):
         # Exception should be raised if neither ssh_hook nor ssh_conn_id is provided
         with pytest.raises(AirflowException, match="Cannot operate without ssh_hook or ssh_conn_id."):
-            task_0 = SSHOperator(task_id="test", command=COMMAND, conn_timeout=TIMEOUT, cmd_timeout=TIMEOUT, dag=self.dag)
+            task_0 = SSHOperator(
+                task_id="test", command=COMMAND, conn_timeout=TIMEOUT, cmd_timeout=TIMEOUT, dag=self.dag
+            )
             task_0.execute(None)
 
         # if ssh_hook is invalid/not provided, use ssh_conn_id to create SSHHook

From 05b62d90448e09d80aa83b946efc5f51290fcbc7 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Wed, 28 Jul 2021 20:10:01 -0400
Subject: [PATCH 25/42] Set deprecated timeout param default value to None; add
 more tests

---
 airflow/providers/ssh/hooks/ssh.py        | 10 ++++---
 airflow/providers/ssh/operators/ssh.py    | 12 +++++----
 tests/providers/ssh/hooks/test_ssh.py     | 32 +++++++++++++++++++++++
 tests/providers/ssh/operators/test_ssh.py |  2 +-
 4 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index 5f9c637b9d18b..3ddae25df8f9c 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -34,6 +34,8 @@
 except ImportError:
     from getpass import getuser
 
+TIMEOUT_DEFAULT = 10
+
 
 class SSHHook(BaseHook):
     """
@@ -56,10 +58,10 @@ class SSHHook(BaseHook):
     :type key_file: str
     :param port: port of remote host to connect (Default is paramiko SSH_PORT)
     :type port: int
-    :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host.
-    :type timeout: int
     :param conn_timeout: timeout for the attempt to connect to the remote_host.
     :type conn_timeout: int
+    :param timeout: (Deprecated). 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
@@ -103,7 +105,7 @@ def __init__(
         password: Optional[str] = None,
         key_file: Optional[str] = None,
         port: Optional[int] = None,
-        timeout: Optional[int] = 10,
+        timeout: Optional[int] = None,
         conn_timeout: Optional[int] = None,
         keepalive_interval: int = 30,
     ) -> None:
@@ -210,7 +212,7 @@ def __init__(
             )
 
         if self.conn_timeout is None:
-            self.conn_timeout = self.timeout
+            self.conn_timeout = self.timeout if self.timeout else TIMEOUT_DEFAULT
 
         if self.pkey and self.key_file:
             raise AirflowException(
diff --git a/airflow/providers/ssh/operators/ssh.py b/airflow/providers/ssh/operators/ssh.py
index 1b9097cecdc33..61f374dbb3515 100644
--- a/airflow/providers/ssh/operators/ssh.py
+++ b/airflow/providers/ssh/operators/ssh.py
@@ -26,6 +26,8 @@
 from airflow.models import BaseOperator
 from airflow.providers.ssh.hooks.ssh import SSHHook
 
+TIMEOUT_DEFAULT = 10
+
 
 class SSHOperator(BaseOperator):
     """
@@ -44,12 +46,12 @@ class SSHOperator(BaseOperator):
     :type remote_host: str
     :param command: command to execute on remote host. (templated)
     :type command: str
-    :param timeout: (deprecated) timeout (in seconds) for executing the command. The default is 10 seconds.
-    :type timeout: int
     :param conn_timeout: timeout (in seconds) for maintaining the connection. The default is 10 seconds.
     :type conn_timeout: int
     :param cmd_timeout: timeout (in seconds) for executing the command. The default is 10 seconds.
     :type cmd_timeout: int
+    :param timeout: (deprecated) timeout (in seconds) for executing the command. The default is 10 seconds.
+    :type timeout: int
     :param environment: a dict of shell environment variables. Note that the
         server will reject them silently if `AcceptEnv` is not set in SSH config.
     :type environment: dict
@@ -71,7 +73,7 @@ def __init__(
         ssh_conn_id: Optional[str] = None,
         remote_host: Optional[str] = None,
         command: Optional[str] = None,
-        timeout: Optional[int] = 10,
+        timeout: Optional[int] = None,
         conn_timeout: Optional[int] = None,
         cmd_timeout: Optional[int] = None,
         environment: Optional[dict] = None,
@@ -87,9 +89,9 @@ def __init__(
         self.conn_timeout = conn_timeout
         self.cmd_timeout = cmd_timeout
         if self.conn_timeout is None:
-            self.conn_timeout = self.timeout
+            self.conn_timeout = self.timeout if self.timeout else TIMEOUT_DEFAULT
         if self.cmd_timeout is None:
-            self.cmd_timeout = self.timeout
+            self.cmd_timeout = self.timeout if self.timeout else TIMEOUT_DEFAULT
         self.environment = environment
         self.get_pty = (self.command.startswith('sudo') or get_pty) if self.command else get_pty
 
diff --git a/tests/providers/ssh/hooks/test_ssh.py b/tests/providers/ssh/hooks/test_ssh.py
index 4afa6fb12e060..139c9d0c18435 100644
--- a/tests/providers/ssh/hooks/test_ssh.py
+++ b/tests/providers/ssh/hooks/test_ssh.py
@@ -81,6 +81,7 @@ class TestSSHHook(unittest.TestCase):
     CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA = 'ssh_with_private_key_passphrase_extra'
     CONN_SSH_WITH_TIMEOUT_EXTRA = 'ssh_with_timeout_extra'
     CONN_SSH_WITH_CONN_TIMEOUT_EXTRA = 'ssh_with_conn_timeout_extra'
+    CONN_SSH_WITH_TIMEOUT_AND_CONN_TIMEOUT_EXTRA = 'ssh_with_timeout_and_conn_timeout_extra'
     CONN_SSH_WITH_EXTRA = 'ssh_with_extra'
     CONN_SSH_WITH_EXTRA_FALSE_LOOK_FOR_KEYS = 'ssh_with_extra_false_look_for_keys'
     CONN_SSH_WITH_HOST_KEY_EXTRA = 'ssh_with_host_key_extra'
@@ -98,6 +99,7 @@ def tearDownClass(cls) -> None:
                 cls.CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA,
                 cls.CONN_SSH_WITH_TIMEOUT_EXTRA,
                 cls.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA,
+                cls.CONN_SSH_WITH_TIMEOUT_AND_CONN_TIMEOUT_EXTRA,
                 cls.CONN_SSH_WITH_EXTRA,
                 cls.CONN_SSH_WITH_HOST_KEY_EXTRA,
                 cls.CONN_SSH_WITH_HOST_KEY_EXTRA_WITH_TYPE,
@@ -170,6 +172,14 @@ def setUpClass(cls) -> None:
                 extra=json.dumps({"conn_timeout": TEST_CONN_TIMEOUT}),
             )
         )
+        db.merge_conn(
+            Connection(
+                conn_id=cls.CONN_SSH_WITH_TIMEOUT_AND_CONN_TIMEOUT_EXTRA,
+                host='localhost',
+                conn_type='ssh',
+                extra=json.dumps({"conn_timeout": TEST_CONN_TIMEOUT, 'timeout': TEST_TIMEOUT}),
+            )
+        )
         db.merge_conn(
             Connection(
                 conn_id=cls.CONN_SSH_WITH_HOST_KEY_EXTRA,
@@ -586,6 +596,28 @@ def test_ssh_connection_with_conn_timeout_extra(self, ssh_mock):
                 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,
+        )
+
+        with hook.get_conn():
+            ssh_mock.return_value.connect.assert_called_once_with(
+                hostname='remote_host',
+                username='username',
+                timeout=30,
+                compress=True,
+                port='port',
+                sock=None,
+                look_for_keys=True,
+            )
+
     def test_openssh_private_key(self):
         # Paramiko behaves differently with OpenSSH generated keys to paramiko
         # generated keys, so we need a test one.
diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index fb7d819b057d6..924ded1c02a9a 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -228,7 +228,7 @@ def test_arg_checking(self):
         assert task_4.cmd_timeout == CMD_TIMEOUT
         assert task_4.ssh_hook.conn_timeout == CONN_TIMEOUT
 
-        # if valid only timeout is provided, use for both conn_timeout and cmd_timeout
+        # if only timeout is provided, use for both conn_timeout and cmd_timeout
         task_5 = SSHOperator(
             task_id="test_5",
             ssh_conn_id=TEST_CONN_ID,

From 83472326abe9acf668793891887531c17f480f3a Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Wed, 28 Jul 2021 21:20:53 -0400
Subject: [PATCH 26/42] Remove unnecessary extra_options.get() calls for
 timeout and conn_timeout

---
 airflow/providers/ssh/hooks/ssh.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index 3ddae25df8f9c..e399fd8baaf75 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -162,10 +162,10 @@ def __init__(
                         DeprecationWarning,
                         stacklevel=2,
                     )
-                    self.timeout = int(extra_options.get('timeout', 10))
+                    self.timeout = int(extra_options['timeout'])
 
                 if "conn_timeout" in extra_options:
-                    self.conn_timeout = int(extra_options.get('conn_timeout'))
+                    self.conn_timeout = int(extra_options['conn_timeout'])
 
                 if "compress" in extra_options and str(extra_options["compress"]).lower() == 'false':
                     self.compress = False

From cb5c30117598102927041c963d95cb4780410a63 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Sun, 1 Aug 2021 18:14:21 -0400
Subject: [PATCH 27/42] Add reference to new timeout parameters in docstrings
 for deprecated timeout

---
 airflow/providers/ssh/hooks/ssh.py     | 2 +-
 airflow/providers/ssh/operators/ssh.py | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index e399fd8baaf75..024a66a46f170 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -60,7 +60,7 @@ class SSHHook(BaseHook):
     :type port: int
     :param conn_timeout: timeout for the attempt to connect to the remote_host.
     :type conn_timeout: int
-    :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host.
+    :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host. Use conn_timeout instead.
     :type timeout: int
     :param keepalive_interval: send a keepalive packet to remote host every
         keepalive_interval seconds
diff --git a/airflow/providers/ssh/operators/ssh.py b/airflow/providers/ssh/operators/ssh.py
index 61f374dbb3515..b0dc19eb20277 100644
--- a/airflow/providers/ssh/operators/ssh.py
+++ b/airflow/providers/ssh/operators/ssh.py
@@ -51,6 +51,7 @@ class SSHOperator(BaseOperator):
     :param cmd_timeout: timeout (in seconds) for executing the command. The default is 10 seconds.
     :type cmd_timeout: int
     :param timeout: (deprecated) timeout (in seconds) for executing the command. The default is 10 seconds.
+        Use conn_timeout and cmd_timeout parameters instead.
     :type timeout: int
     :param environment: a dict of shell environment variables. Note that the
         server will reject them silently if `AcceptEnv` is not set in SSH config.

From df5a5cda24ce2316d7b067574edff4b6d1fc3755 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Sun, 1 Aug 2021 20:21:21 -0400
Subject: [PATCH 28/42] Move timeout docstring text to 2nd line to keep line
 within 110 characters

---
 airflow/providers/ssh/hooks/ssh.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index 024a66a46f170..332de6eff59cf 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -60,7 +60,8 @@ class SSHHook(BaseHook):
     :type port: int
     :param conn_timeout: timeout for the attempt to connect to the remote_host.
     :type conn_timeout: int
-    :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host. Use conn_timeout instead.
+    :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host.
+        Use conn_timeout instead.
     :type timeout: int
     :param keepalive_interval: send a keepalive packet to remote host every
         keepalive_interval seconds

From b7d0115a9556dafdc6da833daf851961028d7877 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Sun, 1 Aug 2021 20:40:02 -0400
Subject: [PATCH 29/42] Favor conn_timeout parameter over extra options

---
 airflow/providers/ssh/hooks/ssh.py    | 2 +-
 tests/providers/ssh/hooks/test_ssh.py | 6 ++++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index 332de6eff59cf..9ea2a86cc4515 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -165,7 +165,7 @@ def __init__(
                     )
                     self.timeout = int(extra_options['timeout'])
 
-                if "conn_timeout" in extra_options:
+                if "conn_timeout" in extra_options and self.conn_timeout is None:
                     self.conn_timeout = int(extra_options['conn_timeout'])
 
                 if "compress" in extra_options and str(extra_options["compress"]).lower() == 'false':
diff --git a/tests/providers/ssh/hooks/test_ssh.py b/tests/providers/ssh/hooks/test_ssh.py
index 139c9d0c18435..7f3c9c2167291 100644
--- a/tests/providers/ssh/hooks/test_ssh.py
+++ b/tests/providers/ssh/hooks/test_ssh.py
@@ -585,11 +585,12 @@ def test_ssh_connection_with_conn_timeout_extra(self, ssh_mock):
             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=30,
+                timeout=15,
                 compress=True,
                 port='port',
                 sock=None,
@@ -607,11 +608,12 @@ def test_ssh_connection_with_timeout_extra_and_conn_timeout_extra(self, ssh_mock
             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=30,
+                timeout=15,
                 compress=True,
                 port='port',
                 sock=None,

From 5c07819b973b6c3e2ff185db76d3e877c3443b24 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Wed, 4 Aug 2021 22:03:13 -0400
Subject: [PATCH 30/42] Add conn_timeout to SSH Connection docs and mark
 timeout as deprecated

---
 docs/apache-airflow-providers-ssh/connections/ssh.rst | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/docs/apache-airflow-providers-ssh/connections/ssh.rst b/docs/apache-airflow-providers-ssh/connections/ssh.rst
index ec9e6d165539e..bcfb01e128e18 100644
--- a/docs/apache-airflow-providers-ssh/connections/ssh.rst
+++ b/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.
+    * ``timeout`` - Deprecated - use conn_timeout instead.
     * ``compress`` - ``true`` to ask the remote client/server to compress traffic; ``false`` to refuse compression. Default is ``true``.
     * ``no_host_key_check`` - Set to ``false`` to restrict connecting to hosts with no entries in ``~/.ssh/known_hosts`` (Hosts file). This provides maximum protection against trojan horse attacks, but can be troublesome when the ``/etc/ssh/ssh_known_hosts`` file is poorly maintained or connections to new hosts are frequently made. This option forces the user to manually add all new hosts. Default is ``true``, ssh will automatically add new host keys to the user known hosts files.
     * ``allow_host_key_change`` - Set to ``true`` if you want to allow connecting to hosts that has host key changed or when you get 'REMOTE HOST IDENTIFICATION HAS CHANGED' error.  This wont protect against Man-In-The-Middle attacks. Other possible solution is to remove the host entry from ``~/.ssh/known_hosts`` file. Default is ``false``.
@@ -60,7 +61,7 @@ Extra (optional)
 
        {
           "key_file": "/home/airflow/.ssh/id_rsa",
-          "timeout": "10",
+          "conn_timeout": "10",
           "compress": "false",
           "look_for_keys": "false",
           "allow_host_key_change": "false",
@@ -75,7 +76,7 @@ Extra (optional)
 
     .. code-block:: bash
 
-        export AIRFLOW_CONN_MAIN_SERVER='ssh://user:pass@localhost:22?timeout=10&compress=false&no_host_key_check=false&allow_host_key_change=true&key_file=%2Fhome%2Fairflow%2F.ssh%2Fid_rsa'
+        export AIRFLOW_CONN_MAIN_SERVER='ssh://user:pass@localhost:22?conn_timeout=10&compress=false&no_host_key_check=false&allow_host_key_change=true&key_file=%2Fhome%2Fairflow%2F.ssh%2Fid_rsa'
 
     Example connection string with ``private_key`` (actual private key provided in connection):
 

From d287e4d0c32696613ed460534fa36bb081404d8a Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Mon, 23 Aug 2021 23:09:19 -0400
Subject: [PATCH 31/42] Fix issue with defaults for timeouts in
 SSHHook/SSHOperator; add parametrized unit tests for SSHHook

---
 airflow/providers/ssh/operators/ssh.py    | 10 ++--
 tests/providers/ssh/hooks/test_ssh.py     | 66 +++++++++++++++++++++++
 tests/providers/ssh/operators/test_ssh.py | 19 +++++++
 3 files changed, 91 insertions(+), 4 deletions(-)

diff --git a/airflow/providers/ssh/operators/ssh.py b/airflow/providers/ssh/operators/ssh.py
index b0dc19eb20277..7f97b031b414d 100644
--- a/airflow/providers/ssh/operators/ssh.py
+++ b/airflow/providers/ssh/operators/ssh.py
@@ -26,7 +26,7 @@
 from airflow.models import BaseOperator
 from airflow.providers.ssh.hooks.ssh import SSHHook
 
-TIMEOUT_DEFAULT = 10
+CMD_TIMEOUT = 10
 
 
 class SSHOperator(BaseOperator):
@@ -47,6 +47,8 @@ class SSHOperator(BaseOperator):
     :param command: command to execute on remote host. (templated)
     :type command: str
     :param conn_timeout: timeout (in seconds) for maintaining the connection. The default is 10 seconds.
+        Nullable. If provided, it will replace the `conn_timeout` which was
+        predefined in the connection of `ssh_conn_id`.
     :type conn_timeout: int
     :param cmd_timeout: timeout (in seconds) for executing the command. The default is 10 seconds.
     :type cmd_timeout: int
@@ -89,10 +91,10 @@ def __init__(
         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
+        if self.conn_timeout is None and self.timeout:
+            self.conn_timeout = self.timeout
         if self.cmd_timeout is None:
-            self.cmd_timeout = self.timeout if self.timeout else TIMEOUT_DEFAULT
+            self.cmd_timeout = self.timeout if self.timeout else CMD_TIMEOUT
         self.environment = environment
         self.get_pty = (self.command.startswith('sudo') or get_pty) if self.command else get_pty
 
diff --git a/tests/providers/ssh/hooks/test_ssh.py b/tests/providers/ssh/hooks/test_ssh.py
index 7f3c9c2167291..be63093ccfbff 100644
--- a/tests/providers/ssh/hooks/test_ssh.py
+++ b/tests/providers/ssh/hooks/test_ssh.py
@@ -25,6 +25,7 @@
 from unittest import mock
 
 import paramiko
+from parameterized import parameterized
 
 from airflow import settings
 from airflow.models import Connection
@@ -76,6 +77,7 @@ def generate_host_key(pkey: paramiko.PKey):
 
 
 class TestSSHHook(unittest.TestCase):
+    CONN_SSH_WITH_NO_EXTRA = 'ssh_with_no_extra'
     CONN_SSH_WITH_PRIVATE_KEY_EXTRA = 'ssh_with_private_key_extra'
     CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA = 'ssh_with_private_key_ecdsa_extra'
     CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA = 'ssh_with_private_key_passphrase_extra'
@@ -94,6 +96,7 @@ class TestSSHHook(unittest.TestCase):
     def tearDownClass(cls) -> None:
         with create_session() as session:
             conns_to_reset = [
+                cls.CONN_SSH_WITH_NO_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA,
@@ -113,6 +116,14 @@ def tearDownClass(cls) -> None:
 
     @classmethod
     def setUpClass(cls) -> None:
+        db.merge_conn(
+            Connection(
+                conn_id=cls.CONN_SSH_WITH_NO_EXTRA,
+                host='localhost',
+                conn_type='ssh',
+                extra=None,
+            )
+        )
         db.merge_conn(
             Connection(
                 conn_id=cls.CONN_SSH_WITH_EXTRA,
@@ -620,6 +631,61 @@ def test_ssh_connection_with_timeout_extra_and_conn_timeout_extra(self, ssh_mock
                 look_for_keys=True,
             )
 
+    @parameterized.expand(
+        [
+            (TEST_TIMEOUT, TEST_CONN_TIMEOUT, True, True, TEST_CONN_TIMEOUT),
+            (TEST_TIMEOUT, TEST_CONN_TIMEOUT, True, False, TEST_CONN_TIMEOUT),
+            (TEST_TIMEOUT, TEST_CONN_TIMEOUT, False, True, TEST_CONN_TIMEOUT),
+            (TEST_TIMEOUT, TEST_CONN_TIMEOUT, False, False, TEST_CONN_TIMEOUT),
+            (TEST_TIMEOUT, None, True, True, TEST_CONN_TIMEOUT),
+            (TEST_TIMEOUT, None, True, False, TEST_TIMEOUT),
+            (TEST_TIMEOUT, None, False, True, TEST_CONN_TIMEOUT),
+            (TEST_TIMEOUT, None, False, False, TEST_TIMEOUT),
+            (None, TEST_CONN_TIMEOUT, True, True, TEST_CONN_TIMEOUT),
+            (None, TEST_CONN_TIMEOUT, True, False, TEST_CONN_TIMEOUT),
+            (None, TEST_CONN_TIMEOUT, False, True, TEST_CONN_TIMEOUT),
+            (None, TEST_CONN_TIMEOUT, False, False, TEST_CONN_TIMEOUT),
+            (None, None, True, True, TEST_CONN_TIMEOUT),
+            (None, None, True, False, TEST_TIMEOUT),
+            (None, None, False, True, TEST_CONN_TIMEOUT),
+            (None, None, False, False, 10),
+        ]
+    )
+    @mock.patch('airflow.providers.ssh.hooks.ssh.paramiko.SSHClient')
+    def test_ssh_connection_with_all_timeout_param_and_extra_combinations(
+        self, timeout, conn_timeout, timeoutextra, conn_timeoutextra, expected_value, ssh_mock
+    ):
+
+        if timeoutextra and conn_timeoutextra:
+            ssh_conn_id = self.CONN_SSH_WITH_TIMEOUT_AND_CONN_TIMEOUT_EXTRA
+        elif timeoutextra and not conn_timeoutextra:
+            ssh_conn_id = self.CONN_SSH_WITH_TIMEOUT_EXTRA
+        elif not timeoutextra and conn_timeoutextra:
+            ssh_conn_id = self.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA
+        else:
+            ssh_conn_id = self.CONN_SSH_WITH_NO_EXTRA
+
+        hook = SSHHook(
+            ssh_conn_id=ssh_conn_id,
+            remote_host='remote_host',
+            port='port',
+            username='username',
+            timeout=timeout,
+            conn_timeout=conn_timeout,
+        )
+
+        # 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=expected_value,
+                compress=True,
+                port='port',
+                sock=None,
+                look_for_keys=True,
+            )
+
     def test_openssh_private_key(self):
         # Paramiko behaves differently with OpenSSH generated keys to paramiko
         # generated keys, so we need a test one.
diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index 924ded1c02a9a..9be7d0b31ff7f 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -31,6 +31,7 @@
 
 TEST_DAG_ID = 'unit_tests_ssh_test_op'
 TEST_CONN_ID = "conn_id_for_testing"
+DEFAULT_TIMEOUT = 10
 CONN_TIMEOUT = 5
 CMD_TIMEOUT = 7
 TIMEOUT = 12
@@ -243,6 +244,24 @@ def test_arg_checking(self):
         assert task_5.cmd_timeout == TIMEOUT
         assert task_5.ssh_hook.conn_timeout == TIMEOUT
 
+        # if no timeouts are provided (conn_timeout, cmd_timeout, timeout), use the hook default
+        # for the conn_timeout and the operator default for the cmd_timeout
+        task_6 = SSHOperator(
+            task_id="task_6",
+            ssh_conn_id=TEST_CONN_ID,
+            command=COMMAND,
+            conn_timeout=None,
+            timeout=None,
+            cmd_timeout=None,
+            dag=self.dag,
+        )
+        try:
+            task_6.execute(None)
+        except Exception:
+            pass
+        assert task_6.cmd_timeout == DEFAULT_TIMEOUT
+        assert task_6.ssh_hook.conn_timeout == DEFAULT_TIMEOUT
+
     @parameterized.expand(
         [
             (COMMAND, False, False),

From 8399a814c1d775148c457dd2c1ed3babb06dac50 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Tue, 24 Aug 2021 07:58:46 -0400
Subject: [PATCH 32/42] Fix spelling mistake in ssh.rst

---
 docs/apache-airflow-providers-ssh/connections/ssh.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/apache-airflow-providers-ssh/connections/ssh.rst b/docs/apache-airflow-providers-ssh/connections/ssh.rst
index bcfb01e128e18..59c8324280eeb 100644
--- a/docs/apache-airflow-providers-ssh/connections/ssh.rst
+++ b/docs/apache-airflow-providers-ssh/connections/ssh.rst
@@ -47,7 +47,7 @@ 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.
-    * ``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.
+    * ``conn_timeout`` - An optional timeout (in seconds) for the TCP connect. Default is ``10``. The parameter to :class:`~airflow.providers.ssh.hooks.ssh.SSHHook` conn_timeout takes precedence.
     * ``timeout`` - Deprecated - use conn_timeout instead.
     * ``compress`` - ``true`` to ask the remote client/server to compress traffic; ``false`` to refuse compression. Default is ``true``.
     * ``no_host_key_check`` - Set to ``false`` to restrict connecting to hosts with no entries in ``~/.ssh/known_hosts`` (Hosts file). This provides maximum protection against trojan horse attacks, but can be troublesome when the ``/etc/ssh/ssh_known_hosts`` file is poorly maintained or connections to new hosts are frequently made. This option forces the user to manually add all new hosts. Default is ``true``, ssh will automatically add new host keys to the user known hosts files.

From fa80b3aeeec9da90e870edd318eac45a8b80fb80 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Wed, 1 Sep 2021 17:23:10 -0400
Subject: [PATCH 33/42] Move details for conn_timeout to SSHHook docstring
 instead of ssh.rst

---
 airflow/providers/ssh/hooks/ssh.py                    | 4 +++-
 docs/apache-airflow-providers-ssh/connections/ssh.rst | 2 +-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
index 9ea2a86cc4515..eda5990288290 100644
--- a/airflow/providers/ssh/hooks/ssh.py
+++ b/airflow/providers/ssh/hooks/ssh.py
@@ -58,7 +58,9 @@ class SSHHook(BaseHook):
     :type key_file: str
     :param port: port of remote host to connect (Default is paramiko SSH_PORT)
     :type port: int
-    :param conn_timeout: timeout for the attempt to connect to the remote_host.
+    :param conn_timeout: timeout (in seconds) for the attempt to connect to the remote_host.
+        The default is 10 seconds. If provided, it will replace the `conn_timeout` which was
+        predefined in the connection of `ssh_conn_id`.
     :type conn_timeout: int
     :param timeout: (Deprecated). timeout for the attempt to connect to the remote_host.
         Use conn_timeout instead.
diff --git a/docs/apache-airflow-providers-ssh/connections/ssh.rst b/docs/apache-airflow-providers-ssh/connections/ssh.rst
index 59c8324280eeb..c9c109ae8720b 100644
--- a/docs/apache-airflow-providers-ssh/connections/ssh.rst
+++ b/docs/apache-airflow-providers-ssh/connections/ssh.rst
@@ -47,7 +47,7 @@ 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.
-    * ``conn_timeout`` - An optional timeout (in seconds) for the TCP connect. Default is ``10``. The parameter to :class:`~airflow.providers.ssh.hooks.ssh.SSHHook` conn_timeout takes precedence.
+    * ``conn_timeout`` - An optional timeout (in seconds) for the TCP connect. Default is ``10``.
     * ``timeout`` - Deprecated - use conn_timeout instead.
     * ``compress`` - ``true`` to ask the remote client/server to compress traffic; ``false`` to refuse compression. Default is ``true``.
     * ``no_host_key_check`` - Set to ``false`` to restrict connecting to hosts with no entries in ``~/.ssh/known_hosts`` (Hosts file). This provides maximum protection against trojan horse attacks, but can be troublesome when the ``/etc/ssh/ssh_known_hosts`` file is poorly maintained or connections to new hosts are frequently made. This option forces the user to manually add all new hosts. Default is ``true``, ssh will automatically add new host keys to the user known hosts files.

From a68185d98fcd0900c2d7445d12a678fbf0de98b7 Mon Sep 17 00:00:00 2001
From: Jed Cunningham <66968678+jedcunningham@users.noreply.github.com>
Date: Wed, 1 Sep 2021 15:35:09 -0600
Subject: [PATCH 34/42] Chart: Make cleanup cronjob cmd/args configurable
 (#17970)

---
 chart/templates/cleanup/cleanup-cronjob.yaml |  8 +++-
 chart/tests/test_cleanup_pods.py             | 47 ++++++++++++++++++++
 chart/values.schema.json                     | 26 +++++++++++
 chart/values.yaml                            |  5 +++
 4 files changed, 84 insertions(+), 2 deletions(-)

diff --git a/chart/templates/cleanup/cleanup-cronjob.yaml b/chart/templates/cleanup/cleanup-cronjob.yaml
index 36e32f435011f..4605dc7055e52 100644
--- a/chart/templates/cleanup/cleanup-cronjob.yaml
+++ b/chart/templates/cleanup/cleanup-cronjob.yaml
@@ -63,8 +63,12 @@ spec:
             - name: airflow-cleanup-pods
               image: {{ template "airflow_image" . }}
               imagePullPolicy: {{ .Values.images.airflow.pullPolicy }}
-              # Don't use entry point here, we don't need to wait on pg-bouncer etc being available.
-              args: ["kubernetes", "cleanup-pods", "--namespace={{ .Release.Namespace }}"]
+              {{- if .Values.cleanup.command }}
+              command: {{ tpl (toYaml .Values.cleanup.command) . | nindent 16 }}
+              {{- end }}
+              {{- if .Values.cleanup.args }}
+              args: {{ tpl (toYaml .Values.cleanup.args) . | nindent 16 }}
+              {{- end }}
               env:
               {{- include "standard_airflow_environment" . | indent 12 }}
               volumeMounts:
diff --git a/chart/tests/test_cleanup_pods.py b/chart/tests/test_cleanup_pods.py
index d28f6cee5c67b..018faa27e004a 100644
--- a/chart/tests/test_cleanup_pods.py
+++ b/chart/tests/test_cleanup_pods.py
@@ -18,6 +18,7 @@
 import unittest
 
 import jmespath
+from parameterized import parameterized
 
 from tests.helm_template_generator import render_chart
 
@@ -104,3 +105,49 @@ def test_should_create_valid_affinity_tolerations_and_node_selector(self):
             "spec.jobTemplate.spec.template.spec.tolerations[0].key",
             docs[0],
         )
+
+    def test_default_command_and_args(self):
+        docs = render_chart(
+            values={"cleanup": {"enabled": True}}, show_only=["templates/cleanup/cleanup-cronjob.yaml"]
+        )
+
+        assert jmespath.search("spec.jobTemplate.spec.template.spec.containers[0].command", docs[0]) is None
+        assert ["bash", "-c", "exec airflow kubernetes cleanup-pods --namespace=default"] == jmespath.search(
+            "spec.jobTemplate.spec.template.spec.containers[0].args", docs[0]
+        )
+
+    @parameterized.expand(
+        [
+            (None, None),
+            (None, ["custom", "args"]),
+            (["custom", "command"], None),
+            (["custom", "command"], ["custom", "args"]),
+        ]
+    )
+    def test_command_and_args_overrides(self, command, args):
+        docs = render_chart(
+            values={"cleanup": {"enabled": True, "command": command, "args": args}},
+            show_only=["templates/cleanup/cleanup-cronjob.yaml"],
+        )
+
+        assert command == jmespath.search(
+            "spec.jobTemplate.spec.template.spec.containers[0].command", docs[0]
+        )
+        assert args == jmespath.search("spec.jobTemplate.spec.template.spec.containers[0].args", docs[0])
+
+    def test_log_groomer_command_and_args_overrides_are_templated(self):
+        docs = render_chart(
+            values={
+                "cleanup": {
+                    "enabled": True,
+                    "command": ["{{ .Release.Name }}"],
+                    "args": ["{{ .Release.Service }}"],
+                }
+            },
+            show_only=["templates/cleanup/cleanup-cronjob.yaml"],
+        )
+
+        assert ["RELEASE-NAME"] == jmespath.search(
+            "spec.jobTemplate.spec.template.spec.containers[0].command", docs[0]
+        )
+        assert ["Helm"] == jmespath.search("spec.jobTemplate.spec.template.spec.containers[0].args", docs[0])
diff --git a/chart/values.schema.json b/chart/values.schema.json
index 491404b426617..a76947e291839 100644
--- a/chart/values.schema.json
+++ b/chart/values.schema.json
@@ -2862,6 +2862,32 @@
                     "type": "string",
                     "default": "*/15 * * * *"
                 },
+                "command": {
+                    "description": "Command to use when running the cleanup cronjob (templated).",
+                    "type": [
+                        "array",
+                        "null"
+                    ],
+                    "items": {
+                        "type": "string"
+                    },
+                    "default": null
+                },
+                "args": {
+                    "description": "Args to use when running the cleanup cronjob (templated).",
+                    "type": [
+                        "array",
+                        "null"
+                    ],
+                    "items": {
+                        "type": "string"
+                    },
+                    "default": [
+                        "bash",
+                        "-c",
+                        "exec airflow kubernetes cleanup-pods --namespace={{ .Release.Namespace }}"
+                    ]
+                },
                 "nodeSelector": {
                     "description": "Select certain nodes for cleanup pods.",
                     "type": "object",
diff --git a/chart/values.yaml b/chart/values.yaml
index a7cbd19d9da9d..efc5db6af8930 100644
--- a/chart/values.yaml
+++ b/chart/values.yaml
@@ -1117,6 +1117,11 @@ cleanup:
   enabled: false
   # Run every 15 minutes
   schedule: "*/15 * * * *"
+  # Command to use when running the cleanup cronjob (templated).
+  command: ~
+  # Args to use when running the cleanup cronjob (templated).
+  args: ["bash", "-c", "exec airflow kubernetes cleanup-pods --namespace={{ .Release.Namespace }}"]
+
 
   # Select certain nodes for airflow cleanup pods.
   nodeSelector: {}

From e177208f0694a7be6ca3875fa6d529be9aeeda2f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= <mik-laj@users.noreply.github.com>
Date: Thu, 2 Sep 2021 08:58:10 +0200
Subject: [PATCH 35/42] Delete unnecessary parameters in EKSPodOperator
 (#17960)

---
 airflow/providers/amazon/aws/hooks/eks.py     | 40 ++++++++++++-------
 airflow/providers/amazon/aws/operators/eks.py | 40 +++++++++----------
 tests/providers/amazon/aws/hooks/test_eks.py  | 36 +++++++++++++++++
 .../amazon/aws/operators/test_eks.py          | 30 ++++++++++++++
 4 files changed, 110 insertions(+), 36 deletions(-)

diff --git a/airflow/providers/amazon/aws/hooks/eks.py b/airflow/providers/amazon/aws/hooks/eks.py
index 8b7cb07715903..9559ca232e56d 100644
--- a/airflow/providers/amazon/aws/hooks/eks.py
+++ b/airflow/providers/amazon/aws/hooks/eks.py
@@ -20,10 +20,11 @@
 import json
 import re
 import tempfile
+import warnings
 from contextlib import contextmanager
 from enum import Enum
 from functools import partial
-from typing import Callable, Dict, List
+from typing import Callable, Dict, List, Optional
 
 import yaml
 from botocore.exceptions import ClientError
@@ -32,10 +33,10 @@
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.utils.json import AirflowJsonEncoder
 
-DEFAULT_CONTEXT_NAME = 'aws'
 DEFAULT_PAGINATION_TOKEN = ''
-DEFAULT_POD_USERNAME = 'aws'
 STS_TOKEN_EXPIRES_IN = 60
+_POD_USERNAME = 'aws'
+_CONTEXT_NAME = 'aws'
 
 
 class ClusterStates(Enum):
@@ -351,21 +352,31 @@ def generate_config_file(
         self,
         eks_cluster_name: str,
         pod_namespace: str,
-        pod_username: str = DEFAULT_POD_USERNAME,
-        pod_context: str = DEFAULT_CONTEXT_NAME,
+        pod_username: Optional[str] = None,
+        pod_context: Optional[str] = None,
     ) -> str:
         """
         Writes the kubeconfig file given an EKS Cluster.
 
-        :param eks_cluster_name: The name of the cluster to create the EKS Managed Nodegroup in.
+        :param eks_cluster_name: The name of the cluster to generate kubeconfig file for.
         :type eks_cluster_name: str
         :param pod_namespace: The namespace to run within kubernetes.
         :type pod_namespace: str
-        :param pod_username: The username under which to execute the pod.
-        :type pod_username: str
-        :param pod_context: The name of the context access parameters to use.
-        :type pod_context: str
         """
+        if pod_username:
+            warnings.warn(
+                "This pod_username parameter is deprecated, because changing the value does not make any "
+                "visible changes to the user.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        if pod_context:
+            warnings.warn(
+                "This pod_context parameter is deprecated, because changing the value does not make any "
+                "visible changes to the user.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
         # Set up the client
         eks_client = self.conn
 
@@ -390,23 +401,22 @@ def generate_config_file(
                     "context": {
                         "cluster": eks_cluster_name,
                         "namespace": pod_namespace,
-                        "user": pod_username,
+                        "user": _POD_USERNAME,
                     },
-                    "name": pod_context,
+                    "name": _CONTEXT_NAME,
                 }
             ],
-            "current-context": pod_context,
+            "current-context": _CONTEXT_NAME,
             "preferences": {},
             "users": [
                 {
-                    "name": pod_username,
+                    "name": _POD_USERNAME,
                     "user": {
                         "token": token,
                     },
                 }
             ],
         }
-
         config_text = yaml.dump(cluster_config, default_flow_style=False)
 
         with tempfile.NamedTemporaryFile(mode='w') as config_file:
diff --git a/airflow/providers/amazon/aws/operators/eks.py b/airflow/providers/amazon/aws/operators/eks.py
index bc831df049d21..9fce83f870ea3 100644
--- a/airflow/providers/amazon/aws/operators/eks.py
+++ b/airflow/providers/amazon/aws/operators/eks.py
@@ -16,17 +16,13 @@
 # under the License.
 
 """This module contains Amazon EKS operators."""
+import warnings
 from datetime import datetime
 from time import sleep
 from typing import Dict, Iterable, List, Optional
 
 from airflow.models import BaseOperator
-from airflow.providers.amazon.aws.hooks.eks import (
-    DEFAULT_CONTEXT_NAME,
-    DEFAULT_POD_USERNAME,
-    ClusterStates,
-    EKSHook,
-)
+from airflow.providers.amazon.aws.hooks.eks import ClusterStates, EKSHook
 from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator
 
 CHECK_INTERVAL_SECONDS = 15
@@ -386,12 +382,8 @@ class EKSPodOperator(KubernetesPodOperator):
     :type in_cluster: bool
     :param namespace: The namespace in which to execute the pod. (templated)
     :type namespace: str
-    :param pod_context: The security context to use while executing the pod. (templated)
-    :type pod_context: str
     :param pod_name: The unique name to give the pod. (templated)
     :type pod_name: str
-    :param pod_username: The username to use while executing the pod. (templated)
-    :type pod_username: str
     :param aws_profile: The named profile containing the credentials for the AWS CLI tool to use.
     :param aws_profile: str
     :param region: Which AWS region the connection should use. (templated)
@@ -409,9 +401,7 @@ class EKSPodOperator(KubernetesPodOperator):
         "cluster_name",
         "in_cluster",
         "namespace",
-        "pod_context",
         "pod_name",
-        "pod_username",
         "aws_conn_id",
         "region",
     } | set(KubernetesPodOperator.template_fields)
@@ -423,9 +413,9 @@ def __init__(
         # file is stored locally in the worker and not in the cluster.
         in_cluster: bool = False,
         namespace: str = DEFAULT_NAMESPACE_NAME,
-        pod_context: str = DEFAULT_CONTEXT_NAME,
+        pod_context: str = None,
         pod_name: str = DEFAULT_POD_NAME,
-        pod_username: str = DEFAULT_POD_USERNAME,
+        pod_username: str = None,
         aws_conn_id: str = DEFAULT_CONN_ID,
         region: Optional[str] = None,
         **kwargs,
@@ -433,9 +423,7 @@ def __init__(
         self.cluster_name = cluster_name
         self.in_cluster = in_cluster
         self.namespace = namespace
-        self.pod_context = pod_context
         self.pod_name = pod_name
-        self.pod_username = pod_username
         self.aws_conn_id = aws_conn_id
         self.region = region
         super().__init__(
@@ -444,17 +432,27 @@ def __init__(
             name=self.pod_name,
             **kwargs,
         )
+        if pod_username:
+            warnings.warn(
+                "This pod_username parameter is deprecated, because changing the value does not make any "
+                "visible changes to the user.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        if pod_context:
+            warnings.warn(
+                "This pod_context parameter is deprecated, because changing the value does not make any "
+                "visible changes to the user.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
 
     def execute(self, context):
         eks_hook = EKSHook(
             aws_conn_id=self.aws_conn_id,
             region_name=self.region,
         )
-
         with eks_hook.generate_config_file(
-            eks_cluster_name=self.cluster_name,
-            pod_namespace=self.namespace,
-            pod_username=self.pod_username,
-            pod_context=self.pod_context,
+            eks_cluster_name=self.cluster_name, pod_namespace=self.namespace
         ) as self.config_file:
             return super().execute(context)
diff --git a/tests/providers/amazon/aws/hooks/test_eks.py b/tests/providers/amazon/aws/hooks/test_eks.py
index 225d03a30e9e1..cb5de1144924d 100644
--- a/tests/providers/amazon/aws/hooks/test_eks.py
+++ b/tests/providers/amazon/aws/hooks/test_eks.py
@@ -17,11 +17,13 @@
 # under the License.
 #
 from copy import deepcopy
+from pathlib import Path
 from typing import Dict, List, Optional, Tuple, Type
 from unittest import mock
 from urllib.parse import ParseResult, urlparse
 
 import pytest
+import yaml
 from _pytest._code import ExceptionInfo
 from botocore.exceptions import ClientError
 from freezegun import freeze_time
@@ -711,6 +713,40 @@ def test_create_nodegroup_handles_launch_template_combinations(
                 )
 
 
+class TestEKSHook:
+    @mock.patch('airflow.providers.amazon.aws.hooks.eks.EKSHook._fetch_sts_token')
+    @mock.patch('airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.conn')
+    def test_generate_config_file(self, mock_conn, mock_fetch_sts_token):
+        mock_fetch_sts_token.return_value = 'STS-token'
+        mock_conn.describe_cluster.return_value = {
+            'cluster': {'certificateAuthority': {'data': 'test-cert'}, 'endpoint': 'test-endpoint'}
+        }
+        hook = EKSHook()
+        with hook.generate_config_file(
+            eks_cluster_name='test-cluster', pod_namespace='k8s-namespace'
+        ) as config_file:
+            config = yaml.load(Path(config_file).read_text())
+            assert config == {
+                'apiVersion': 'v1',
+                'kind': 'Config',
+                'clusters': [
+                    {
+                        'cluster': {'server': 'test-endpoint', 'certificate-authority-data': 'test-cert'},
+                        'name': 'test-cluster',
+                    }
+                ],
+                'contexts': [
+                    {
+                        'context': {'cluster': 'test-cluster', 'namespace': 'k8s-namespace', 'user': 'aws'},
+                        'name': 'aws',
+                    }
+                ],
+                'current-context': 'aws',
+                'preferences': {},
+                'users': [{'name': 'aws', 'user': {'token': 'STS-token'}}],
+            }
+
+
 # Helper methods for repeated assert combinations.
 def assert_all_arn_values_are_valid(expected_arn_values, pattern, arn_under_test) -> None:
     """
diff --git a/tests/providers/amazon/aws/operators/test_eks.py b/tests/providers/amazon/aws/operators/test_eks.py
index 4a9064d8f68e1..864bb673f599c 100644
--- a/tests/providers/amazon/aws/operators/test_eks.py
+++ b/tests/providers/amazon/aws/operators/test_eks.py
@@ -24,6 +24,7 @@
     EKSCreateNodegroupOperator,
     EKSDeleteClusterOperator,
     EKSDeleteNodegroupOperator,
+    EKSPodOperator,
 )
 from tests.providers.amazon.aws.utils.eks_test_constants import (
     NODEROLE_ARN,
@@ -160,3 +161,32 @@ def test_existing_nodegroup(self, mock_delete_nodegroup):
         mock_delete_nodegroup.assert_called_once_with(
             clusterName=self.cluster_name, nodegroupName=self.nodegroup_name
         )
+
+
+class TestEKSPodOperator(unittest.TestCase):
+    @mock.patch('airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.execute')
+    @mock.patch('airflow.providers.amazon.aws.hooks.eks.EKSHook.generate_config_file')
+    @mock.patch('airflow.providers.amazon.aws.hooks.eks.EKSHook.__init__', return_value=None)
+    def test_existing_nodegroup(
+        self, mock_eks_hook, mock_generate_config_file, mock_k8s_pod_operator_execute
+    ):
+        ti_context = mock.MagicMock(name="ti_context")
+
+        op = EKSPodOperator(
+            task_id="run_pod",
+            cluster_name=CLUSTER_NAME,
+            image="amazon/aws-cli:latest",
+            cmds=["sh", "-c", "ls"],
+            labels={"demo": "hello_world"},
+            get_logs=True,
+            # Delete the pod when it reaches its final state, or the execution is interrupted.
+            is_delete_operator_pod=True,
+        )
+        op_return_value = op.execute(ti_context)
+        mock_k8s_pod_operator_execute.assert_called_once_with(ti_context)
+        mock_eks_hook.assert_called_once_with(aws_conn_id='aws_default', region_name=None)
+        mock_generate_config_file.assert_called_once_with(
+            eks_cluster_name=CLUSTER_NAME, pod_namespace='default'
+        )
+        self.assertEqual(mock_k8s_pod_operator_execute.return_value, op_return_value)
+        self.assertEqual(mock_generate_config_file.return_value.__enter__.return_value, op.config_file)

From a5f724916e86e24a93e28fcadc8536c12a1b910e Mon Sep 17 00:00:00 2001
From: Josh Fell <48934154+josh-fell@users.noreply.github.com>
Date: Thu, 2 Sep 2021 05:16:57 -0400
Subject: [PATCH 36/42] Remove default_args pattern + added
 get_current_context() use for Core Airflow example DAGs (#16866)

Co-authored-by: Ash Berlin-Taylor <ash_github@firemirror.com>
---
 airflow/example_dags/example_bash_operator.py |  5 ---
 .../example_branch_datetime_operator.py       |  6 ---
 .../example_branch_day_of_week_operator.py    |  5 ---
 .../example_dags/example_branch_operator.py   |  5 ---
 .../example_branch_python_dop_operator_3.py   |  7 +--
 airflow/example_dags/example_dag_decorator.py |  4 +-
 .../example_kubernetes_executor.py            |  5 ---
 .../example_kubernetes_executor_config.py     |  5 ---
 ...example_passing_params_via_test_command.py |  3 --
 .../example_dags/example_python_operator.py   |  5 ---
 .../example_short_circuit_operator.py         |  5 ---
 airflow/example_dags/example_skip_dag.py      |  6 +--
 .../example_trigger_controller_dag.py         |  1 -
 .../example_trigger_target_dag.py             |  9 ++--
 airflow/example_dags/example_xcom.py          | 43 ++++++++-----------
 airflow/example_dags/example_xcomargs.py      |  2 -
 .../example_dags/tutorial_taskflow_api_etl.py | 10 +----
 .../tutorial_taskflow_api_etl_virtualenv.py   | 14 ++----
 18 files changed, 30 insertions(+), 110 deletions(-)

diff --git a/airflow/example_dags/example_bash_operator.py b/airflow/example_dags/example_bash_operator.py
index 6906fc85ce352..b44dbcb883936 100644
--- a/airflow/example_dags/example_bash_operator.py
+++ b/airflow/example_dags/example_bash_operator.py
@@ -25,13 +25,8 @@
 from airflow.operators.dummy import DummyOperator
 from airflow.utils.dates import days_ago
 
-args = {
-    'owner': 'airflow',
-}
-
 with DAG(
     dag_id='example_bash_operator',
-    default_args=args,
     schedule_interval='0 0 * * *',
     start_date=days_ago(2),
     dagrun_timeout=timedelta(minutes=60),
diff --git a/airflow/example_dags/example_branch_datetime_operator.py b/airflow/example_dags/example_branch_datetime_operator.py
index c0ddeef2cdae5..a9e9c24410b6b 100644
--- a/airflow/example_dags/example_branch_datetime_operator.py
+++ b/airflow/example_dags/example_branch_datetime_operator.py
@@ -27,14 +27,9 @@
 from airflow.operators.dummy import DummyOperator
 from airflow.utils.dates import days_ago
 
-args = {
-    "owner": "airflow",
-}
-
 dag = DAG(
     dag_id="example_branch_datetime_operator",
     start_date=days_ago(2),
-    default_args=args,
     tags=["example"],
     schedule_interval="@daily",
 )
@@ -60,7 +55,6 @@
 dag = DAG(
     dag_id="example_branch_datetime_operator_2",
     start_date=days_ago(2),
-    default_args=args,
     tags=["example"],
     schedule_interval="@daily",
 )
diff --git a/airflow/example_dags/example_branch_day_of_week_operator.py b/airflow/example_dags/example_branch_day_of_week_operator.py
index 597ea902a5de1..4a7a87e368a54 100644
--- a/airflow/example_dags/example_branch_day_of_week_operator.py
+++ b/airflow/example_dags/example_branch_day_of_week_operator.py
@@ -25,14 +25,9 @@
 from airflow.operators.weekday import BranchDayOfWeekOperator
 from airflow.utils.dates import days_ago
 
-args = {
-    "owner": "airflow",
-}
-
 with DAG(
     dag_id="example_weekday_branch_operator",
     start_date=days_ago(2),
-    default_args=args,
     tags=["example"],
     schedule_interval="@daily",
 ) as dag:
diff --git a/airflow/example_dags/example_branch_operator.py b/airflow/example_dags/example_branch_operator.py
index ca85e6dd69ae3..69e63d2418e48 100644
--- a/airflow/example_dags/example_branch_operator.py
+++ b/airflow/example_dags/example_branch_operator.py
@@ -26,13 +26,8 @@
 from airflow.utils.dates import days_ago
 from airflow.utils.edgemodifier import Label
 
-args = {
-    'owner': 'airflow',
-}
-
 with DAG(
     dag_id='example_branch_operator',
-    default_args=args,
     start_date=days_ago(2),
     schedule_interval="@daily",
     tags=['example', 'example2'],
diff --git a/airflow/example_dags/example_branch_python_dop_operator_3.py b/airflow/example_dags/example_branch_python_dop_operator_3.py
index badad5a2e93f8..fd69ab1cc2015 100644
--- a/airflow/example_dags/example_branch_python_dop_operator_3.py
+++ b/airflow/example_dags/example_branch_python_dop_operator_3.py
@@ -26,11 +26,6 @@
 from airflow.operators.python import BranchPythonOperator
 from airflow.utils.dates import days_ago
 
-args = {
-    'owner': 'airflow',
-    'depends_on_past': True,
-}
-
 
 def should_run(**kwargs):
     """
@@ -55,7 +50,7 @@ def should_run(**kwargs):
     dag_id='example_branch_dop_operator_v3',
     schedule_interval='*/1 * * * *',
     start_date=days_ago(2),
-    default_args=args,
+    default_args={'depends_on_past': True},
     tags=['example'],
 ) as dag:
 
diff --git a/airflow/example_dags/example_dag_decorator.py b/airflow/example_dags/example_dag_decorator.py
index fe1ecf828d33c..6bc150eeb38fa 100644
--- a/airflow/example_dags/example_dag_decorator.py
+++ b/airflow/example_dags/example_dag_decorator.py
@@ -25,8 +25,6 @@
 from airflow.operators.email import EmailOperator
 from airflow.utils.dates import days_ago
 
-DEFAULT_ARGS = {"owner": "airflow"}
-
 
 class GetRequestOperator(BaseOperator):
     """Custom operator to sand GET request to provided url"""
@@ -40,7 +38,7 @@ def execute(self, context):
 
 
 # [START dag_decorator_usage]
-@dag(default_args=DEFAULT_ARGS, schedule_interval=None, start_date=days_ago(2), tags=['example'])
+@dag(schedule_interval=None, start_date=days_ago(2), tags=['example'])
 def example_dag_decorator(email: str = 'example@example.com'):
     """
     DAG to send server IP to email.
diff --git a/airflow/example_dags/example_kubernetes_executor.py b/airflow/example_dags/example_kubernetes_executor.py
index 35180e18a19fe..c76c1191e6907 100644
--- a/airflow/example_dags/example_kubernetes_executor.py
+++ b/airflow/example_dags/example_kubernetes_executor.py
@@ -25,13 +25,8 @@
 from airflow.operators.python import PythonOperator
 from airflow.utils.dates import days_ago
 
-args = {
-    'owner': 'airflow',
-}
-
 with DAG(
     dag_id='example_kubernetes_executor',
-    default_args=args,
     schedule_interval=None,
     start_date=days_ago(2),
     tags=['example', 'example2'],
diff --git a/airflow/example_dags/example_kubernetes_executor_config.py b/airflow/example_dags/example_kubernetes_executor_config.py
index 5290dd8e1eb8e..201220b31d9d7 100644
--- a/airflow/example_dags/example_kubernetes_executor_config.py
+++ b/airflow/example_dags/example_kubernetes_executor_config.py
@@ -27,10 +27,6 @@
 from airflow.settings import AIRFLOW_HOME
 from airflow.utils.dates import days_ago
 
-default_args = {
-    'owner': 'airflow',
-}
-
 log = logging.getLogger(__name__)
 
 try:
@@ -38,7 +34,6 @@
 
     with DAG(
         dag_id='example_kubernetes_executor_config',
-        default_args=default_args,
         schedule_interval=None,
         start_date=days_ago(2),
         tags=['example3'],
diff --git a/airflow/example_dags/example_passing_params_via_test_command.py b/airflow/example_dags/example_passing_params_via_test_command.py
index d2ecdfaf1e7ae..e9e004a6fb93b 100644
--- a/airflow/example_dags/example_passing_params_via_test_command.py
+++ b/airflow/example_dags/example_passing_params_via_test_command.py
@@ -59,9 +59,6 @@ def print_env_vars(test_mode):
 
 with DAG(
     "example_passing_params_via_test_command",
-    default_args={
-        "owner": "airflow",
-    },
     schedule_interval='*/1 * * * *',
     start_date=days_ago(1),
     dagrun_timeout=timedelta(minutes=4),
diff --git a/airflow/example_dags/example_python_operator.py b/airflow/example_dags/example_python_operator.py
index a9db34254ebcc..4da7564cc7c8b 100644
--- a/airflow/example_dags/example_python_operator.py
+++ b/airflow/example_dags/example_python_operator.py
@@ -24,13 +24,8 @@
 from airflow.operators.python import PythonOperator, PythonVirtualenvOperator
 from airflow.utils.dates import days_ago
 
-args = {
-    'owner': 'airflow',
-}
-
 with DAG(
     dag_id='example_python_operator',
-    default_args=args,
     schedule_interval=None,
     start_date=days_ago(2),
     tags=['example'],
diff --git a/airflow/example_dags/example_short_circuit_operator.py b/airflow/example_dags/example_short_circuit_operator.py
index 3836ef9490257..d919f91a69131 100644
--- a/airflow/example_dags/example_short_circuit_operator.py
+++ b/airflow/example_dags/example_short_circuit_operator.py
@@ -23,13 +23,8 @@
 from airflow.operators.python import ShortCircuitOperator
 from airflow.utils import dates
 
-args = {
-    'owner': 'airflow',
-}
-
 with DAG(
     dag_id='example_short_circuit_operator',
-    default_args=args,
     start_date=dates.days_ago(2),
     tags=['example'],
 ) as dag:
diff --git a/airflow/example_dags/example_skip_dag.py b/airflow/example_dags/example_skip_dag.py
index 77fbf4a6fcbc4..ba90e225d1131 100644
--- a/airflow/example_dags/example_skip_dag.py
+++ b/airflow/example_dags/example_skip_dag.py
@@ -23,10 +23,6 @@
 from airflow.operators.dummy import DummyOperator
 from airflow.utils.dates import days_ago
 
-args = {
-    'owner': 'airflow',
-}
-
 
 # Create some placeholder operators
 class DummySkipOperator(DummyOperator):
@@ -56,6 +52,6 @@ def create_test_pipeline(suffix, trigger_rule, dag_):
     join >> final
 
 
-with DAG(dag_id='example_skip_dag', default_args=args, start_date=days_ago(2), tags=['example']) as dag:
+with DAG(dag_id='example_skip_dag', start_date=days_ago(2), tags=['example']) as dag:
     create_test_pipeline('1', 'all_success', dag)
     create_test_pipeline('2', 'one_success', dag)
diff --git a/airflow/example_dags/example_trigger_controller_dag.py b/airflow/example_dags/example_trigger_controller_dag.py
index 9d0239919b43b..2987b8b1833ce 100644
--- a/airflow/example_dags/example_trigger_controller_dag.py
+++ b/airflow/example_dags/example_trigger_controller_dag.py
@@ -27,7 +27,6 @@
 
 with DAG(
     dag_id="example_trigger_controller_dag",
-    default_args={"owner": "airflow"},
     start_date=days_ago(2),
     schedule_interval="@once",
     tags=['example'],
diff --git a/airflow/example_dags/example_trigger_target_dag.py b/airflow/example_dags/example_trigger_target_dag.py
index 2b4661b417c62..46125394822f6 100644
--- a/airflow/example_dags/example_trigger_target_dag.py
+++ b/airflow/example_dags/example_trigger_target_dag.py
@@ -28,19 +28,18 @@
 from airflow.utils.dates import days_ago
 
 
-def run_this_func(**context):
+def run_this_func(dag_run):
     """
     Print the payload "message" passed to the DagRun conf attribute.
 
-    :param context: The execution context
-    :type context: dict
+    :param dag_run: The DagRun object
+    :type dag_run: DagRun
     """
-    print(f"Remotely received value of {context['dag_run'].conf['message']} for key=message")
+    print(f"Remotely received value of {dag_run.conf['message']} for key=message")
 
 
 with DAG(
     dag_id="example_trigger_target_dag",
-    default_args={"owner": "airflow"},
     start_date=days_ago(2),
     schedule_interval=None,
     tags=['example'],
diff --git a/airflow/example_dags/example_xcom.py b/airflow/example_dags/example_xcom.py
index de5e14ee0c0e9..60cf3f526891d 100644
--- a/airflow/example_dags/example_xcom.py
+++ b/airflow/example_dags/example_xcom.py
@@ -36,26 +36,17 @@ def push_by_returning(**kwargs):
     return value_2
 
 
-def puller(**kwargs):
-    """Pull all previously pushed XComs and check if the pushed values match the pulled values."""
-    ti = kwargs['ti']
+def _compare_values(pulled_value, check_value):
+    if pulled_value != check_value:
+        raise ValueError(f'The two values differ {pulled_value} and {check_value}')
 
-    # get value_1
-    pulled_value_1 = ti.xcom_pull(key=None, task_ids='push')
-    if pulled_value_1 != value_1:
-        raise ValueError(f'The two values differ {pulled_value_1} and {value_1}')
 
-    # get value_2
-    pulled_value_2 = ti.xcom_pull(task_ids='push_by_returning')
-    if pulled_value_2 != value_2:
-        raise ValueError(f'The two values differ {pulled_value_2} and {value_2}')
+def puller(pulled_value_1, pulled_value_2, **kwargs):
+    """Pull all previously pushed XComs and check if the pushed values match the pulled values."""
 
-    # get both value_1 and value_2
-    pulled_value_1, pulled_value_2 = ti.xcom_pull(key=None, task_ids=['push', 'push_by_returning'])
-    if pulled_value_1 != value_1:
-        raise ValueError(f'The two values differ {pulled_value_1} and {value_1}')
-    if pulled_value_2 != value_2:
-        raise ValueError(f'The two values differ {pulled_value_2} and {value_2}')
+    # Check pulled values from function args
+    _compare_values(pulled_value_1, value_1)
+    _compare_values(pulled_value_2, value_2)
 
 
 def pull_value_from_bash_push(**kwargs):
@@ -70,7 +61,6 @@ def pull_value_from_bash_push(**kwargs):
     'example_xcom',
     schedule_interval="@once",
     start_date=days_ago(2),
-    default_args={'owner': 'airflow'},
     tags=['example'],
 ) as dag:
 
@@ -87,10 +77,12 @@ def pull_value_from_bash_push(**kwargs):
     pull = PythonOperator(
         task_id='puller',
         python_callable=puller,
+        op_kwargs={
+            'pulled_value_1': push1.output['value from pusher 1'],
+            'pulled_value_2': push2.output,
+        },
     )
 
-    pull << [push1, push2]
-
     bash_push = BashOperator(
         task_id='bash_push',
         bash_command='echo "bash_push demo"  && '
@@ -98,13 +90,12 @@ def pull_value_from_bash_push(**kwargs):
         '{{ ti.xcom_push(key="manually_pushed_value", value="manually_pushed_value") }}" && '
         'echo "value_by_return"',
     )
+
     bash_pull = BashOperator(
         task_id='bash_pull',
         bash_command='echo "bash pull demo" && '
-        'echo "The xcom pushed manually is '
-        '"{{ ti.xcom_pull(task_ids="bash_push", key="manually_pushed_value") }}" && '
-        'echo "The returned_value xcom is '
-        '"{{ ti.xcom_pull(task_ids="bash_push", key="return_value") }}" && '
+        f'echo "The xcom pushed manually is {bash_push.output["manually_pushed_value"]}" && '
+        f'echo "The returned_value xcom is {bash_push.output}" && '
         'echo "finished"',
         do_xcom_push=False,
     )
@@ -115,3 +106,7 @@ def pull_value_from_bash_push(**kwargs):
     )
 
     [bash_pull, python_pull_from_bash] << bash_push
+
+    # Task dependencies created via `XComArgs`:
+    #   push1 >> pull
+    #   push2 >> pull
diff --git a/airflow/example_dags/example_xcomargs.py b/airflow/example_dags/example_xcomargs.py
index 0fb728ef47e55..bf7dc8a242816 100644
--- a/airflow/example_dags/example_xcomargs.py
+++ b/airflow/example_dags/example_xcomargs.py
@@ -42,7 +42,6 @@ def print_value(value):
 
 with DAG(
     dag_id='example_xcom_args',
-    default_args={'owner': 'airflow'},
     start_date=days_ago(2),
     schedule_interval=None,
     tags=['example'],
@@ -57,7 +56,6 @@ def print_value(value):
 
 with DAG(
     "example_xcom_args_with_operators",
-    default_args={'owner': 'airflow'},
     start_date=days_ago(2),
     schedule_interval=None,
     tags=['example'],
diff --git a/airflow/example_dags/tutorial_taskflow_api_etl.py b/airflow/example_dags/tutorial_taskflow_api_etl.py
index edc589ce4f117..bc24c3ff82a32 100644
--- a/airflow/example_dags/tutorial_taskflow_api_etl.py
+++ b/airflow/example_dags/tutorial_taskflow_api_etl.py
@@ -26,17 +26,9 @@
 
 # [END import_module]
 
-# [START default_args]
-# These args will get passed on to each operator
-# You can override them on a per-task basis during operator initialization
-default_args = {
-    'owner': 'airflow',
-}
-# [END default_args]
-
 
 # [START instantiate_dag]
-@dag(default_args=default_args, schedule_interval=None, start_date=days_ago(2), tags=['example'])
+@dag(schedule_interval=None, start_date=days_ago(2), tags=['example'])
 def tutorial_taskflow_api_etl():
     """
     ### TaskFlow API Tutorial Documentation
diff --git a/airflow/example_dags/tutorial_taskflow_api_etl_virtualenv.py b/airflow/example_dags/tutorial_taskflow_api_etl_virtualenv.py
index d2723beee6d81..bdbf47627435e 100644
--- a/airflow/example_dags/tutorial_taskflow_api_etl_virtualenv.py
+++ b/airflow/example_dags/tutorial_taskflow_api_etl_virtualenv.py
@@ -19,24 +19,14 @@
 
 # [START tutorial]
 # [START import_module]
-import json
-
 from airflow.decorators import dag, task
 from airflow.utils.dates import days_ago
 
 # [END import_module]
 
-# [START default_args]
-# These args will get passed on to each operator
-# You can override them on a per-task basis during operator initialization
-default_args = {
-    'owner': 'airflow',
-}
-# [END default_args]
-
 
 # [START instantiate_dag]
-@dag(default_args=default_args, schedule_interval=None, start_date=days_ago(2), tags=['example'])
+@dag(schedule_interval=None, start_date=days_ago(2), tags=['example'])
 def tutorial_taskflow_api_etl_virtualenv():
     """
     ### TaskFlow API Tutorial Documentation
@@ -61,6 +51,8 @@ def extract():
         pipeline. In this case, getting data is simulated by reading from a
         hardcoded JSON string.
         """
+        import json
+
         data_string = '{"1001": 301.27, "1002": 433.21, "1003": 502.22}'
 
         order_data_dict = json.loads(data_string)

From e9513b2feacb9ae159e0fe57d538a6c35b4497b3 Mon Sep 17 00:00:00 2001
From: Jarek Potiuk <jarek@potiuk.com>
Date: Thu, 2 Sep 2021 11:24:16 +0200
Subject: [PATCH 37/42] Fix constraint generation properly (#17964)

The #17939 did not fix the problem finally. It turned out that
one more change was needed - since we now always upgrade to latest
dependencies in `push` and `schedule` type of build we do not need
to check for the variable UPGRADE_TO_NEWER_DEPENDENCIES (which
was not set in "Build Image" step.

This fixes it, but also changes the constraint generation to add
comments in the generated constraint files, describing how and
why the files are generated.
---
 CI.rst                                        |  8 ++--
 scripts/ci/selective_ci_checks.sh             | 21 ++++----
 .../in_container/run_generate_constraints.sh  | 48 ++++++++++++++++---
 3 files changed, 56 insertions(+), 21 deletions(-)

diff --git a/CI.rst b/CI.rst
index 61b0d503cf739..6cc97cccb0bb2 100644
--- a/CI.rst
+++ b/CI.rst
@@ -124,7 +124,7 @@ You can use those variables when you try to reproduce the build locally.
 |                                         |             |              |            | case of Postgres or MySQL. However,             |
 |                                         |             |              |            | it requires to perform manual init/reset        |
 |                                         |             |              |            | if you stop the environment.                    |
-+-----------------------------------------+-------------+--------------+-------------+-------------------------------------------------+
++-----------------------------------------+-------------+--------------+-------------+------------------------------------------------+
 |                                                           Mount variables                                                           |
 +-----------------------------------------+-------------+--------------+------------+-------------------------------------------------+
 | ``MOUNT_SELECTED_LOCAL_SOURCES``        |     true    |    false     |    false   | Determines whether local sources are            |
@@ -209,7 +209,7 @@ You can use those variables when you try to reproduce the build locally.
 +-----------------------------------------+-------------+--------------+------------+-------------------------------------------------+
 |                                                            Git variables                                                            |
 +-----------------------------------------+-------------+--------------+------------+-------------------------------------------------+
-| COMMIT_SHA                              |             | GITHUB_SHA   | GITHUB_SHA | SHA of the commit of the build is run           |
+| ``COMMIT_SHA``                          |             | GITHUB_SHA   | GITHUB_SHA | SHA of the commit of the build is run           |
 +-----------------------------------------+-------------+--------------+------------+-------------------------------------------------+
 |                                                         Verbosity variables                                                         |
 +-----------------------------------------+-------------+--------------+------------+-------------------------------------------------+
@@ -255,7 +255,7 @@ You can use those variables when you try to reproduce the build locally.
 |                                         |             |              |            | of the airflow repository                       |
 |                                         |             |              |            | ("constraints-main, "constraints-2-0")          |
 |                                         |             |              |            | but when this flag is set to anything but false |
-|                                         |             |              |            | (for example commit SHA), they are not used     |
+|                                         |             |              |            | (for example random value), they are not used   |
 |                                         |             |              |            | used and "eager" upgrade strategy is used       |
 |                                         |             |              |            | when installing dependencies. We set it         |
 |                                         |             |              |            | to true in case of direct pushes (merges)       |
@@ -265,7 +265,7 @@ You can use those variables when you try to reproduce the build locally.
 |                                         |             |              |            | we automatically push latest set of             |
 |                                         |             |              |            | "tested" constraints to the repository.         |
 |                                         |             |              |            |                                                 |
-|                                         |             |              |            | Setting the value to commit SHA is best way     |
+|                                         |             |              |            | Setting the value to random value is best way   |
 |                                         |             |              |            | to assure that constraints are upgraded even if |
 |                                         |             |              |            | there is no change to setup.py                  |
 |                                         |             |              |            |                                                 |
diff --git a/scripts/ci/selective_ci_checks.sh b/scripts/ci/selective_ci_checks.sh
index 15020ad366538..10aa128c7614d 100755
--- a/scripts/ci/selective_ci_checks.sh
+++ b/scripts/ci/selective_ci_checks.sh
@@ -43,16 +43,14 @@ else
 fi
 
 function check_upgrade_to_newer_dependencies_needed() {
-    # shellcheck disable=SC2153
-    if [[ "${UPGRADE_TO_NEWER_DEPENDENCIES}" != "false" ||
-            ${GITHUB_EVENT_NAME=} == 'push' || ${GITHUB_EVENT_NAME=} == "scheduled" ]]; then
-        # Trigger upgrading to latest constraints when we are in push or schedule event or when it is forced
-        # By UPGRADE_TO_NEWER_DEPENDENCIES set to non-false. The variable is set to
-        # SHA of the merge commit - so that it always triggers rebuilding layer in the docker image
+    if [[ ${GITHUB_EVENT_NAME=} == 'push' || ${GITHUB_EVENT_NAME=} == "scheduled" ]]; then
+        # Trigger upgrading to latest constraints when we are in push or schedule event. We use the
+        # random string so that it always triggers rebuilding layer in the docker image
         # Each build that upgrades to latest constraints will get truly latest constraints, not those
-        # Cached in the image if we set it to "true". This upgrade_to_newer_dependencies variable
-        # can later be overridden in case we find that setup.* files changed (see below)
-        upgrade_to_newer_dependencies="${INCOMING_COMMIT_SHA}"
+        # cached in the image because docker layer will get invalidated.
+        # This upgrade_to_newer_dependencies variable can later be overridden
+        # in case we find that any of the setup.* files changed (see below)
+        upgrade_to_newer_dependencies="${RANDOM}"
     fi
 }
 
@@ -353,8 +351,9 @@ function check_if_setup_files_changed() {
 
     if [[ $(count_changed_files) != "0" ]]; then
         # In case the setup files changed, we automatically force upgrading to newer dependencies
-        # no matter what was set before
-        upgrade_to_newer_dependencies="${INCOMING_COMMIT_SHA}"
+        # no matter what was set before. We set it to random number to make sure that it will be
+        # always invalidating the layer in Docker that triggers installing the dependencies
+        upgrade_to_newer_dependencies="${RANDOM}"
     fi
     start_end::group_end
 }
diff --git a/scripts/in_container/run_generate_constraints.sh b/scripts/in_container/run_generate_constraints.sh
index d71fc9d9a8bac..402ac853541a8 100755
--- a/scripts/in_container/run_generate_constraints.sh
+++ b/scripts/in_container/run_generate_constraints.sh
@@ -20,6 +20,12 @@
 
 CONSTRAINTS_DIR="/files/constraints-${PYTHON_MAJOR_MINOR_VERSION}"
 
+LATEST_CONSTRAINT_FILE="${CONSTRAINTS_DIR}/original-${AIRFLOW_CONSTRAINTS}-${PYTHON_MAJOR_MINOR_VERSION}.txt"
+CURRENT_CONSTRAINT_FILE="${CONSTRAINTS_DIR}/${AIRFLOW_CONSTRAINTS}-${PYTHON_MAJOR_MINOR_VERSION}.txt"
+
+mkdir -pv "${CONSTRAINTS_DIR}"
+
+
 if [[ ${GENERATE_CONSTRAINTS_MODE} == "no-providers" ]]; then
     AIRFLOW_CONSTRAINTS="constraints-no-providers"
     NO_PROVIDERS_EXTRAS=$(python -c 'import setup; print(",".join(setup.CORE_EXTRAS_REQUIREMENTS.keys()))')
@@ -31,11 +37,34 @@ if [[ ${GENERATE_CONSTRAINTS_MODE} == "no-providers" ]]; then
     echo "Install airflow with [${NO_PROVIDERS_EXTRAS}] extras only (uninstall all packages first)."
     echo
     install_local_airflow_with_eager_upgrade "[${NO_PROVIDERS_EXTRAS}]"
+    cat <<EOF >"${CURRENT_CONSTRAINT_FILE}"
+#
+# This constraints file was automatically generated on $(date -u +'%Y-%m-%dT%H:%M:%SZ')
+# via "eager-upgrade" mechanism of PIP. For the "${DEFAULT_BRANCH}" branch of Airflow.
+# This variant of constraints install just the 'bare' 'apache-airflow' package build from the HEAD of
+# the branch, without installing any of the providers.
+#
+# Those constraints represent the "newest" dependencies airflow could use, if providers did not limit
+# Airflow in any way.
+#
+EOF
 elif [[ ${GENERATE_CONSTRAINTS_MODE} == "source-providers" ]]; then
     AIRFLOW_CONSTRAINTS="constraints-source-providers"
     echo
     echo "Providers are already installed from sources."
     echo
+    cat <<EOF >"${CURRENT_CONSTRAINT_FILE}"
+#
+# This constraints file was automatically generated on $(date -u +'%Y-%m-%dT%H:%M:%SZ')
+# via "eager-upgrade" mechanism of PIP. For the "${DEFAULT_BRANCH}" branch of Airflow.
+# This variant of constraints install uses the HEAD of the branch version of both
+# 'apache-airflow' package nd all available community provider packages.
+#
+# Those constraints represent the dependencies that are used by all pull requests when they are build in CI.
+# They represent "latest" and greatest set of constraints that HEAD of the "apache-airflow" package should
+# Install with "HEAD" of providers. Those are the only constraints that are used by our CI builds.
+#
+EOF
 elif [[ ${GENERATE_CONSTRAINTS_MODE} == "pypi-providers" ]]; then
     AIRFLOW_CONSTRAINTS="constraints"
     echo
@@ -43,6 +72,18 @@ elif [[ ${GENERATE_CONSTRAINTS_MODE} == "pypi-providers" ]]; then
     echo
     install_all_providers_from_pypi_with_eager_upgrade
 else
+    cat <<EOF >"${CURRENT_CONSTRAINT_FILE}"
+#
+# This constraints file was automatically generated on $(date -u +'%Y-%m-%dT%H:%M:%SZ')
+# via "eager-upgrade" mechanism of PIP. For the "${DEFAULT_BRANCH}" branch of Airflow.
+# This variant of constraints install uses the HEAD of the branch version for 'apache-airflow' but installs
+# the providers from PIP-released packages at the moment of the constraint generation.
+#
+# Those constraints are actually those that that regular users use to install released version of Airflow.
+# We also use those constraints after "apache-airflow" is released and the constraints are tagged with
+# "constraints-X.Y.Z" tag to build the production image for that version.
+#
+EOF
     echo
     echo "${COLOR_RED}Error! GENERATE_CONSTRAINTS_MODE has wrong value: '${GENERATE_CONSTRAINTS_MODE}' ${COLOR_RESET}"
     echo
@@ -51,11 +92,6 @@ fi
 
 readonly AIRFLOW_CONSTRAINTS
 
-LATEST_CONSTRAINT_FILE="${CONSTRAINTS_DIR}/original-${AIRFLOW_CONSTRAINTS}-${PYTHON_MAJOR_MINOR_VERSION}.txt"
-CURRENT_CONSTRAINT_FILE="${CONSTRAINTS_DIR}/${AIRFLOW_CONSTRAINTS}-${PYTHON_MAJOR_MINOR_VERSION}.txt"
-
-mkdir -pv "${CONSTRAINTS_DIR}"
-
 CONSTRAINTS_LOCATION="https://raw.githubusercontent.com/${CONSTRAINTS_GITHUB_REPOSITORY}/${DEFAULT_CONSTRAINTS_BRANCH}/${AIRFLOW_CONSTRAINTS}-${PYTHON_MAJOR_MINOR_VERSION}.txt"
 readonly CONSTRAINTS_LOCATION
 
@@ -69,7 +105,7 @@ echo
 pip freeze | sort | \
     grep -v "apache_airflow" | \
     grep -v "@" | \
-    grep -v "/opt/airflow" >"${CURRENT_CONSTRAINT_FILE}"
+    grep -v "/opt/airflow" >>"${CURRENT_CONSTRAINT_FILE}"
 
 echo
 echo "Constraints generated in ${CURRENT_CONSTRAINT_FILE}"

From 1e8ce4fb11d068f2caf50666a3d28f2705493fd5 Mon Sep 17 00:00:00 2001
From: Ephraim Anierobi <splendidzigy24@gmail.com>
Date: Thu, 2 Sep 2021 13:14:36 +0100
Subject: [PATCH 38/42] Fix max_active_runs not allowing moving of queued
 dagruns to running (#17945)

Currently, if you set max_active_runs for a dag and that dag has many queued dagruns
with execution dates older than another dag's queued dagruns, airflow will not move
the newer queued dagruns to running with the effect that only one dagruns would be in running at any time

This PR fixes this by updating the DagRun.last_scheduling_decision whenever a decision of scheduling
was made
---
 airflow/jobs/scheduler_job.py    |  1 +
 tests/jobs/test_scheduler_job.py | 37 ++++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/airflow/jobs/scheduler_job.py b/airflow/jobs/scheduler_job.py
index b3cb5314a4091..add07d96a2310 100644
--- a/airflow/jobs/scheduler_job.py
+++ b/airflow/jobs/scheduler_job.py
@@ -1003,6 +1003,7 @@ def _update_state(dag_run):
                 continue
             active_runs = active_runs_of_dags[dag_run.dag_id]
             if dag.max_active_runs and active_runs >= dag.max_active_runs:
+                dag_run.last_scheduling_decision = timezone.utcnow()
                 self.log.debug(
                     "DAG %s already has %d active runs, not moving any more runs to RUNNING state %s",
                     dag.dag_id,
diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py
index 964f7375558c3..0be44789b914c 100644
--- a/tests/jobs/test_scheduler_job.py
+++ b/tests/jobs/test_scheduler_job.py
@@ -2881,6 +2881,43 @@ def test_do_schedule_max_active_runs_and_manual_trigger(self, dag_maker):
         # Assert that the other one is queued
         assert len(DagRun.find(dag_id=dag.dag_id, state=State.QUEUED, session=session)) == 1
 
+    def test_max_active_runs_in_a_dag_doesnt_stop_running_dagruns_in_otherdags(self, dag_maker):
+        session = settings.Session()
+        with dag_maker('test_dag1', max_active_runs=1) as dag:
+            DummyOperator(task_id='mytask')
+        date = dag.following_schedule(DEFAULT_DATE)
+        for _ in range(30):
+            dr = dag.create_dagrun(run_type=DagRunType.SCHEDULED, state=State.QUEUED, execution_date=date)
+            date = dr.execution_date + timedelta(hours=1)
+
+        date = timezone.datetime(2020, 1, 1)
+        with dag_maker('test_dag2', start_date=date) as dag2:
+            DummyOperator(task_id='mytask')
+
+        for _ in range(10):
+            dr = dag2.create_dagrun(run_type=DagRunType.SCHEDULED, state=State.QUEUED, execution_date=date)
+            date = dr.execution_date + timedelta(hours=1)
+
+        self.scheduler_job = SchedulerJob(subdir=os.devnull)
+        self.scheduler_job.executor = MockExecutor(do_update=False)
+        self.scheduler_job.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent)
+
+        self.scheduler_job._start_queued_dagruns(session)
+        session.flush()
+        self.scheduler_job._start_queued_dagruns(session)
+        session.flush()
+
+        assert (
+            len(
+                session.query(DagRun)
+                .filter(DagRun.dag_id == dag.dag_id)
+                .filter(DagRun.state == State.RUNNING)
+                .all()
+            )
+            == 1
+        )
+        assert len(session.query(DagRun).filter(DagRun.state == State.RUNNING).all()) == 11
+
     @pytest.mark.parametrize(
         "state, start_date, end_date",
         [

From ff925ae3cc5b46199edbe15139913de1fa900f53 Mon Sep 17 00:00:00 2001
From: bbenshalom <benny1886@gmail.com>
Date: Thu, 2 Sep 2021 22:32:45 +0300
Subject: [PATCH 39/42] Add DAG run endpoint for marking a dagrun success or
 failed(#17839)

Co-authored-by: bbenshalom <bbenshalom@outbrain.com>
Co-authored-by: Tzu-ping Chung <uranusjr@gmail.com>
Co-authored-by: Ephraim Anierobi <splendidzigy24@gmail.com>
---
 .../endpoints/dag_run_endpoint.py             | 31 ++++++-
 airflow/api_connexion/openapi/v1.yaml         | 39 ++++++++
 .../api_connexion/schemas/dag_run_schema.py   | 10 ++-
 .../endpoints/test_dag_run_endpoint.py        | 88 +++++++++++++++++++
 4 files changed, 166 insertions(+), 2 deletions(-)

diff --git a/airflow/api_connexion/endpoints/dag_run_endpoint.py b/airflow/api_connexion/endpoints/dag_run_endpoint.py
index feb6d24cf3247..e816aac27e157 100644
--- a/airflow/api_connexion/endpoints/dag_run_endpoint.py
+++ b/airflow/api_connexion/endpoints/dag_run_endpoint.py
@@ -14,6 +14,8 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from typing import Optional
+
 from flask import current_app, g, request
 from marshmallow import ValidationError
 from sqlalchemy import or_
@@ -27,11 +29,12 @@
     dagrun_collection_schema,
     dagrun_schema,
     dagruns_batch_form_schema,
+    set_dagrun_state_form_schema,
 )
 from airflow.models import DagModel, DagRun
 from airflow.security import permissions
 from airflow.utils.session import provide_session
-from airflow.utils.state import State
+from airflow.utils.state import DagRunState, State
 from airflow.utils.types import DagRunType
 
 
@@ -271,3 +274,29 @@ def post_dag_run(dag_id, session):
         )
 
     raise AlreadyExists(detail=f"DAGRun with DAG ID: '{dag_id}' and DAGRun ID: '{run_id}' already exists")
+
+
+@security.requires_access(
+    [
+        (permissions.ACTION_CAN_READ, permissions.RESOURCE_DAG),
+        (permissions.ACTION_CAN_EDIT, permissions.RESOURCE_DAG_RUN),
+    ]
+)
+@provide_session
+def update_dag_run_state(dag_id: str, dag_run_id: str, session) -> dict:
+    """Set a state of a dag run."""
+    dag_run: Optional[DagRun] = (
+        session.query(DagRun).filter(DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id).one_or_none()
+    )
+    if dag_run is None:
+        error_message = f'Dag Run id {dag_run_id} not found in dag {dag_id}'
+        raise NotFound(error_message)
+    try:
+        post_body = set_dagrun_state_form_schema.load(request.json)
+    except ValidationError as err:
+        raise BadRequest(detail=str(err))
+
+    state = post_body['state']
+    dag_run.set_state(state=DagRunState(state))
+    session.merge(dag_run)
+    return dagrun_schema.dump(dag_run)
diff --git a/airflow/api_connexion/openapi/v1.yaml b/airflow/api_connexion/openapi/v1.yaml
index 2ce6804bbd888..8dae5f2fbc6b6 100644
--- a/airflow/api_connexion/openapi/v1.yaml
+++ b/airflow/api_connexion/openapi/v1.yaml
@@ -674,6 +674,35 @@ paths:
         '404':
           $ref: '#/components/responses/NotFound'
 
+    patch:
+      summary: Modify a DAG run
+      description: Modify a DAG run
+      x-openapi-router-controller: airflow.api_connexion.endpoints.dag_run_endpoint
+      operationId: update_dag_run_state
+      tags: [ UpdateDagRunState ]
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/UpdateDagRunState'
+
+      responses:
+        '200':
+          description: Success.
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/DAGRun'
+        '400':
+          $ref: '#/components/responses/BadRequest'
+        '401':
+          $ref: '#/components/responses/Unauthenticated'
+        '403':
+          $ref: '#/components/responses/PermissionDenied'
+        '404':
+          $ref: '#/components/responses/NotFound'
+
   /eventLogs:
     get:
       summary: List log entries
@@ -1996,6 +2025,16 @@ components:
       required:
         - dag_id
 
+    UpdateDagRunState:
+      type: object
+      properties:
+        state:
+          description: The state to set this DagRun
+          type: string
+          enum:
+            - success
+            - failed
+
     DAGRunCollection:
       type: object
       description: Collection of DAG runs.
diff --git a/airflow/api_connexion/schemas/dag_run_schema.py b/airflow/api_connexion/schemas/dag_run_schema.py
index 62ac172ce5525..fa51f10ed48c0 100644
--- a/airflow/api_connexion/schemas/dag_run_schema.py
+++ b/airflow/api_connexion/schemas/dag_run_schema.py
@@ -18,7 +18,7 @@
 import json
 from typing import List, NamedTuple
 
-from marshmallow import fields, post_dump, pre_load
+from marshmallow import fields, post_dump, pre_load, validate
 from marshmallow.schema import Schema
 from marshmallow.validate import Range
 from marshmallow_sqlalchemy import SQLAlchemySchema, auto_field
@@ -29,6 +29,7 @@
 from airflow.api_connexion.schemas.enum_schemas import DagStateField
 from airflow.models.dagrun import DagRun
 from airflow.utils import timezone
+from airflow.utils.state import DagRunState
 from airflow.utils.types import DagRunType
 
 
@@ -104,6 +105,12 @@ def autofill(self, data, **kwargs):
         return data
 
 
+class SetDagRunStateFormSchema(Schema):
+    """Schema for handling the request of setting state of DAG run"""
+
+    state = DagStateField(validate=validate.OneOf([DagRunState.SUCCESS.value, DagRunState.FAILED.value]))
+
+
 class DAGRunCollection(NamedTuple):
     """List of DAGRuns with metadata"""
 
@@ -141,4 +148,5 @@ class Meta:
 
 dagrun_schema = DAGRunSchema()
 dagrun_collection_schema = DAGRunCollectionSchema()
+set_dagrun_state_form_schema = SetDagRunStateFormSchema()
 dagruns_batch_form_schema = DagRunsBatchFormSchema()
diff --git a/tests/api_connexion/endpoints/test_dag_run_endpoint.py b/tests/api_connexion/endpoints/test_dag_run_endpoint.py
index 71b52747dcff7..47d89cba7c44a 100644
--- a/tests/api_connexion/endpoints/test_dag_run_endpoint.py
+++ b/tests/api_connexion/endpoints/test_dag_run_endpoint.py
@@ -18,10 +18,12 @@
 from unittest import mock
 
 import pytest
+from freezegun import freeze_time
 from parameterized import parameterized
 
 from airflow.api_connexion.exceptions import EXCEPTIONS_LINK_MAP
 from airflow.models import DAG, DagModel, DagRun
+from airflow.operators.dummy import DummyOperator
 from airflow.security import permissions
 from airflow.utils import timezone
 from airflow.utils.session import create_session, provide_session
@@ -1154,3 +1156,89 @@ def test_should_raises_403_unauthorized(self):
             environ_overrides={'REMOTE_USER': "test_view_dags"},
         )
         assert response.status_code == 403
+
+
+class TestPostSetDagRunState(TestDagRunEndpoint):
+    @pytest.mark.parametrize("state", ["failed", "success"])
+    @freeze_time(TestDagRunEndpoint.default_time)
+    def test_should_respond_200(self, state, dag_maker):
+        dag_id = "TEST_DAG_ID"
+        dag_run_id = 'TEST_DAG_RUN_ID'
+        with dag_maker(dag_id):
+            DummyOperator(task_id='task_id')
+        dag_maker.create_dagrun(run_id=dag_run_id)
+
+        request_json = {"state": state}
+
+        response = self.client.patch(
+            f"api/v1/dags/{dag_id}/dagRuns/{dag_run_id}",
+            json=request_json,
+            environ_overrides={"REMOTE_USER": "test"},
+        )
+
+        assert response.status_code == 200
+        assert response.json == {
+            'conf': {},
+            'dag_id': dag_id,
+            'dag_run_id': dag_run_id,
+            'end_date': self.default_time,
+            'execution_date': dag_maker.start_date.isoformat(),
+            'external_trigger': False,
+            'logical_date': dag_maker.start_date.isoformat(),
+            'start_date': dag_maker.start_date.isoformat(),
+            'state': state,
+        }
+
+    @pytest.mark.parametrize('invalid_state', ["running", "queued"])
+    @freeze_time(TestDagRunEndpoint.default_time)
+    def test_should_response_400_for_non_existing_dag_run_state(self, invalid_state, dag_maker):
+        dag_id = "TEST_DAG_ID"
+        dag_run_id = 'TEST_DAG_RUN_ID'
+        with dag_maker(dag_id):
+            DummyOperator(task_id='task_id')
+        dag_maker.create_dagrun(run_id=dag_run_id)
+
+        request_json = {"state": invalid_state}
+
+        response = self.client.patch(
+            "api/v1/dags/TEST_DAG_ID/dagRuns/TEST_DAG_RUN_ID_1",
+            json=request_json,
+            environ_overrides={"REMOTE_USER": "test"},
+        )
+        assert response.status_code == 400
+        assert response.json == {
+            'detail': f"'{invalid_state}' is not one of ['success', 'failed'] - 'state'",
+            'status': 400,
+            'title': 'Bad Request',
+            'type': EXCEPTIONS_LINK_MAP[400],
+        }
+
+    def test_should_raises_401_unauthenticated(self, session):
+        response = self.client.patch(
+            "api/v1/dags/TEST_DAG_ID/dagRuns/TEST_DAG_RUN_ID_1",
+            json={
+                "state": 'success',
+            },
+        )
+
+        assert_401(response)
+
+    def test_should_raise_403_forbidden(self):
+        response = self.client.patch(
+            "api/v1/dags/TEST_DAG_ID/dagRuns/TEST_DAG_RUN_ID_1",
+            json={
+                "state": 'success',
+            },
+            environ_overrides={'REMOTE_USER': "test_no_permissions"},
+        )
+        assert response.status_code == 403
+
+    def test_should_respond_404(self):
+        response = self.client.patch(
+            "api/v1/dags/INVALID_DAG_ID/dagRuns/TEST_DAG_RUN_ID_1",
+            json={
+                "state": 'success',
+            },
+            environ_overrides={"REMOTE_USER": "test"},
+        )
+        assert response.status_code == 404

From 68228337487cf9ad866922933703de02ac86cdc2 Mon Sep 17 00:00:00 2001
From: Brent Bovenzi <brent.bovenzi@gmail.com>
Date: Thu, 2 Sep 2021 23:07:07 +0200
Subject: [PATCH 40/42] Fix blank dag dependencies view (#17990)

* Fix blank dag dependencies view

* calculate graph if node and edges are empty
---
 airflow/www/templates/airflow/dag_dependencies.html | 2 +-
 airflow/www/views.py                                | 7 +++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/airflow/www/templates/airflow/dag_dependencies.html b/airflow/www/templates/airflow/dag_dependencies.html
index feca4bff8fc29..c893c75bd06ce 100644
--- a/airflow/www/templates/airflow/dag_dependencies.html
+++ b/airflow/www/templates/airflow/dag_dependencies.html
@@ -44,7 +44,7 @@ <h2>
     <span class="legend-item trigger">trigger</span>
     <span class="legend-item sensor">sensor</span>
   </div>
-  <div style="float:right">Last refresh: {{ last_refresh }}</div>
+  <div style="float:right">Last refresh: <time datetime="{{ last_refresh }}">{{ last_refresh }}</time></div>
 </div>
 <div id="error" style="display: none; margin-top: 10px;" class="alert alert-danger" role="alert">
   <span class="glyphicon glyphicon-exclamation-sign" aria-hidden="true"></span>
diff --git a/airflow/www/views.py b/airflow/www/views.py
index 146145b2abc98..d5428b3629a46 100644
--- a/airflow/www/views.py
+++ b/airflow/www/views.py
@@ -4294,7 +4294,10 @@ def list(self):
         """Display DAG dependencies"""
         title = "DAG Dependencies"
 
-        if timezone.utcnow() > self.last_refresh + self.refresh_interval:
+        if not self.nodes or not self.edges:
+            self._calculate_graph()
+            self.last_refresh = timezone.utcnow()
+        elif timezone.utcnow() > self.last_refresh + self.refresh_interval:
             max_last_updated = SerializedDagModel.get_max_last_updated_datetime()
             if max_last_updated is None or max_last_updated > self.last_refresh:
                 self._calculate_graph()
@@ -4305,7 +4308,7 @@ def list(self):
             title=title,
             nodes=self.nodes,
             edges=self.edges,
-            last_refresh=self.last_refresh.strftime("%Y-%m-%d %H:%M:%S"),
+            last_refresh=self.last_refresh,
             arrange=conf.get("webserver", "dag_orientation"),
             width=request.args.get("width", "100%"),
             height=request.args.get("height", "800"),

From 81cfaa10831793564c52a450526ff36493b65fbb Mon Sep 17 00:00:00 2001
From: Bas Harenslak <BasPH@users.noreply.github.com>
Date: Fri, 3 Sep 2021 00:54:32 +0200
Subject: [PATCH 41/42] Serialize the template_ext attribute to show it in UI
 (#17985)

Co-authored-by: Bas Harenslak <bas@astronomer.io>
---
 airflow/models/baseoperator.py                | 1 +
 airflow/serialization/serialized_objects.py   | 1 +
 tests/serialization/test_dag_serialization.py | 5 +++++
 3 files changed, 7 insertions(+)

diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py
index 3feeb5b1e2851..60d6621fe583b 100644
--- a/airflow/models/baseoperator.py
+++ b/airflow/models/baseoperator.py
@@ -1572,6 +1572,7 @@ def get_serialized_fields(cls):
                     'subdag',
                     'ui_color',
                     'ui_fgcolor',
+                    'template_ext',
                     'template_fields',
                     'template_fields_renderers',
                 }
diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py
index 4978736d15d57..35561aeb1bb9b 100644
--- a/airflow/serialization/serialized_objects.py
+++ b/airflow/serialization/serialized_objects.py
@@ -434,6 +434,7 @@ def __init__(self, *args, **kwargs):
         # Move class attributes into object attributes.
         self.ui_color = BaseOperator.ui_color
         self.ui_fgcolor = BaseOperator.ui_fgcolor
+        self.template_ext = BaseOperator.template_ext
         self.template_fields = BaseOperator.template_fields
         self.operator_extra_links = BaseOperator.operator_extra_links
 
diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py
index 450bed93f0b19..cb68722f6a171 100644
--- a/tests/serialization/test_dag_serialization.py
+++ b/tests/serialization/test_dag_serialization.py
@@ -99,6 +99,7 @@
                 "_outlets": [],
                 "ui_color": "#f0ede4",
                 "ui_fgcolor": "#000",
+                "template_ext": ['.sh', '.bash'],
                 "template_fields": ['bash_command', 'env'],
                 "template_fields_renderers": {'bash_command': 'bash', 'env': 'json'},
                 "bash_command": "echo {{ task.task_id }}",
@@ -130,6 +131,7 @@
                 "_operator_extra_links": [{"tests.test_utils.mock_operators.CustomOpLink": {}}],
                 "ui_color": "#fff",
                 "ui_fgcolor": "#000",
+                "template_ext": [],
                 "template_fields": ['bash_command'],
                 "template_fields_renderers": {},
                 "_task_type": "CustomOperator",
@@ -447,6 +449,7 @@ def validate_deserialized_task(
             # Type is excluded, so don't check it
             '_log',
             # List vs tuple. Check separately
+            'template_ext',
             'template_fields',
             # We store the string, real dag has the actual code
             'on_failure_callback',
@@ -457,6 +460,8 @@ def validate_deserialized_task(
         }
 
         assert serialized_task.task_type == task.task_type
+
+        assert set(serialized_task.template_ext) == set(task.template_ext)
         assert set(serialized_task.template_fields) == set(task.template_fields)
 
         assert serialized_task.upstream_task_ids == task.upstream_task_ids

From 40f3cfe90aed0e8cf601dc8cd17028ceb967cfa4 Mon Sep 17 00:00:00 2001
From: fatmumuhomer <avoicelikerunningwater@gmail.com>
Date: Thu, 2 Sep 2021 22:50:40 -0400
Subject: [PATCH 42/42] Remove commented out line

---
 tests/providers/ssh/operators/test_ssh.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/providers/ssh/operators/test_ssh.py b/tests/providers/ssh/operators/test_ssh.py
index 9be7d0b31ff7f..4c7da74caba9e 100644
--- a/tests/providers/ssh/operators/test_ssh.py
+++ b/tests/providers/ssh/operators/test_ssh.py
@@ -276,7 +276,6 @@ def test_get_pyt_set_correctly(self, command, get_pty_in, get_pty_out):
             task_id="test",
             ssh_hook=self.hook,
             command=command,
-            # timeout=TIMEOUT,
             conn_timeout=TIMEOUT,
             cmd_timeout=TIMEOUT,
             get_pty=get_pty_in,
