pyinfra 3.4__py2.py3-none-any.whl → 3.5__py2.py3-none-any.whl

pyinfra/api/arguments.py

@@ -72,6 +72,11 @@ class ConnectorArguments(TypedDict, total=False):
     _get_pty: bool
     _stdin: Union[str, Iterable[str]]
 
+    # Retry arguments
+    _retries: int
+    _retry_delay: Union[int, float]
+    _retry_until: Optional[Callable[[dict], bool]]
+
 
 def generate_env(config: "Config", value: dict) -> dict:
     env = config.ENV.copy()
@@ -232,11 +237,28 @@ def all_global_arguments() -> List[tuple[str, Type]]:
     return list(get_type_hints(AllArguments).items())
 
 
+# Create a dictionary for retry arguments
+retry_argument_meta: dict[str, ArgumentMeta] = {
+    "_retries": ArgumentMeta(
+        "Number of times to retry failed operations.",
+        default=lambda config: config.RETRY,
+    ),
+    "_retry_delay": ArgumentMeta(
+        "Delay in seconds between retry attempts.",
+        default=lambda config: config.RETRY_DELAY,
+    ),
+    "_retry_until": ArgumentMeta(
+        "Callable taking output data that returns True to continue retrying.",
+        default=lambda config: None,
+    ),
+}
+
 all_argument_meta: dict[str, ArgumentMeta] = {
     **auth_argument_meta,
     **shell_argument_meta,
     **meta_argument_meta,
     **execution_argument_meta,
+    **retry_argument_meta,  # Add retry arguments
 }
 
 EXECUTION_KWARG_KEYS = list(ExecutionArguments.__annotations__.keys())
@@ -286,6 +308,45 @@ __argument_docs__ = {
     ),
     "Operation meta & callbacks": (meta_argument_meta, "", ""),
     "Execution strategy": (execution_argument_meta, "", ""),
+    "Retry behavior": (
+        retry_argument_meta,
+        """
+        Retry arguments allow you to automatically retry operations that fail. You can specify
+        how many times to retry, the delay between retries, and optionally a condition
+        function to determine when to stop retrying.
+        """,
+        """
+        .. code:: python
+
+            # Retry a command up to 3 times with the default 5 second delay
+            server.shell(
+                name="Run flaky command with retries",
+                commands=["flaky_command"],
+                _retries=3,
+            )
+            # Retry with a custom delay
+            server.shell(
+                name="Run flaky command with custom delay",
+                commands=["flaky_command"],
+                _retries=2,
+                _retry_delay=10,  # 10 second delay between retries
+            )
+            # Retry with a custom condition
+            def retry_on_specific_error(output_data):
+                # Retry if stderr contains "temporary failure"
+                for line in output_data["stderr_lines"]:
+                    if "temporary failure" in line.lower():
+                        return True
+                return False
+
+            server.shell(
+                name="Run command with conditional retry",
+                commands=["flaky_command"],
+                _retries=5,
+                _retry_until=retry_on_specific_error,
+            )
+        """,
+    ),
 }
 
 
@@ -305,7 +366,8 @@ def pop_global_arguments(
 
     config = state.config
     if ctx_config.isset():
-        config = config
+        config = ctx_config.get()
+        assert config is not None
 
     cdkwargs = host.current_deploy_kwargs
     meta_kwargs: dict[str, Any] = cdkwargs or {}  # type: ignore[assignment]

pyinfra/api/config.py

@@ -53,6 +53,12 @@ class ConfigDefaults:
     IGNORE_ERRORS: bool = False
     # Shell to use to execute commands
     SHELL: str = "sh"
+    # Whether to display full diffs for files
+    DIFF: bool = False
+    # Number of times to retry failed operations
+    RETRY: int = 0
+    # Delay in seconds between retry attempts
+    RETRY_DELAY: int = 5
 
 
 config_defaults = {key: value for key, value in ConfigDefaults.__dict__.items() if key.isupper()}

pyinfra/api/connect.py

@@ -46,5 +46,22 @@ def connect_all(state: "State"):
 
 
 def disconnect_all(state: "State"):
-    for host in state.activated_hosts:  # only hosts we connected to please!
-        host.disconnect()  # normally a noop
+    """
+    Disconnect from all of the configured servers in parallel. Reads/writes state.inventory.
+
+    Args:
+        state (``pyinfra.api.State`` obj): the state containing an inventory to connect to
+    """
+    greenlet_to_host = {
+        state.pool.spawn(host.disconnect): host
+        for host in state.activated_hosts  # only hosts we connected to please!
+    }
+
+    with progress_spinner(greenlet_to_host.values()) as progress:
+        for greenlet in gevent.iwait(greenlet_to_host.keys()):
+            host = greenlet_to_host[greenlet]
+            progress(host)
+
+    for greenlet, host in greenlet_to_host.items():
+        # Raise any unexpected exception
+        greenlet.get()

pyinfra/api/operation.py

@@ -47,6 +47,9 @@ class OperationMeta:
     _commands: Optional[list[Any]] = None
     _maybe_is_change: Optional[bool] = None
     _success: Optional[bool] = None
+    _retry_attempts: int = 0
+    _max_retries: int = 0
+    _retry_succeeded: Optional[bool] = None
 
     def __init__(self, hash, is_change: Optional[bool]):
         self._hash = hash
@@ -59,9 +62,17 @@ class OperationMeta:
         """
 
         if self._commands is not None:
+            retry_info = ""
+            if self._retry_attempts > 0:
+                retry_result = "succeeded" if self._retry_succeeded else "failed"
+                retry_info = (
+                    f", retries={self._retry_attempts}/{self._max_retries} ({retry_result})"
+                )
+
             return (
                 "OperationMeta(executed=True, "
-                f"success={self.did_succeed()}, hash={self._hash}, commands={len(self._commands)})"
+                f"success={self.did_succeed()}, hash={self._hash}, "
+                f"commands={len(self._commands)}{retry_info})"
             )
         return (
             "OperationMeta(executed=False, "
@@ -74,12 +85,20 @@ class OperationMeta:
         success: bool,
         commands: list[Any],
         combined_output: "CommandOutput",
+        retry_attempts: int = 0,
+        max_retries: int = 0,
     ) -> None:
         if self.is_complete():
             raise RuntimeError("Cannot complete an already complete operation")
         self._success = success
         self._commands = commands
         self._combined_output = combined_output
+        self._retry_attempts = retry_attempts
+        self._max_retries = max_retries
+
+        # Determine if operation succeeded after retries
+        if retry_attempts > 0:
+            self._retry_succeeded = success
 
     def is_complete(self) -> bool:
         return self._success is not None
@@ -150,6 +169,40 @@ class OperationMeta:
     def stderr(self) -> str:
         return "\n".join(self.stderr_lines)
 
+    @property
+    def retry_attempts(self) -> int:
+        return self._retry_attempts
+
+    @property
+    def max_retries(self) -> int:
+        return self._max_retries
+
+    @property
+    def was_retried(self) -> bool:
+        """
+        Returns whether this operation was retried at least once.
+        """
+        return self._retry_attempts > 0
+
+    @property
+    def retry_succeeded(self) -> Optional[bool]:
+        """
+        Returns whether this operation succeeded after retries.
+        Returns None if the operation was not retried.
+        """
+        return self._retry_succeeded
+
+    def get_retry_info(self) -> dict[str, Any]:
+        """
+        Returns a dictionary with all retry-related information.
+        """
+        return {
+            "retry_attempts": self._retry_attempts,
+            "max_retries": self._max_retries,
+            "was_retried": self.was_retried,
+            "retry_succeeded": self._retry_succeeded,
+        }
+
 
 def add_op(state: State, op_func, *args, **kwargs):
     """

pyinfra/api/operations.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import time
 import traceback
 from itertools import product
 from socket import error as socket_error, timeout as timeout_error
@@ -66,6 +67,11 @@ def _run_host_op(state: "State", host: "Host", op_hash: str) -> Optional[bool]:
     continue_on_error = global_arguments["_continue_on_error"]
     timeout = global_arguments.get("_timeout", 0)
 
+    # Extract retry arguments
+    retries = global_arguments.get("_retries", 0)
+    retry_delay = global_arguments.get("_retry_delay", 5)
+    retry_until = global_arguments.get("_retry_until", None)
+
     executor_kwarg_keys = CONNECTOR_ARGUMENT_KEYS
     # See: https://github.com/python/mypy/issues/10371
     base_connector_arguments: ConnectorArguments = cast(
@@ -73,67 +79,114 @@ def _run_host_op(state: "State", host: "Host", op_hash: str) -> Optional[bool]:
         {key: global_arguments[key] for key in executor_kwarg_keys if key in global_arguments},  # type: ignore[literal-required] # noqa
     )
 
+    retry_attempt = 0
     did_error = False
     executed_commands = 0
-    commands = []
+    commands: list[PyinfraCommand] = []
     all_output_lines: list[OutputLine] = []
 
-    for command in op_data.command_generator():
-        commands.append(command)
-
-        status = False
-
-        connector_arguments = base_connector_arguments.copy()
-        connector_arguments.update(command.connector_arguments)
-
-        if not isinstance(command, PyinfraCommand):
-            raise TypeError("{0} is an invalid pyinfra command!".format(command))
-
-        if isinstance(command, FunctionCommand):
-            try:
-                status = command.execute(state, host, connector_arguments)
-            except Exception as e:
-                # Custom functions could do anything, so expect anything!
-                logger.warning(traceback.format_exc())
-                host.log_styled(
-                    f"Unexpected error in Python callback: {format_exception(e)}",
-                    fg="red",
-                    log_func=logger.warning,
-                )
-
-        elif isinstance(command, StringCommand):
-            output_lines = CommandOutput([])
-            try:
-                status, output_lines = command.execute(
-                    state,
-                    host,
-                    connector_arguments,
-                )
-            except (timeout_error, socket_error, SSHException) as e:
-                log_host_command_error(host, e, timeout=timeout)
-            all_output_lines.extend(output_lines)
-            # If we failed and have not already printed the stderr, print it
-            if status is False and not state.print_output:
-                print_host_combined_output(host, output_lines)
-
-        else:
-            try:
-                status = command.execute(state, host, connector_arguments)
-            except (timeout_error, socket_error, SSHException, IOError) as e:
-                log_host_command_error(host, e, timeout=timeout)
-
-        # Break the loop to trigger a failure
-        if status is False:
-            did_error = True
-            if continue_on_error is True:
-                continue
-            break
+    # Retry loop
+    while retry_attempt <= retries:
+        did_error = False
+        executed_commands = 0
+        commands = []
+        all_output_lines = []
+
+        for command in op_data.command_generator():
+            commands.append(command)
+            status = False
+            connector_arguments = base_connector_arguments.copy()
+            connector_arguments.update(command.connector_arguments)
+
+            if not isinstance(command, PyinfraCommand):
+                raise TypeError("{0} is an invalid pyinfra command!".format(command))
+
+            if isinstance(command, FunctionCommand):
+                try:
+                    status = command.execute(state, host, connector_arguments)
+                except Exception as e:
+                    # Custom functions could do anything, so expect anything!
+                    logger.warning(traceback.format_exc())
+                    host.log_styled(
+                        f"Unexpected error in Python callback: {format_exception(e)}",
+                        fg="red",
+                        log_func=logger.warning,
+                    )
+
+            elif isinstance(command, StringCommand):
+                output_lines = CommandOutput([])
+                try:
+                    status, output_lines = command.execute(
+                        state,
+                        host,
+                        connector_arguments,
+                    )
+                except (timeout_error, socket_error, SSHException) as e:
+                    log_host_command_error(host, e, timeout=timeout)
+                all_output_lines.extend(output_lines)
+                # If we failed and have not already printed the stderr, print it
+                if status is False and not state.print_output:
+                    print_host_combined_output(host, output_lines)
+
+            else:
+                try:
+                    status = command.execute(state, host, connector_arguments)
+                except (timeout_error, socket_error, SSHException, IOError) as e:
+                    log_host_command_error(host, e, timeout=timeout)
+
+            # Break the loop to trigger a failure
+            if status is False:
+                did_error = True
+                if continue_on_error is True:
+                    continue
+                break
+
+            executed_commands += 1
+
+        # Check if we should retry
+        should_retry = False
+        if retry_attempt < retries:
+            # Retry on error
+            if did_error:
+                should_retry = True
+            # Retry on condition if no error
+            elif retry_until and not did_error:
+                try:
+                    output_data = {
+                        "stdout_lines": [
+                            line.line for line in all_output_lines if line.buffer_name == "stdout"
+                        ],
+                        "stderr_lines": [
+                            line.line for line in all_output_lines if line.buffer_name == "stderr"
+                        ],
+                        "commands": [str(command) for command in commands],
+                        "executed_commands": executed_commands,
+                        "host": host.name,
+                        "operation": ", ".join(state.get_op_meta(op_hash).names) or "Operation",
+                    }
+                    should_retry = retry_until(output_data)
+                except Exception as e:
+                    host.log_styled(
+                        f"Error in retry_until function: {format_exception(e)}",
+                        fg="red",
+                        log_func=logger.warning,
+                    )
+
+        if should_retry:
+            retry_attempt += 1
+            state.trigger_callbacks("operation_host_retry", host, op_hash, retry_attempt, retries)
+            op_name = ", ".join(state.get_op_meta(op_hash).names) or "Operation"
+            host.log_styled(
+                f"Retrying {op_name} (attempt {retry_attempt}/{retries}) after {retry_delay}s...",
+                fg="yellow",
+                log_func=logger.info,
+            )
+            time.sleep(retry_delay)
+            continue
 
-        executed_commands += 1
+        break
 
     # Handle results
-    #
-
     op_success = return_status = not did_error
     host_results = state.get_results_for_host(host)
 
@@ -142,10 +195,13 @@ def _run_host_op(state: "State", host: "Host", op_hash: str) -> Optional[bool]:
         host_results.success_ops += 1
 
         _status_log = "Success" if executed_commands > 0 else "No changes"
+        if retry_attempt > 0:
+            _status_log = f"{_status_log} on retry {retry_attempt}"
+
         _click_log_status = click.style(_status_log, "green")
         logger.info("{0}{1}".format(host.print_prefix, _click_log_status))
 
-        state.trigger_callbacks("operation_host_success", host, op_hash)
+        state.trigger_callbacks("operation_host_success", host, op_hash, retry_attempt)
     else:
         if ignore_errors:
             host_results.ignored_error_ops += 1
@@ -156,6 +212,11 @@ def _run_host_op(state: "State", host: "Host", op_hash: str) -> Optional[bool]:
             host_results.partial_ops += 1
 
         _command_description = f"executed {executed_commands} commands"
+        if retry_attempt > 0:
+            _command_description = (
+                f"{_command_description} (failed after {retry_attempt}/{retries} retries)"
+            )
+
         log_error_or_warning(host, ignore_errors, _command_description, continue_on_error)
 
         # Ignored, op "completes" w/ ignored error
@@ -164,12 +225,14 @@ def _run_host_op(state: "State", host: "Host", op_hash: str) -> Optional[bool]:
             return_status = True
 
         # Unignored error -> False
-        state.trigger_callbacks("operation_host_error", host, op_hash)
+        state.trigger_callbacks("operation_host_error", host, op_hash, retry_attempt, retries)
 
     op_data.operation_meta.set_complete(
         op_success,
         commands,
         CommandOutput(all_output_lines),
+        retry_attempts=retry_attempt,
+        max_retries=retries,
     )
 
     return return_status

pyinfra/api/state.py

@@ -70,11 +70,19 @@ class BaseStateCallback:
         pass
 
     @staticmethod
-    def operation_host_success(state: "State", host: "Host", op_hash):
+    def operation_host_success(state: "State", host: "Host", op_hash, retry_count: int = 0):
         pass
 
     @staticmethod
-    def operation_host_error(state: "State", host: "Host", op_hash):
+    def operation_host_error(
+        state: "State", host: "Host", op_hash, retry_count: int = 0, max_retries: int = 0
+    ):
+        pass
+
+    @staticmethod
+    def operation_host_retry(
+        state: "State", host: "Host", op_hash, retry_num: int, max_retries: int
+    ):
         pass
 
     @staticmethod

pyinfra/connectors/scp/__init__.py

@@ -0,0 +1 @@
+from .client import SCPClient  # noqa: F401

pyinfra/connectors/scp/client.py

@@ -0,0 +1,204 @@
+from __future__ import annotations
+
+import ntpath
+import os
+from pathlib import PurePath
+from shlex import quote
+from socket import timeout as SocketTimeoutError
+from typing import IO, AnyStr
+
+from paramiko import Channel
+from paramiko.transport import Transport
+
+SCP_COMMAND = b"scp"
+
+
+# Unicode conversion functions; assume UTF-8
+def asbytes(s: bytes | str | PurePath) -> bytes:
+    """Turns unicode into bytes, if needed.
+
+    Assumes UTF-8.
+    """
+    if isinstance(s, bytes):
+        return s
+    elif isinstance(s, PurePath):
+        return bytes(s)
+    else:
+        return s.encode("utf-8")
+
+
+def asunicode(s: bytes | str) -> str:
+    """Turns bytes into unicode, if needed.
+
+    Uses UTF-8.
+    """
+    if isinstance(s, bytes):
+        return s.decode("utf-8", "replace")
+    else:
+        return s
+
+
+class SCPClient:
+    """
+    An scp1 implementation, compatible with openssh scp.
+    Raises SCPException for all transport related errors. Local filesystem
+    and OS errors pass through.
+
+    Main public methods are .putfo and .getfo
+    """
+
+    def __init__(
+        self,
+        transport: Transport,
+        buff_size: int = 16384,
+        socket_timeout: float = 10.0,
+    ):
+        self.transport = transport
+        self.buff_size = buff_size
+        self.socket_timeout = socket_timeout
+        self._channel: Channel | None = None
+        self.scp_command = SCP_COMMAND
+
+    @property
+    def channel(self) -> Channel:
+        """Return an open Channel, (re)opening if needed."""
+        if self._channel is None or self._channel.closed:
+            self._channel = self.transport.open_session()
+        return self._channel
+
+    def __enter__(self):
+        _ = self.channel  # triggers opening if not already open
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.close()
+
+    def putfo(
+        self,
+        fl: IO[AnyStr],
+        remote_path: str | bytes,
+        mode: str | bytes = "0644",
+        size: int | None = None,
+    ) -> None:
+        if size is None:
+            pos = fl.tell()
+            fl.seek(0, os.SEEK_END)  # Seek to end
+            size = fl.tell() - pos
+            fl.seek(pos, os.SEEK_SET)  # Seek back
+
+        self.channel.settimeout(self.socket_timeout)
+        self.channel.exec_command(
+            self.scp_command + b" -t " + asbytes(quote(asunicode(remote_path)))
+        )
+        self._recv_confirm()
+        self._send_file(fl, remote_path, mode, size=size)
+        self.close()
+
+    def getfo(self, remote_path: str, fl: IO):
+        remote_path_sanitized = quote(remote_path)
+        if os.name == "nt":
+            remote_file_name = ntpath.basename(remote_path_sanitized)
+        else:
+            remote_file_name = os.path.basename(remote_path_sanitized)
+        self.channel.settimeout(self.socket_timeout)
+        self.channel.exec_command(self.scp_command + b" -f " + asbytes(remote_path_sanitized))
+        self._recv_all(fl, remote_file_name)
+        self.close()
+        return fl
+
+    def close(self):
+        """close scp channel"""
+        if self._channel is not None:
+            self._channel.close()
+            self._channel = None
+
+    def _send_file(self, fl, name, mode, size):
+        basename = asbytes(os.path.basename(name))
+        # The protocol can't handle \n in the filename.
+        # Quote them as the control sequence \^J for now,
+        # which is how openssh handles it.
+        self.channel.sendall(
+            ("C%s %d " % (mode, size)).encode("ascii") + basename.replace(b"\n", b"\\^J") + b"\n"
+        )
+        self._recv_confirm()
+        file_pos = 0
+        buff_size = self.buff_size
+        chan = self.channel
+        while file_pos < size:
+            chan.sendall(fl.read(buff_size))
+            file_pos = fl.tell()
+        chan.sendall(b"\x00")
+        self._recv_confirm()
+
+    def _recv_confirm(self):
+        # read scp response
+        msg = b""
+        try:
+            msg = self.channel.recv(512)
+        except SocketTimeoutError:
+            raise SCPException("Timeout waiting for scp response")
+        # slice off the first byte, so this compare will work in py2 and py3
+        if msg and msg[0:1] == b"\x00":
+            return
+        elif msg and msg[0:1] == b"\x01":
+            raise SCPException(asunicode(msg[1:]))
+        elif self.channel.recv_stderr_ready():
+            msg = self.channel.recv_stderr(512)
+            raise SCPException(asunicode(msg))
+        elif not msg:
+            raise SCPException("No response from server")
+        else:
+            raise SCPException("Invalid response from server", msg)
+
+    def _recv_all(self, fh: IO, remote_file_name: str) -> None:
+        # loop over scp commands, and receive as necessary
+        commands = (b"C",)
+        while not self.channel.closed:
+            # wait for command as long as we're open
+            self.channel.sendall(b"\x00")
+            msg = self.channel.recv(1024)
+            if not msg:  # chan closed while receiving
+                break
+            assert msg[-1:] == b"\n"
+            msg = msg[:-1]
+            code = msg[0:1]
+            if code not in commands:
+                raise SCPException(asunicode(msg[1:]))
+            self._recv_file(msg[1:], fh, remote_file_name)
+
+    def _recv_file(self, cmd: bytes, fh: IO, remote_file_name: str) -> None:
+        chan = self.channel
+        parts = cmd.strip().split(b" ", 2)
+
+        try:
+            size = int(parts[1])
+        except (ValueError, IndexError):
+            chan.send(b"\x01")
+            chan.close()
+            raise SCPException("Bad file format")
+
+        buff_size = self.buff_size
+        pos = 0
+        chan.send(b"\x00")
+        try:
+            while pos < size:
+                # we have to make sure we don't read the final byte
+                if size - pos <= buff_size:
+                    buff_size = size - pos
+                data = chan.recv(buff_size)
+                if not data:
+                    raise SCPException("Underlying channel was closed")
+                fh.write(data)
+                pos = fh.tell()
+            msg = chan.recv(512)
+            if msg and msg[0:1] != b"\x00":
+                raise SCPException(asunicode(msg[1:]))
+        except SocketTimeoutError:
+            chan.close()
+            raise SCPException("Error receiving, socket.timeout")
+
+
+class SCPException(Exception):
+    """SCP exception class"""
+
+    pass