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

pyinfra/connectors/docker.py

@@ -1,25 +1,4 @@
-"""
-The ``@docker`` connector allows you to build Docker images, or modify running
-Docker containers, using ``pyinfra``. You can pass either an image name or
-existing container ID:
-
-+ Image - will create a container from the image, execute operations and save
-    into a new image
-+ Existing container ID - will simply execute operations against the container,
-    leaving it up afterwards
-
-
-.. code:: shell
-
-    # A Docker base image must be provided
-    pyinfra @docker/alpine:3.8 ...
-
-    # pyinfra can run on multiple Docker images in parallel
-    pyinfra @docker/alpine:3.8,@docker/ubuntu:bionic ...
-
-    # Execute against a running container
-    pyinfra @docker/2beb8c15a1b1 ...
-"""
+from __future__ import annotations
 
 import json
 import os
@@ -27,51 +6,42 @@ from tempfile import mkstemp
 from typing import TYPE_CHECKING
 
 import click
+from typing_extensions import TypedDict, Unpack
 
 from pyinfra import local, 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
 from pyinfra.progress import progress_spinner
 
-from .local import run_shell_command as run_local_shell_command
-from .util import make_unix_command_for_host
+from .base import BaseConnector, DataMeta
+from .local import LocalConnector
+from .util import CommandOutput, 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: bool = True
-    keys_prefix: str = "docker"
-
-    class DataKeys:
-        identifier = "ID of container or image to target"
-        container_id = "ID of container to target, overrides ``docker_identifier``"
-
-
-DATA_KEYS = Meta.keys()
-
+class ConnectorData(TypedDict):
+    docker_identifier: str
 
-def make_names_data(identifier=None):
-    if not identifier:
-        raise InventoryError("No docker base ID provided!")
 
-    yield (
-        "@docker/{0}".format(identifier),
-        {DATA_KEYS.identifier: identifier},
-        ["@docker"],
-    )
+connector_data_meta: dict[str, DataMeta] = {
+    "docker_identifier": DataMeta("ID of container or image to start from"),
+}
 
 
-def _find_start_docker_container(container_id):
+def _find_start_docker_container(container_id) -> tuple[str, bool]:
     docker_info = local.shell("docker container inspect {0}".format(container_id))
+    assert isinstance(docker_info, str)
     docker_info = json.loads(docker_info)[0]
     if docker_info["State"]["Running"] is False:
         logger.info("Starting stopped container: {0}".format(container_id))
         local.shell("docker container start {0}".format(container_id))
+        return container_id, False
+    return container_id, True
 
 
 def _start_docker_image(image_name):
@@ -86,220 +56,233 @@ def _start_docker_image(image_name):
         raise ConnectError(e.args[0])
 
 
-def connect(state: "State", host: "Host"):
-    docker_container_id = host.data.get(DATA_KEYS.container_id)
-    if docker_container_id:  # user can provide a docker_container_id
-        host.connector_data["docker_container_no_disconnect"] = True
-        host.connector_data["docker_container_id"] = docker_container_id
-        return True
+class DockerConnector(BaseConnector):
+    """
+    The docker connector allows you to build Docker images or modify running
+    Docker containers. You can pass either an image name or existing container ID:
 
-    docker_identifier = getattr(host.data, DATA_KEYS.identifier)
-    with progress_spinner({"prepare docker container"}):
-        try:
-            # Check if the provided @docker/X is an existing container ID
-            _find_start_docker_container(docker_identifier)
-        except PyinfraError:
-            container_id = _start_docker_image(docker_identifier)
-        else:
-            container_id = docker_identifier
-            host.connector_data["docker_container_no_disconnect"] = True
+    + Image - will create a new container from the image, execute operations \
+        against it, save into a new Docker image and remove the container
+    + Existing container ID - will execute operations against the running \
+        container, leaving it running
 
-    host.connector_data["docker_container_id"] = container_id
-    return True
+    .. code:: shell
 
+        # A Docker base image must be provided
+        pyinfra @docker/alpine:3.8 ...
 
-def disconnect(state, host):
-    container_id = host.connector_data["docker_container_id"]
+        # pyinfra can run on multiple Docker images in parallel
+        pyinfra @docker/alpine:3.8,@docker/ubuntu:bionic ...
 
-    if host.connector_data.get("docker_container_no_disconnect"):
-        logger.info(
-            "{0}docker build complete, container left running: {1}".format(
-                host.print_prefix,
-                click.style(container_id, bold=True),
-            ),
-        )
-        return
+        # Execute against a running container
+        pyinfra @docker/2beb8c15a1b1 ...
+    """
 
-    with progress_spinner({"docker commit"}):
-        image_id = local.shell("docker commit {0}".format(container_id), splitlines=True)[-1][
-            7:19
-        ]  # last line is the image ID, get sha256:[XXXXXXXXXX]...
+    handles_execution = True
 
-    with progress_spinner({"docker rm"}):
-        local.shell(
-            "docker rm -f {0}".format(container_id),
-        )
+    data_cls = ConnectorData
+    data_meta = connector_data_meta
+    data: ConnectorData
 
-    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=False,
-    timeout=None,
-    stdin=None,
-    success_exit_codes=None,
-    print_output=False,
-    print_input=False,
-    return_combined_output=False,
-    **command_kwargs,
-):
-    container_id = host.connector_data["docker_container_id"]
-
-    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",
-        command,
-    )
-
-    return run_local_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,  # ignored
-    print_output=False,
-    print_input=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``.
-    """
+    local: LocalConnector
 
-    fd, temp_filename = mkstemp()
+    container_id: str
+    no_stop: bool = False
 
-    try:
-        # 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(temp_filename, "wb") as temp_f:
-                data = file_io.read()
+    def __init__(self, state: "State", host: "Host"):
+        super().__init__(state, host)
+        self.local = LocalConnector(state, host)
 
-                if isinstance(data, str):
-                    data = data.encode()
+    @staticmethod
+    def make_names_data(name=None):
+        if not name:
+            raise InventoryError("No docker base ID provided!")
 
-                temp_f.write(data)
-
-        docker_id = host.connector_data["docker_container_id"]
-        docker_command = "docker cp {0} {1}:{2}".format(
-            temp_filename,
-            docker_id,
-            remote_filename,
+        yield (
+            "@docker/{0}".format(name),
+            {"docker_identifier": name},
+            ["@docker"],
         )
 
-        status, _, stderr = run_local_shell_command(
-            state,
-            host,
-            docker_command,
-            print_output=print_output,
-            print_input=print_input,
-        )
-    finally:
-        os.close(fd)
-        os.remove(temp_filename)
-
-    if not status:
-        raise IOError("\n".join(stderr))
-
-    if print_output:
-        click.echo(
-            "{0}file uploaded to container: {1}".format(
-                host.print_prefix,
-                remote_filename,
+    def connect(self) -> None:
+        self.local.connect()
+
+        docker_identifier = self.data["docker_identifier"]
+        with progress_spinner({"prepare docker container"}):
+            try:
+                self.container_id, was_running = _find_start_docker_container(docker_identifier)
+                if was_running:
+                    self.no_stop = True
+            except PyinfraError:
+                self.container_id = _start_docker_image(docker_identifier)
+
+    def disconnect(self):
+        container_id = self.container_id
+
+        if self.no_stop:
+            logger.info(
+                "{0}docker build complete, container left running: {1}".format(
+                    self.host.print_prefix,
+                    click.style(container_id, bold=True),
+                ),
+            )
+            return
+
+        with progress_spinner({"docker commit"}):
+            image_id = local.shell("docker commit {0}".format(container_id), splitlines=True)[-1][
+                7:19
+            ]  # last line is the image ID, get sha256:[XXXXXXXXXX]...
+
+        with progress_spinner({"docker rm"}):
+            local.shell(
+                "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),
             ),
-            err=True,
         )
 
-    return status
-
-
-def get_file(
-    state: "State",
-    host: "Host",
-    remote_filename,
-    filename_or_io,
-    remote_temp_filename=None,  # ignored
-    print_output=False,
-    print_input=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.
-    """
-
-    fd, temp_filename = mkstemp()
-
-    try:
-        docker_id = host.connector_data["docker_container_id"]
-        docker_command = "docker cp {0}:{1} {2}".format(
-            docker_id,
-            remote_filename,
-            temp_filename,
+    def run_shell_command(
+        self,
+        command: StringCommand,
+        print_output: bool = False,
+        print_input: bool = False,
+        **arguments: Unpack["ConnectorArguments"],
+    ) -> tuple[bool, CommandOutput]:
+        local_arguments = extract_control_arguments(arguments)
+
+        container_id = self.container_id
+
+        command = make_unix_command_for_host(self.state, self.host, command, **arguments)
+        command = StringCommand(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 = run_local_shell_command(
-            state,
-            host,
+        return self.local.run_shell_command(
             docker_command,
             print_output=print_output,
             print_input=print_input,
+            **local_arguments,
         )
 
-        # Load the temporary file and write it to our file or IO object
-        with open(temp_filename, encoding="utf-8") as temp_f:
-            with get_file_io(filename_or_io, "wb") as file_io:
-                data = temp_f.read()
-                data_bytes: bytes
-
-                if isinstance(data, str):
-                    data_bytes = data.encode()
-                else:
-                    data_bytes = data
-
-                file_io.write(data_bytes)
-    finally:
-        os.close(fd)
-        os.remove(temp_filename)
-
-    if not status:
-        raise IOError("\n".join(stderr))
-
-    if print_output:
-        click.echo(
-            "{0}file downloaded from container: {1}".format(
-                host.print_prefix,
-                remote_filename,
-            ),
-            err=True,
-        )
+    def put_file(
+        self,
+        filename_or_io,
+        remote_filename,
+        remote_temp_filename=None,  # ignored
+        print_output=False,
+        print_input=False,
+        **kwargs,  # ignored (sudo/etc)
+    ) -> bool:
+        """
+        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, temp_filename = mkstemp()
+
+        try:
+            # 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(temp_filename, "wb") as temp_f:
+                    data = file_io.read()
+
+                    if isinstance(data, str):
+                        data = data.encode()
+
+                    temp_f.write(data)
+
+            docker_command = StringCommand(
+                "docker",
+                "cp",
+                temp_filename,
+                f"{self.container_id}:{remote_filename}",
+            )
+
+            status, output = self.local.run_shell_command(
+                docker_command,
+                print_output=print_output,
+                print_input=print_input,
+            )
+        finally:
+            os.close(fd)
+            os.remove(temp_filename)
+
+        if not status:
+            raise IOError(output.stderr)
+
+        if print_output:
+            click.echo(
+                "{0}file uploaded to container: {1}".format(
+                    self.host.print_prefix,
+                    remote_filename,
+                ),
+                err=True,
+            )
+
+        return status
+
+    def get_file(
+        self,
+        remote_filename,
+        filename_or_io,
+        remote_temp_filename=None,  # ignored
+        print_output=False,
+        print_input=False,
+        **kwargs,  # ignored (sudo/etc)
+    ) -> bool:
+        """
+        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.
+        """
+
+        fd, temp_filename = mkstemp()
 
-    return status
+        try:
+            docker_command = StringCommand(
+                "docker",
+                "cp",
+                f"{self.container_id}:{remote_filename}",
+                temp_filename,
+            )
+
+            status, output = self.local.run_shell_command(
+                docker_command,
+                print_output=print_output,
+                print_input=print_input,
+            )
+
+            # Load the temporary file and write it to our file or IO object
+            with open(temp_filename, "rb") as temp_f:
+                with get_file_io(filename_or_io, "wb") as file_io:
+                    data = temp_f.read()
+                    file_io.write(data)
+        finally:
+            os.close(fd)
+            os.remove(temp_filename)
+
+        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