pyinfra 2.9.2__py2.py3-none-any.whl → 3.0b1__py2.py3-none-any.whl

pyinfra/connectors/dockerssh.py

@@ -1,313 +1,291 @@
-"""
-**Note**: this connector is in beta!
-
-The ``@dockerssh`` connector allows you to run commands on Docker containers on a remote machine.
-
-.. code:: shell
-
-    # A Docker base image must be provided
-    pyinfra @dockerssh/remotehost:alpine:3.8 ...
-
-    # pyinfra can run on multiple Docker images in parallel
-    pyinfra @dockerssh/remotehost:alpine:3.8,@dockerssh/remotehost:ubuntu:bionic ...
-"""
-
 import os
 from tempfile import mkstemp
 from typing import TYPE_CHECKING
 
 import click
+from typing_extensions import Unpack
 
 from pyinfra import logger
 from pyinfra.api import QuoteString, StringCommand
-from pyinfra.api.connectors import BaseConnectorMeta
 from pyinfra.api.exceptions import ConnectError, InventoryError, PyinfraError
 from pyinfra.api.util import get_file_io, memoize
 from pyinfra.progress import progress_spinner
 
-from . import ssh
-from .util import make_unix_command_for_host
+from .base import BaseConnector
+from .ssh import SSHConnector
+from .util import extract_control_arguments, make_unix_command_for_host
 
 if TYPE_CHECKING:
+    from pyinfra.api.arguments import ConnectorArguments
     from pyinfra.api.host import Host
     from pyinfra.api.state import State
 
 
-class Meta(BaseConnectorMeta):
-    handles_execution = True
-
-
-def remote_remove(
-    state: "State", host: "Host", filename, print_output: bool = False, print_input: bool = False
-):
-    """
-    Deletes a file on a remote machine over ssh.
-    """
-    remove_status, _, remove_stderr = ssh.run_shell_command(
-        state,
-        host,
-        "rm -f {0}".format(filename),
-        print_output=print_output,
-        print_input=print_input,
-    )
-
-    if not remove_status:
-        raise IOError("\n".join(remove_stderr))
-
-
 @memoize
 def show_warning():
     logger.warning("The @dockerssh connector is in beta!")
 
 
-def make_names_data(host_image_str):
-    try:
-        hostname, image = host_image_str.split(":", 1)
-    except (AttributeError, ValueError):  # failure to parse the host_image_str
-        raise InventoryError("No ssh host or docker base image provided!")
+class DockerSSHConnector(BaseConnector):
+    """
+    **Note**: this connector is in beta!
 
-    if not image:
-        raise InventoryError("No docker base image provided!")
+    The ``@dockerssh`` connector allows you to run commands on Docker containers \
+    on a remote machine.
 
-    show_warning()
+    .. code:: shell
 
-    yield (
-        "@dockerssh/{0}:{1}".format(hostname, image),
-        {"ssh_hostname": hostname, "docker_image": image},
-        ["@dockerssh"],
-    )
+        # A Docker base image must be provided
+        pyinfra @dockerssh/remotehost:alpine:3.8 ...
 
+        # pyinfra can run on multiple Docker images in parallel
+        pyinfra @dockerssh/remotehost:alpine:3.8,@dockerssh/remotehost:ubuntu:bionic ...
+    """
 
-def connect(state: "State", host: "Host"):
-    if not host.connection:
-        host.connection = ssh.connect(state, host)
+    handles_execution = True
 
-    if "docker_container_id" in host.host_data:  # user can provide a docker_container_id
-        return host.connection
+    ssh: SSHConnector
 
-    try:
-        with progress_spinner({"docker run"}):
-            # last line is the container ID
-            status, stdout, stderr = ssh.run_shell_command(
-                state,
-                host,
-                "docker run -d {0} tail -f /dev/null".format(host.data.docker_image),
-            )
-            if not status:
-                raise IOError("\n".join(stderr))
-            container_id = stdout[-1]
+    def __init__(self, state: "State", host: "Host"):
+        super().__init__(state, host)
+        self.ssh = SSHConnector(state, host)
 
-    except PyinfraError as e:
-        host.connection = None  # fail connection
-        raise ConnectError(e.args[0])
+    @staticmethod
+    def make_names_data(name):
+        try:
+            hostname, image = name.split(":", 1)
+        except (AttributeError, ValueError):  # failure to parse the name
+            raise InventoryError("No ssh host or docker base image provided!")
 
-    host.host_data["docker_container_id"] = container_id
-    return host.connection
+        if not image:
+            raise InventoryError("No docker base image provided!")
 
+        show_warning()
 
-def disconnect(state: "State", host: "Host"):
-    container_id = host.host_data["docker_container_id"][:12]
+        yield (
+            "@dockerssh/{0}:{1}".format(hostname, image),
+            {"ssh_hostname": hostname, "docker_image": image},
+            ["@dockerssh"],
+        )
 
-    with progress_spinner({"docker commit"}):
-        image_id = ssh.run_shell_command(state, host, "docker commit {0}".format(container_id))[1][
-            -1
-        ][
-            7:19
-        ]  # last line is the image ID, get sha256:[XXXXXXXXXX]...
+    def connect(self) -> None:
+        self.ssh.connect()
+
+        if "docker_container_id" in self.host.host_data:  # user can provide a docker_container_id
+            return
+
+        try:
+            with progress_spinner({"docker run"}):
+                # last line is the container ID
+                status, output = self.ssh.run_shell_command(
+                    StringCommand(
+                        "docker",
+                        "run",
+                        "-d",
+                        self.host.data.docker_image,
+                        "tail",
+                        "-f",
+                        "/dev/null",
+                    ),
+                )
+                if not status:
+                    raise IOError(output.stderr)
+                container_id = output.stdout_lines[-1]
+
+        except PyinfraError as e:
+            raise ConnectError(e.args[0])
+
+        self.host.host_data["docker_container_id"] = container_id
+
+    def disconnect(self) -> None:
+        container_id = self.host.host_data["docker_container_id"][:12]
+
+        with progress_spinner({"docker commit"}):
+            _, output = self.ssh.run_shell_command(StringCommand("docker", "commit", container_id))
+
+            # Last line is the image ID, get sha256:[XXXXXXXXXX]...
+            image_id = output.stdout_lines[-1][7:19]
+
+        with progress_spinner({"docker rm"}):
+            self.ssh.run_shell_command(
+                StringCommand("docker", "rm", "-f", container_id),
+            )
 
-    with progress_spinner({"docker rm"}):
-        ssh.run_shell_command(
-            state,
-            host,
-            "docker rm -f {0}".format(container_id),
+        logger.info(
+            "{0}docker build complete, image ID: {1}".format(
+                self.host.print_prefix,
+                click.style(image_id, bold=True),
+            ),
         )
 
-    logger.info(
-        "{0}docker build complete, image ID: {1}".format(
-            host.print_prefix,
-            click.style(image_id, bold=True),
-        ),
-    )
-
-
-def run_shell_command(
-    state: "State",
-    host: "Host",
-    command,
-    get_pty: bool = False,
-    timeout=None,
-    stdin=None,
-    success_exit_codes=None,
-    print_output: bool = False,
-    print_input: bool = False,
-    return_combined_output=False,
-    **command_kwargs,
-):
-    container_id = host.host_data["docker_container_id"]
-
-    # Don't sudo/su in Docker - is this the right thing to do? Makes deploys that
-    # target SSH systems work w/Docker out of the box (ie most docker commands
-    # are run as root).
-    for key in ("sudo", "su_user"):
-        command_kwargs.pop(key, None)
-
-    command = make_unix_command_for_host(state, host, command, **command_kwargs)
-    command = QuoteString(command)
-
-    docker_flags = "-it" if get_pty else "-i"
-    docker_command = StringCommand(
-        "docker",
-        "exec",
-        docker_flags,
-        container_id,
-        "sh",
-        "-c",
+    def run_shell_command(
+        self,
         command,
-    )
-
-    return ssh.run_shell_command(
-        state,
-        host,
-        docker_command,
-        timeout=timeout,
-        stdin=stdin,
-        success_exit_codes=success_exit_codes,
-        print_output=print_output,
-        print_input=print_input,
-        return_combined_output=return_combined_output,
-    )
-
-
-def put_file(
-    state: "State",
-    host: "Host",
-    filename_or_io,
-    remote_filename,
-    remote_temp_filename=None,
-    print_output: bool = False,
-    print_input: bool = False,
-    **kwargs,  # ignored (sudo/etc)
-):
-    """
-    Upload a file/IO object to the target Docker container by copying it to a
-    temporary location and then uploading it into the container using ``docker cp``.
-    """
-
-    fd, local_temp_filename = mkstemp()
-    remote_temp_filename = remote_temp_filename or state.get_temp_filename(local_temp_filename)
-
-    # Load our file or IO object and write it to the temporary file
-    with get_file_io(filename_or_io) as file_io:
-        with open(local_temp_filename, "wb") as temp_f:
-            data = file_io.read()
-
-            if isinstance(data, str):
-                data = data.encode()
-
-            temp_f.write(data)
-
-    # upload file to remote server
-    ssh_status = ssh.put_file(state, host, local_temp_filename, remote_temp_filename)
-    if not ssh_status:
-        raise IOError("Failed to copy file over ssh")
-
-    try:
-        docker_id = host.host_data["docker_container_id"]
-        docker_command = "docker cp {0} {1}:{2}".format(
-            remote_temp_filename,
-            docker_id,
-            remote_filename,
+        print_output: bool = False,
+        print_input: bool = False,
+        **arguments: Unpack["ConnectorArguments"],
+    ):
+        local_arguments = extract_control_arguments(arguments)
+
+        container_id = self.host.host_data["docker_container_id"]
+
+        command = make_unix_command_for_host(self.state, self.host, command, **arguments)
+        command = QuoteString(command)
+
+        docker_flags = "-it" if local_arguments.get("_get_pty") else "-i"
+        docker_command = StringCommand(
+            "docker",
+            "exec",
+            docker_flags,
+            container_id,
+            "sh",
+            "-c",
+            command,
         )
 
-        status, _, stderr = ssh.run_shell_command(
-            state,
-            host,
+        return self.ssh.run_shell_command(
             docker_command,
             print_output=print_output,
             print_input=print_input,
+            **local_arguments,
         )
-    finally:
-        os.close(fd)
-        os.remove(local_temp_filename)
-        remote_remove(
-            state,
-            host,
-            local_temp_filename,
-            print_output=print_output,
-            print_input=print_input,
+
+    def put_file(
+        self,
+        filename_or_io,
+        remote_filename,
+        remote_temp_filename=None,
+        print_output: bool = False,
+        print_input: bool = False,
+        **kwargs,  # ignored (sudo/etc)
+    ):
+        """
+        Upload a file/IO object to the target Docker container by copying it to a
+        temporary location and then uploading it into the container using ``docker cp``.
+        """
+
+        fd, local_temp_filename = mkstemp()
+        remote_temp_filename = remote_temp_filename or self.host.get_temp_filename(
+            local_temp_filename
         )
 
-    if not status:
-        raise IOError("\n".join(stderr))
+        # Load our file or IO object and write it to the temporary file
+        with get_file_io(filename_or_io) as file_io:
+            with open(local_temp_filename, "wb") as temp_f:
+                data = file_io.read()
 
-    if print_output:
-        click.echo(
-            "{0}file uploaded to container: {1}".format(
-                host.print_prefix,
-                remote_filename,
-            ),
-            err=True,
-        )
+                if isinstance(data, str):
+                    data = data.encode()
 
-    return status
+                temp_f.write(data)
 
+        # upload file to remote server
+        ssh_status = self.ssh.put_file(local_temp_filename, remote_temp_filename)
+        if not ssh_status:
+            raise IOError("Failed to copy file over ssh")
 
-def get_file(
-    state: "State",
-    host: "Host",
-    remote_filename,
-    filename_or_io,
-    remote_temp_filename=None,
-    print_output: bool = False,
-    print_input: bool = False,
-    **kwargs,  # ignored (sudo/etc)
-):
-    """
-    Download a file from the target Docker container by copying it to a temporary
-    location and then reading that into our final file/IO object.
-    """
+        try:
+            docker_id = self.host.host_data["docker_container_id"]
+            docker_command = StringCommand(
+                "docker",
+                "cp",
+                remote_temp_filename,
+                f"{docker_id}:{remote_filename}",
+            )
 
-    remote_temp_filename = remote_temp_filename or state.get_temp_filename(remote_filename)
+            status, output = self.ssh.run_shell_command(
+                docker_command,
+                print_output=print_output,
+                print_input=print_input,
+            )
+        finally:
+            os.close(fd)
+            os.remove(local_temp_filename)
+            self.remote_remove(
+                local_temp_filename,
+                print_output=print_output,
+                print_input=print_input,
+            )
 
-    try:
-        docker_id = host.host_data["docker_container_id"]
-        docker_command = "docker cp {0}:{1} {2}".format(
-            docker_id,
-            remote_filename,
-            remote_temp_filename,
-        )
+        if not status:
+            raise IOError(output.stderr)
 
-        status, _, stderr = ssh.run_shell_command(
-            state,
-            host,
-            docker_command,
-            print_output=print_output,
-            print_input=print_input,
-        )
+        if print_output:
+            click.echo(
+                "{0}file uploaded to container: {1}".format(
+                    self.host.print_prefix,
+                    remote_filename,
+                ),
+                err=True,
+            )
 
-        ssh_status = ssh.get_file(state, host, remote_temp_filename, filename_or_io)
-    finally:
-        remote_remove(
-            state,
-            host,
-            remote_temp_filename,
-            print_output=print_output,
-            print_input=print_input,
-        )
+        return status
+
+    def get_file(
+        self,
+        remote_filename,
+        filename_or_io,
+        remote_temp_filename=None,
+        print_output: bool = False,
+        print_input: bool = False,
+        **kwargs,  # ignored (sudo/etc)
+    ):
+        """
+        Download a file from the target Docker container by copying it to a temporary
+        location and then reading that into our final file/IO object.
+        """
+
+        remote_temp_filename = remote_temp_filename or self.host.get_temp_filename(remote_filename)
+
+        try:
+            docker_id = self.host.host_data["docker_container_id"]
+            docker_command = StringCommand(
+                "docker",
+                "cp",
+                f"{docker_id}:{remote_filename}",
+                remote_temp_filename,
+            )
 
-    if not ssh_status:
-        raise IOError("failed to copy file over ssh")
+            status, output = self.ssh.run_shell_command(
+                docker_command,
+                print_output=print_output,
+                print_input=print_input,
+            )
 
-    if not status:
-        raise IOError("\n".join(stderr))
+            ssh_status = self.ssh.get_file(remote_temp_filename, filename_or_io)
+        finally:
+            self.remote_remove(
+                remote_temp_filename,
+                print_output=print_output,
+                print_input=print_input,
+            )
 
-    if print_output:
-        click.echo(
-            "{0}file downloaded from container: {1}".format(
-                host.print_prefix,
-                remote_filename,
-            ),
-            err=True,
+        if not ssh_status:
+            raise IOError("failed to copy file over ssh")
+
+        if not status:
+            raise IOError(output.stderr)
+
+        if print_output:
+            click.echo(
+                "{0}file downloaded from container: {1}".format(
+                    self.host.print_prefix,
+                    remote_filename,
+                ),
+                err=True,
+            )
+
+        return status
+
+    def remote_remove(self, filename, print_output: bool = False, print_input: bool = False):
+        """
+        Deletes a file on a remote machine over ssh.
+        """
+        remove_status, output = self.ssh.run_shell_command(
+            StringCommand("rm", "-f", filename),
+            print_output=print_output,
+            print_input=print_input,
         )
 
-    return status
+        if not remove_status:
+            raise IOError(output.stderr)