baolsen commented on a change in pull request #17378:
URL: https://github.com/apache/airflow/pull/17378#discussion_r691183619
##########
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:
Build has passed. Was flaky CI before now.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
