[GitHub] [airflow] baolsen commented on a change in pull request #17378: Refactor SSHOperator so a subclass can run many commands (#10874)

[GitBox <gi...@apache.org>]

baolsen commented on a change in pull request #17378:
URL: https://github.com/apache/airflow/pull/17378#discussion_r699137347



##########
File path: airflow/providers/ssh/operators/ssh.py
##########
@@ -80,103 +82,120 @@ def __init__(
         self.environment = environment
         self.get_pty = (self.command.startswith('sudo') or get_pty) if self.command else get_pty
 
-    def execute(self, context) -> Union[bytes, str, bool]:
+    def get_hook(self) -> SSHHook:
+        if self.ssh_conn_id:
+            if self.ssh_hook and isinstance(self.ssh_hook, SSHHook):
+                self.log.info("ssh_conn_id is ignored when ssh_hook is provided.")
+            else:
+                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)
+        if not self.ssh_hook:
+            raise AirflowException("Cannot operate without ssh_hook or ssh_conn_id.")
+
+        if self.remote_host is not None:
+            self.log.info(
+                "remote_host is provided explicitly. "
+                "It will replace the remote_host which was defined "
+                "in ssh_hook or predefined in connection of ssh_conn_id."
+            )
+            self.ssh_hook.remote_host = self.remote_host
+
+        return self.ssh_hook
+
+    def get_ssh_client(self) -> SSHClient:
+        # Remember to call close_ssh_client on this when done!
+        self.log.info('Creating ssh_client')
+        return self.get_hook().get_conn()
+
+    def exec_ssh_client_command(self, ssh_client: SSHClient, command: str) -> Tuple[int, bytes, bytes]:

Review comment:
       Hey @potiuk, thanks for the feedback.
   Hey @uranusjr, please review again and let me know :)




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org