pyinfra 0.11.dev3__py3-none-any.whl → 3.5.1__py3-none-any.whl

pyinfra/api/state.py

@@ -1,99 +1,214 @@
-from contextlib import contextmanager
+from __future__ import annotations
+
+from collections import defaultdict
+from dataclasses import dataclass
+from enum import IntEnum
+from graphlib import CycleError, TopologicalSorter
 from multiprocessing import cpu_count
-from uuid import uuid4
+from typing import TYPE_CHECKING, Callable, Iterator, Optional
 
 from gevent.pool import Pool
-from pkg_resources import parse_version
+from paramiko import PKey
 
-from pyinfra import __version__, logger
+from pyinfra import logger
 
 from .config import Config
 from .exceptions import PyinfraError
-from .util import get_caller_frameinfo, sha1_hash
+
+if TYPE_CHECKING:
+    from pyinfra.api.arguments import AllArguments
+    from pyinfra.api.command import PyinfraCommand
+    from pyinfra.api.host import Host
+    from pyinfra.api.inventory import Inventory
+    from pyinfra.api.operation import OperationMeta
+
 
 # Work out the max parallel we can achieve with the open files limit of the user/process,
 # take 10 for opening Python files and /3 for ~3 files per host during op runs.
 # See: https://github.com/Fizzadar/pyinfra/issues/44
 try:
-    from resource import getrlimit, RLIMIT_NOFILE
+    from resource import RLIMIT_NOFILE, getrlimit
+
     nofile_limit, _ = getrlimit(RLIMIT_NOFILE)
     MAX_PARALLEL = round((nofile_limit - 10) / 3)
 
 # Resource isn't available on Windows
 except ImportError:
     nofile_limit = 0
-    MAX_PARALLEL = None
+    MAX_PARALLEL = 100000
+
+
+class BaseStateCallback:
+    # Host callbacks
+    #
+
+    @staticmethod
+    def host_before_connect(state: "State", host: "Host"):
+        pass
+
+    @staticmethod
+    def host_connect(state: "State", host: "Host"):
+        pass
+
+    @staticmethod
+    def host_connect_error(state: "State", host: "Host", error):
+        pass
+
+    @staticmethod
+    def host_disconnect(state: "State", host: "Host"):
+        pass
+
+    # Operation callbacks
+    #
+
+    @staticmethod
+    def operation_start(state: "State", op_hash):
+        pass
+
+    @staticmethod
+    def operation_host_start(state: "State", host: "Host", op_hash):
+        pass
+
+    @staticmethod
+    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, 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
 
-def _make_name(current, new):
-    '''
-    Stops duplication between similarly named nested deploys, eg:
+    @staticmethod
+    def operation_end(state: "State", op_hash):
+        pass
 
-    Turn:
-        Deploy Kubernetes master/Configure Kubernetes
-    Into:
-        Deploy Kubernetes master/Configure
-    '''
 
-    current_tokens = current.split()
-    new_tokens = new.split()
+class StateStage(IntEnum):
+    # Setup - collect inventory & data
+    Setup = 1
+    # Connect - connect to the inventory
+    Connect = 2
+    # Prepare - detect operation changes
+    Prepare = 3
+    # Execute - execute operations
+    Execute = 4
+    # Disconnect - disconnect from the inventory
+    Disconnect = 5
 
-    new = ' '.join(
-        new_token for new_token in new_tokens
-        if new_token not in current_tokens
-    )
 
-    return '/'.join((current, new))
+class StateOperationMeta:
+    names: set[str]
+    args: list[str]
+    op_order: tuple[int, ...]
+    global_arguments: "AllArguments"
 
+    def __init__(self, op_order: tuple[int, ...]):
+        self.op_order = op_order
+        self.names = set()
+        self.args = []
+        self.global_arguments = {}  # type: ignore
 
-class State(object):
-    '''
+
+@dataclass
+class StateOperationHostData:
+    command_generator: Callable[[], Iterator["PyinfraCommand"]]
+    global_arguments: "AllArguments"
+    operation_meta: "OperationMeta"
+    parent_op_hash: Optional[str] = None
+
+
+class StateHostMeta:
+    ops = 0
+    ops_change = 0
+    ops_no_change = 0
+    op_hashes: set[str]
+
+    def __init__(self) -> None:
+        self.op_hashes = set()
+
+
+class StateHostResults:
+    ops = 0
+    success_ops = 0
+    error_ops = 0
+    ignored_error_ops = 0
+    partial_ops = 0
+
+
+class State:
+    """
     Manages state for a pyinfra deploy.
-    '''
+    """
 
-    initialised = False
+    initialised: bool = False
 
     # A pyinfra.api.Inventory which stores all our pyinfra.api.Host's
-    inventory = None
+    inventory: "Inventory"
 
     # A pyinfra.api.Config
-    config = None
+    config: "Config"
 
     # Main gevent pool
-    pool = None
-
-    # Whether we are in an @operation (so inner ops aren't wrapped)
-    in_op = False
-    # Whether we are deploying (ie hosts are all ready)
-    deploying = False
+    pool: "Pool"
 
-    # Current op hash for use w/facts
-    current_op_hash = None
+    # Current stage this state is in
+    current_stage: StateStage = StateStage.Setup
+    # Warning counters by stage
+    stage_warnings: dict[StateStage, int] = defaultdict(int)
 
-    loop_counter = None
-    loop_line = None
+    # Whether we are executing operations (ie hosts are all ready)
+    is_executing: bool = False
 
-    # Name of the current deploy
-    in_deploy = False
-    deploy_name = None
-    deploy_kwargs = None
-    deploy_data = None
-    deploy_line_numbers = None
+    # Whether we should check for operation changes as part of the operation ordering phase, this
+    # allows us to guesstimate which ops will result in changes on which hosts.
+    check_for_changes: bool = True
 
-    # Flags for printing
-    print_output = False  # print output from the actual deploy (-v)
-    print_fact_info = False  # log fact gathering as INFO > DEBUG (-v)
-    print_fact_output = False  # print output from facts (-vv)
+    print_noop_info: bool = False  # print "[host] noop: reason for noop"
+    print_fact_info: bool = False  # print "loaded fact X"
+    print_input: bool = False
+    print_fact_input: bool = False
+    print_output: bool = False
+    print_fact_output: bool = False
 
     # Used in CLI
-    deploy_dir = None  # base directory for locating files/templates/etc
-    current_op_file = 0  # increments for each file excuted to maintain order
+    cwd: Optional[str] = None  # base directory for locating files/templates/etc
+    current_deploy_filename: Optional[str] = None
+    current_exec_filename: Optional[str] = None
+    current_op_file_number: int = 0
+    should_raise_failed_hosts: Optional[Callable[["State"], bool]] = None
+
+    def __init__(
+        self,
+        inventory: Optional["Inventory"] = None,
+        config: Optional["Config"] = None,
+        check_for_changes: bool = True,
+        **kwargs,
+    ):
+        """
+        Initializes the state, the main Pyinfra
+
+        Args:
+            inventory (Optional[Inventory], optional): The inventory. Defaults to None.
+            config (Optional[Config], optional): The config object. Defaults to None.
+        """
+        self.check_for_changes = check_for_changes
 
-    def __init__(self, inventory=None, config=None, **kwargs):
         if inventory:
             self.init(inventory, config, **kwargs)
 
-    def init(self, inventory, config, initial_limit=None):
+    def init(
+        self,
+        inventory: "Inventory",
+        config: Optional["Config"],
+        initial_limit=None,
+    ):
         # Config validation
         #
 
@@ -101,23 +216,6 @@ class State(object):
         if config is None:
             config = Config()
 
-        # Error if our min version is not met
-        if config.MIN_PYINFRA_VERSION is not None:
-            running_version = parse_version(__version__)
-            needed_version = parse_version(
-                # Version must be a string
-                str(config.MIN_PYINFRA_VERSION),
-            )
-
-            if needed_version > running_version:
-                raise PyinfraError((
-                    'Minimum pyinfra version not met '
-                    '(minimum={0}, running={1})'
-                ).format(
-                    config.MIN_PYINFRA_VERSION,
-                    __version__,
-                ))
-
         if not config.PARALLEL:
             # TODO: benchmark this
             # In my own tests the optimum number of parallel SSH processes is
@@ -125,215 +223,182 @@ class State(object):
             cpus = cpu_count()
             ideal_parallel = cpus * 20
 
-            config.PARALLEL = (
-                min(ideal_parallel, len(inventory), MAX_PARALLEL)
-                if MAX_PARALLEL is not None
-                else min(ideal_parallel, len(inventory))
-            )
+            config.PARALLEL = min(ideal_parallel, len(inventory), MAX_PARALLEL)
 
         # If explicitly set, just issue a warning
-        elif MAX_PARALLEL is not None and config.PARALLEL > MAX_PARALLEL:
-            logger.warning((
-                'Parallel set to {0}, but this may hit the open files limit of {1}.\n'
-                '    Max recommended value: {2}'
-            ).format(config.PARALLEL, nofile_limit, MAX_PARALLEL))
+        elif config.PARALLEL > MAX_PARALLEL:
+            logger.warning(
+                (
+                    "Parallel set to {0}, but this may hit the open files limit of {1}.\n"
+                    "    Max recommended value: {2}"
+                ).format(config.PARALLEL, nofile_limit, MAX_PARALLEL),
+            )
 
         # Actually initialise the state object
         #
 
+        self.callback_handlers: list[BaseStateCallback] = []
+
         # Setup greenlet pools
         self.pool = Pool(config.PARALLEL)
         self.fact_pool = Pool(config.PARALLEL)
 
-        # Connection storage
-        self.ssh_connections = {}
-        self.sftp_connections = {}
-
         # Private keys
-        self.private_keys = {}
-
-        # Facts storage
-        self.facts = {}
-        self.fact_locks = {}
+        self.private_keys: dict[str, PKey] = {}
 
         # Assign inventory/config
         self.inventory = inventory
         self.config = config
 
         # Hosts we've activated at any time
-        self.activated_hosts = set()
+        self.activated_hosts: set["Host"] = set()
         # Active hosts that *haven't* failed yet
-        self.active_hosts = set()
-        # Hosts that are ready to be deployed to
-        self.ready_hosts = set()
+        self.active_hosts: set["Host"] = set()
+        # Hosts that have failed
+        self.failed_hosts: set["Host"] = set()
 
         # Limit hosts changes dynamically to limit operations to a subset of hosts
-        self.limit_hosts = initial_limit
+        self.limit_hosts: list["Host"] = initial_limit
 
         # Op basics
-        self.op_line_numbers_to_hash = {}
-        self.op_meta = {}  # maps operation hash -> names/etc
-        self.ops_run = set()  # list of ops which have been started/run
+        self.op_meta: dict[str, StateOperationMeta] = {}  # maps operation hash -> names/etc
 
         # Op dict for each host
-        self.ops = {
-            host: {}
-            for host in inventory
-        }
-
-        # Facts dict for each host
-        self.facts = {
-            host: {}
-            for host in inventory
-        }
+        self.ops: dict["Host", dict[str, StateOperationHostData]] = {host: {} for host in inventory}
 
         # Meta dict for each host
-        self.meta = {
-            host: {
-                'ops': 0,  # one function call in a deploy file
-                'commands': 0,  # actual # of commands to run
-                'op_hashes': set(),
-            }
-            for host in inventory
-        }
+        self.meta: dict["Host", StateHostMeta] = {host: StateHostMeta() for host in inventory}
 
         # Results dict for each host
-        self.results = {
-            host: {
-                'ops': 0,  # success_ops + failed ops w/ignore_errors
-                'success_ops': 0,
-                'error_ops': 0,
-                'commands': 0,
-            }
-            for host in inventory
+        self.results: dict["Host", StateHostResults] = {
+            host: StateHostResults() for host in inventory
         }
 
         # Assign state back references to inventory & config
         inventory.state = config.state = self
+        for host in inventory:
+            host.init(self)
 
         self.initialised = True
 
-    @contextmanager
-    def deploy(self, name, kwargs, data, line_number, in_deploy=True):
-        '''
-        Wraps a group of operations as a deploy, this should not be used
-        directly, instead use ``pyinfra.api.deploy.deploy``.
-        '''
-
-        # Handle nested deploy names
-        if self.deploy_name:
-            name = _make_name(self.deploy_name, name)
-
-        # Store the previous values
-        old_in_deploy = self.in_deploy
-        old_deploy_name = self.deploy_name
-        old_deploy_kwargs = self.deploy_kwargs
-        old_deploy_data = self.deploy_data
-        old_deploy_line_numbers = self.deploy_line_numbers
-        self.in_deploy = in_deploy
-
-        # Limit the new hosts to a subset of the old hosts if they existed
-        if (
-            old_deploy_kwargs
-            and old_deploy_kwargs.get('hosts') is not None
-        ):
-            # If we have hosts - subset them based on the old hosts
-            if 'hosts' in kwargs:
-                kwargs['hosts'] = [
-                    host for host in kwargs['hosts']
-                    if host in old_deploy_kwargs['hosts']
-                ]
-            # Otherwise simply carry the previous hosts
-            else:
-                kwargs['hosts'] = old_deploy_kwargs['hosts']
-
-        # Make new line numbers - note convert from and back to tuple to avoid
-        # keeping deploy_line_numbers mutable.
-        new_line_numbers = list(self.deploy_line_numbers or [])
-        new_line_numbers.append(line_number)
-        new_line_numbers = tuple(new_line_numbers)
-
-        # Set the new values
-        self.deploy_name = name
-        self.deploy_kwargs = kwargs
-        self.deploy_data = data
-        self.deploy_line_numbers = new_line_numbers
-        logger.debug('Starting deploy {0} (args={1}, data={2})'.format(
-            name, kwargs, data,
-        ))
-
-        yield
-
-        # Restore the previous values
-        self.in_deploy = old_in_deploy
-        self.deploy_name = old_deploy_name
-        self.deploy_kwargs = old_deploy_kwargs
-        self.deploy_data = old_deploy_data
-        self.deploy_line_numbers = old_deploy_line_numbers
-
-        logger.debug('Reset deploy to {0} (args={1}, data={2})'.format(
-            old_deploy_name, old_deploy_kwargs, old_deploy_data,
-        ))
-
-    @contextmanager
-    def preserve_loop_order(self, items):
-        frameinfo = get_caller_frameinfo(frame_offset=1)  # escape contextlib
-        self.loop_line = frameinfo.lineno
-
-        def item_generator():
-            for i, item in enumerate(items, 1):
-                self.loop_counter = i
-                yield item
-
-        yield item_generator
-
-        self.loop_counter = None
-        self.loop_line = None
+    def set_stage(self, stage: StateStage) -> None:
+        if stage < self.current_stage:
+            raise Exception("State stage cannot go backwards!")
+        self.current_stage = stage
+
+    def increment_warning_counter(self) -> None:
+        self.stage_warnings[self.current_stage] += 1
+
+    def get_warning_counter(self) -> int:
+        return self.stage_warnings[self.current_stage]
+
+    def should_check_for_changes(self):
+        return self.check_for_changes
+
+    def add_callback_handler(self, handler):
+        if not isinstance(handler, BaseStateCallback):
+            raise TypeError(
+                ("{0} is not a valid callback handler (use `BaseStateCallback`)").format(handler),
+            )
+        self.callback_handlers.append(handler)
+
+    def trigger_callbacks(self, method_name: str, *args, **kwargs):
+        for handler in self.callback_handlers:
+            func = getattr(handler, method_name)
+            func(self, *args, **kwargs)
 
     def get_op_order(self):
-        line_numbers_to_hash = self.op_line_numbers_to_hash
-        sorted_line_numbers = sorted(list(line_numbers_to_hash.keys()))
+        ts: TopologicalSorter = TopologicalSorter()
+
+        for host in self.inventory:
+            for i, op_hash in enumerate(host.op_hash_order):
+                if not i:
+                    ts.add(op_hash)
+                else:
+                    ts.add(op_hash, host.op_hash_order[i - 1])
+
+        final_op_order = []
+
+        try:
+            ts.prepare()
+        except CycleError as e:
+            raise PyinfraError(
+                (
+                    "Cycle detected in operation ordering DAG.\n"
+                    f"    Error: {e}\n\n"
+                    "    This can happen when using loops in operation code, "
+                    "please see: https://docs.pyinfra.com/en/latest/deploy-process.html#loops-cycle-errors"  # noqa: E501
+                ),
+            )
+
+        while ts.is_active():
+            # Ensure that where we have multiple different operations that can be executed in any
+            # dependency order we order them by line numbers.
+            node_group = sorted(
+                ts.get_ready(),
+                key=lambda op_hash: self.op_meta[op_hash].op_order,
+            )
+            ts.done(*node_group)
+            final_op_order.extend(node_group)
+
+        return final_op_order
 
-        return [
-            line_numbers_to_hash[numbers]
-            for numbers in sorted_line_numbers
-        ]
+    def get_op_meta(self, op_hash: str) -> StateOperationMeta:
+        return self.op_meta[op_hash]
 
-    def activate_host(self, host):
-        '''
+    def get_meta_for_host(self, host: "Host") -> StateHostMeta:
+        return self.meta[host]
+
+    def get_results_for_host(self, host: "Host") -> StateHostResults:
+        return self.results[host]
+
+    def get_op_data_for_host(
+        self,
+        host: "Host",
+        op_hash: str,
+    ) -> StateOperationHostData:
+        return self.ops[host][op_hash]
+
+    def set_op_data_for_host(
+        self,
+        host: "Host",
+        op_hash: str,
+        op_data: StateOperationHostData,
+    ):
+        self.ops[host][op_hash] = op_data
+
+    def activate_host(self, host: "Host"):
+        """
         Flag a host as active.
-        '''
+        """
 
-        logger.debug('Activating host: {0}'.format(host))
+        logger.debug("Activating host: %s", host)
 
         # Add to *both* activated and active - active will reduce as hosts fail
         # but connected will not, enabling us to track failed %.
         self.activated_hosts.add(host)
         self.active_hosts.add(host)
 
-    # def ready_host(self, host):
-    #     '''
-    #     Flag a host as ready, after which facts will not be gathered for it.
-    #     '''
-
-    #     logger.debug('Readying host: {0}'.format(host))
-    #     self.ready_hosts.add(host)
-
     def fail_hosts(self, hosts_to_fail, activated_count=None):
-        '''
+        """
         Flag a ``set`` of hosts as failed, error for ``config.FAIL_PERCENT``.
-        '''
+        """
 
         if not hosts_to_fail:
             return
 
         activated_count = activated_count or len(self.activated_hosts)
 
-        logger.debug('Failing hosts: {0}'.format(', '.join(
-            (host.name for host in hosts_to_fail),
-        )))
+        logger.debug(
+            "Failing hosts: {0}".format(
+                ", ".join(
+                    (host.name for host in hosts_to_fail),
+                ),
+            ),
+        )
+
+        self.failed_hosts.update(hosts_to_fail)
 
-        # Remove the failed hosts from the inventory
         self.active_hosts -= hosts_to_fail
 
         # Check we're not above the fail percent
@@ -341,40 +406,29 @@ class State(object):
 
         # No hosts left!
         if not active_hosts:
-            raise PyinfraError('No hosts remaining!')
+            raise PyinfraError("No hosts remaining!")
 
         if self.config.FAIL_PERCENT is not None:
-            percent_failed = (
-                1 - len(active_hosts) / activated_count
-            ) * 100
+            percent_failed = (1 - len(active_hosts) / activated_count) * 100
 
             if percent_failed > self.config.FAIL_PERCENT:
-                raise PyinfraError('Over {0}% of hosts failed ({1}%)'.format(
-                    self.config.FAIL_PERCENT,
-                    int(round(percent_failed)),
-                ))
-
-    def is_host_in_limit(self, host):
-        '''
+                if self.should_raise_failed_hosts and self.should_raise_failed_hosts(self) is False:
+                    return
+
+                raise PyinfraError(
+                    "Over {0}% of hosts failed ({1}%)".format(
+                        self.config.FAIL_PERCENT,
+                        int(round(percent_failed)),
+                    ),
+                )
+
+    def is_host_in_limit(self, host: "Host"):
+        """
         Returns a boolean indicating if the host is within the current state limit.
-        '''
+        """
 
         limit_hosts = self.limit_hosts
 
         if not isinstance(limit_hosts, list):
             return True
         return host in limit_hosts
-
-    def get_temp_filename(self, hash_key=None):
-        '''
-        Generate a temporary filename for this deploy.
-        '''
-
-        if not hash_key:
-            hash_key = str(uuid4())
-
-        temp_filename = '{0}/{1}'.format(
-            self.config.TEMP_DIR, sha1_hash(hash_key),
-        )
-
-        return temp_filename