ssh_hook — Airflow Documentation

# Source code for ssh_hook

    # -*- coding: utf-8 -*-
    #
    # Copyright 2012-2015 Spotify AB
    # Ported to Airflow by Bolke de Bruin
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    # http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    #
    # This is a port of Luigi's ssh implementation. All credits go there.
    import subprocess
    from contextlib import contextmanager

    from airflow.hooks.base_hook import BaseHook
    from airflow.exceptions import AirflowException

    import logging

    class SSHHook(BaseHook):
        """
        Light-weight remote execution library and utilities.

        Using this hook (which is just a convenience wrapper for subprocess),
        is created to let you stream data from a remotely stored file.

        As a bonus, :class:`SSHHook` also provides a really cool feature that let's you
        set up ssh tunnels super easily using a python context manager (there is an example
        in the integration part of unittests).

        :param key_file: Typically the SSHHook uses the keys that are used by the user
            airflow is running under. This sets the behavior to use another file instead.
        :type key_file: str
        :param connect_timeout: sets the connection timeout for this connection.
        :type connect_timeout: int
        :param no_host_key_check: whether to check to host key. If True host keys will not
            be checked, but are also not stored in the current users's known_hosts file.
        :type no_host_key_check: bool
        :param tty: allocate a tty.
        :type tty: bool
        :param sshpass: Use to non-interactively perform password authentication by using
            sshpass.
        :type sshpass: bool
        """
        def __init__(self, conn_id='ssh_default'):
            conn = self.get_connection(conn_id)
            self.key_file = conn.extra_dejson.get('key_file', None)
            self.connect_timeout = conn.extra_dejson.get('connect_timeout', None)
            self.tcp_keepalive = conn.extra_dejson.get('tcp_keepalive', False)
            self.server_alive_interval = conn.extra_dejson.get('server_alive_interval', 60)
            self.no_host_key_check = conn.extra_dejson.get('no_host_key_check', False)
            self.tty = conn.extra_dejson.get('tty', False)
            self.sshpass = conn.extra_dejson.get('sshpass', False)
            self.conn = conn

        def get_conn(self):
            pass

        def _host_ref(self):
            if self.conn.login:
                return "{0}@{1}".format(self.conn.login, self.conn.host)
            else:
                return self.conn.host

        def _prepare_command(self, cmd):
            connection_cmd = ["ssh", self._host_ref(), "-o", "ControlMaster=no"]
            if self.sshpass:
                connection_cmd = ["sshpass", "-e"] + connection_cmd
            else:
                connection_cmd += ["-o", "BatchMode=yes"]  # no password prompts

            if self.conn.port:
                connection_cmd += ["-p", str(self.conn.port)]

            if self.connect_timeout:
                connection_cmd += ["-o", "ConnectionTimeout={}".format(self.connect_timeout)]

            if self.tcp_keepalive:
                connection_cmd += ["-o", "TCPKeepAlive=yes"]
                connection_cmd += ["-o", "ServerAliveInterval={}".format(self.server_alive_interval)]

            if self.no_host_key_check:
                connection_cmd += ["-o", "UserKnownHostsFile=/dev/null",
                                   "-o", "StrictHostKeyChecking=no"]

            if self.key_file:
                connection_cmd += ["-i", self.key_file]

            if self.tty:
                connection_cmd += ["-t"]

            connection_cmd += cmd
            logging.debug("SSH cmd: {} ".format(connection_cmd))

            return connection_cmd

        def Popen(self, cmd, **kwargs):
            """
            Remote Popen

            :param cmd: command to remotely execute
            :param kwargs: extra arguments to Popen (see subprocess.Popen)
            :return: handle to subprocess
            """
            prefixed_cmd = self._prepare_command(cmd)
            return subprocess.Popen(prefixed_cmd, **kwargs)

        def check_output(self, cmd):
            """
            Executes a remote command and returns the stdout a remote process.
            Simplified version of Popen when you only want the output as a string and detect any errors.

            :param cmd: command to remotely execute
            :return: stdout
            """
            p = self.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            output, stderr = p.communicate()

            if p.returncode != 0:
                # I like this better: RemoteCalledProcessError(p.returncode, cmd, self.host, output=output)
                raise AirflowException("Cannot execute {} on {}. Error code is: {}. Output: {}, Stderr: {}".format(
                                       cmd, self.conn.host, p.returncode, output, stderr))

            return output

        @contextmanager
        def tunnel(self, local_port, remote_port=None, remote_host="localhost"):
            """
            Creates a tunnel between two hosts. Like ssh -L <LOCAL_PORT>:host:<REMOTE_PORT>.
            Remember to close() the returned "tunnel" object in order to clean up
            after yourself when you are done with the tunnel.

            :param local_port:
            :type local_port: int
            :param remote_port:
            :type remote_port: int
            :param remote_host:
            :type remote_host: str
            :return:
            """
            tunnel_host = "{0}:{1}:{2}".format(local_port, remote_host, remote_port)
            proc = self.Popen(["-L", tunnel_host, "echo -n ready && cat"],
                              stdin=subprocess.PIPE, stdout=subprocess.PIPE)

            ready = proc.stdout.read(5)
            assert ready == b"ready", "Did not get 'ready' from remote"
            yield
            proc.communicate()
            assert proc.returncode == 0, "Tunnel process did unclean exit (returncode {}".format(proc.returncode)