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

pyinfra/api/facts.py

@@ -1,259 +1,310 @@
-'''
+"""
 The pyinfra facts API. Facts enable pyinfra to collect remote server state which
 is used to "diff" with the desired state, producing the final commands required
 for a deploy.
-'''
 
-from inspect import isfunction, ismethod
-from socket import (
-    error as socket_error,
-    timeout as timeout_error,
-)
+Note that the facts API does *not* use the global currently in context host so
+it's possible to call facts on hosts out of context (ie give me the IP of this
+other host B while I operate on this host A).
+"""
+
+from __future__ import annotations
+
+import inspect
+import re
+from inspect import getcallargs
+from socket import error as socket_error, timeout as timeout_error
+from typing import TYPE_CHECKING, Any, Callable, Generic, Iterable, Optional, Type, TypeVar, cast
 
 import click
 import gevent
-
-from gevent.lock import BoundedSemaphore
 from paramiko import SSHException
+from typing_extensions import override
 
 from pyinfra import logger
+from pyinfra.api import StringCommand
+from pyinfra.api.arguments import all_global_arguments, pop_global_arguments
 from pyinfra.api.util import (
-    get_arg_value,
+    get_kwargs_str,
+    log_error_or_warning,
     log_host_command_error,
-    make_hash,
-    underscore,
+    print_host_combined_output,
 )
+from pyinfra.connectors.util import CommandOutput
+from pyinfra.context import ctx_host, ctx_state
 from pyinfra.progress import progress_spinner
 
+from .arguments import CONNECTOR_ARGUMENT_KEYS
 
-# Index of snake_case facts -> CamelCase classes
-FACTS = {}
-FACT_LOCK = BoundedSemaphore()
+if TYPE_CHECKING:
+    from pyinfra.api import Host, State
 
+SUDO_REGEX = r"^sudo: unknown user"
+SU_REGEXES = (
+    r"^su: user .+ does not exist",
+    r"^su: unknown login",
+)
 
-def is_fact(name):
-    return name in FACTS
 
+T = TypeVar("T")
 
-def get_fact_names():
-    '''
-    Returns a list of available facts in camel_case format.
-    '''
 
-    return list(FACTS.keys())
+class FactBase(Generic[T]):
+    name: str
 
+    abstract: bool = True
 
-class FactMeta(type):
-    '''
-    Metaclass to dynamically build the facts index.
-    '''
+    shell_executable: str | None = None
 
-    def __init__(cls, name, bases, attrs):
-        if attrs.get('abstract'):
-            return
+    command: Callable[..., str | StringCommand]
 
-        fact_name = underscore(name)
-        cls.name = fact_name
+    def requires_command(self, *args, **kwargs) -> str | None:
+        return None
 
-        # Get the an instance of the fact, attach to facts
-        FACTS[fact_name] = cls
+    @override
+    def __init_subclass__(cls) -> None:
+        super().__init_subclass__()
+        module_name = cls.__module__.replace("pyinfra.facts.", "")
+        cls.name = f"{module_name}.{cls.__name__}"
 
+        # Check that fact's `command` method does not inadvertently take a global
+        # argument, most commonly `name`.
+        if hasattr(cls, "command") and callable(cls.command):
+            command_args = set(inspect.signature(cls.command).parameters.keys())
+            global_args = set([name for name, _ in all_global_arguments()])
+            command_global_args = command_args & global_args
 
-class FactBase(object, metaclass=FactMeta):
-    abstract = True
+            if len(command_global_args) > 0:
+                names = ", ".join(command_global_args)
+                raise TypeError(f"{cls.name}'s arguments {names} are reserved for global arguments")
 
     @staticmethod
-    def default():
-        '''
+    def default() -> T:
+        """
         Set the default attribute to be a type (eg list/dict).
-        '''
+        """
 
-    @staticmethod
-    def process(output):
-        return '\n'.join(output)
+        return cast(T, None)
 
-    def process_pipeline(self, args, output):
-        return {
-            arg: self.process([output[i]])
-            for i, arg in enumerate(args)
-        }
+    def process(self, output: Iterable[str]) -> T:
+        # NOTE: TypeVar does not support a default, so we have to cast this str -> T
+        return cast(T, "\n".join(output))
 
+    def process_pipeline(self, args, output):
+        return {arg: self.process([output[i]]) for i, arg in enumerate(args)}
 
-class ShortFactBase(object, metaclass=FactMeta):
-    fact = None
 
+class ShortFactBase(Generic[T]):
+    name: str
+    fact: Type[FactBase]
 
-def get_short_facts(state, short_fact, **kwargs):
-    facts = get_facts(state, short_fact.fact.name, **kwargs)
+    @override
+    def __init_subclass__(cls) -> None:
+        super().__init_subclass__()
+        module_name = cls.__module__.replace("pyinfra.facts.", "")
+        cls.name = f"{module_name}.{cls.__name__}"
 
-    return {
-        host: short_fact.process_data(data)
-        for host, data in facts.items()
-    }
+    def process_data(self, data):
+        return data
 
 
-def get_facts(state, name, args=None, ensure_hosts=None):
-    '''
-    Get a single fact for all hosts in the state.
-    '''
+def get_short_facts(state: "State", host: "Host", short_fact, **kwargs):
+    fact_data = get_fact(state, host, short_fact.fact, **kwargs)
+    return short_fact().process_data(fact_data)
 
-    # Create an instance of the fact
-    fact = FACTS[name]()
 
-    if isinstance(fact, ShortFactBase):
-        return get_short_facts(state, fact, args=args, ensure_hosts=ensure_hosts)
+def _make_command(command_attribute, host_args):
+    if callable(command_attribute):
+        host_args.pop("self", None)
+        return command_attribute(**host_args)
+    return command_attribute
 
-    logger.debug('Getting fact: {0} (ensure_hosts: {1})'.format(
-        name, ensure_hosts,
-    ))
 
+def _handle_fact_kwargs(state: "State", host: "Host", cls, args, kwargs):
     args = args or []
+    kwargs = kwargs or {}
 
-    # Apply args or defaults
-    sudo = state.config.SUDO
-    sudo_user = state.config.SUDO_USER
-    su_user = state.config.SU_USER
-    ignore_errors = state.config.IGNORE_ERRORS
-
-    # Timeout for operations !== timeout for connect (config.CONNECT_TIMEOUT)
-    timeout = None
-
-    # Get the current op meta
-    current_op_hash = state.current_op_hash
-    current_op_meta = state.op_meta.get(current_op_hash)
-
-    # If inside an operation, fetch config meta
-    if current_op_meta:
-        sudo = current_op_meta['sudo']
-        sudo_user = current_op_meta['sudo_user']
-        su_user = current_op_meta['su_user']
-        ignore_errors = current_op_meta['ignore_errors']
-        timeout = current_op_meta['timeout']
-
-    # Make a hash which keeps facts unique - but usable cross-deploy/threads.
-    # Locks are used to maintain order.
-    fact_hash = make_hash((name, args, sudo, sudo_user, su_user, ignore_errors))
-
-    # Already got this fact? Unlock and return 'em
-    current_facts = state.facts.get(fact_hash, {})
-    if current_facts:
-        if not ensure_hosts or all(
-            host in current_facts for host in ensure_hosts
-        ):
-            return current_facts
-
-    with FACT_LOCK:
-        # Add any hosts we must have, whether considered in the inventory or not
-        # (these hosts might be outside the --limit or current op limit_hosts).
-        hosts = set(state.inventory)
-        if ensure_hosts:
-            hosts.update(ensure_hosts)
-
-        # Execute the command for each state inventory in a greenlet
-        greenlet_to_host = {}
-
-        for host in hosts:
-            if host in current_facts:
-                continue
-
-            # if host in state.ready_hosts:
-            #     continue
-
-            # Work out the command
-            command = fact.command
-
-            if ismethod(command) or isfunction(command):
-                # Generate actual arguments by pasing strings as jinja2 templates
-                host_args = [get_arg_value(state, host, arg) for arg in args]
-
-                command = command(*host_args)
-
-            greenlet = state.fact_pool.spawn(
-                host.run_shell_command, state, command,
-                sudo=sudo, sudo_user=sudo_user,
-                su_user=su_user, timeout=timeout,
-                print_output=state.print_fact_output,
-            )
-            greenlet_to_host[greenlet] = host
-
-        # Wait for all the commands to execute
-        progress_prefix = 'fact: {0}'.format(name)
-        if args:
-            progress_prefix = '{0}{1}'.format(progress_prefix, args)
-
-        with progress_spinner(
-            greenlet_to_host.values(),
-            prefix_message=progress_prefix,
-        ) as progress:
-            for greenlet in gevent.iwait(greenlet_to_host.keys()):
-                host = greenlet_to_host[greenlet]
-                progress(host)
-
-        hostname_facts = {}
-        failed_hosts = set()
-
-        # Collect the facts and any failures
-        for greenlet, host in greenlet_to_host.items():
-            status = False
-            stdout = []
-
-            try:
-                status, stdout, _ = greenlet.get()
-
-            except (timeout_error, socket_error, SSHException) as e:
-                failed_hosts.add(host)
-                log_host_command_error(
-                    host, e,
-                    timeout=timeout,
-                )
-
-            data = fact.default()
-
-            if status and stdout:
-                data = fact.process(stdout)
-
-            hostname_facts[host] = data
-
-        log_name = click.style(name, bold=True)
-
-        filtered_args = list(filter(None, args))
-        if filtered_args:
-            log = 'Loaded fact {0}: {1}'.format(log_name, tuple(filtered_args))
-        else:
-            log = 'Loaded fact {0}'.format(log_name)
+    # Start with a (shallow) copy of current operation kwargs if any
+    ctx_kwargs: dict[str, Any] = (
+        cast(dict[str, Any], host.current_op_global_arguments) or {}
+    ).copy()
+    # Update with the input kwargs (overrides)
+    ctx_kwargs.update(kwargs)
 
-        if state.print_fact_info:
-            logger.info(log)
-        else:
-            logger.debug(log)
+    # Pop executor kwargs, pass remaining
+    global_kwargs, _ = pop_global_arguments(state, host, cast(dict[str, Any], ctx_kwargs))
 
-        # Check we've not failed
-        if not ignore_errors:
-            state.fail_hosts(failed_hosts)
+    fact_kwargs = {key: value for key, value in kwargs.items() if key not in global_kwargs}
 
-        # Assign the facts
-        state.facts.setdefault(fact_hash, {}).update(hostname_facts)
+    if args or fact_kwargs:
+        # Merges args & kwargs into a single kwargs dictionary
+        fact_kwargs = getcallargs(cls().command, *args, **fact_kwargs)
 
-    return state.facts[fact_hash]
+    return fact_kwargs, global_kwargs
 
 
-def get_fact(state, host, name):
-    '''
-    Wrapper around ``get_facts`` returning facts for one host or a function
-    that does.
-    '''
+def get_facts(state, *args, **kwargs):
+    def get_host_fact(host, *args, **kwargs):
+        with ctx_host.use(host):
+            return get_fact(state, host, *args, **kwargs)
 
-    # Expecting a function to return
-    if callable(getattr(FACTS[name], 'command', None)):
-        def wrapper(*args):
-            fact_data = get_facts(state, name, args=args, ensure_hosts=(host,))
+    with ctx_state.use(state):
+        greenlet_to_host = {
+            state.pool.spawn(get_host_fact, host, *args, **kwargs): host
+            for host in state.inventory.iter_active_hosts()
+        }
 
-            return fact_data.get(host)
-        return wrapper
+    results = {}
+
+    with progress_spinner(greenlet_to_host.values()) as progress:
+        for greenlet in gevent.iwait(greenlet_to_host.keys()):
+            host = greenlet_to_host[greenlet]
+            results[host] = greenlet.get()
+            progress(host)
+
+    return results
+
+
+def get_fact(
+    state: "State",
+    host: "Host",
+    cls: type[FactBase],
+    args: Optional[Any] = None,
+    kwargs: Optional[Any] = None,
+    ensure_hosts: Optional[Any] = None,
+    apply_failed_hosts: bool = True,
+) -> Any:
+    if issubclass(cls, ShortFactBase):
+        return get_short_facts(
+            state,
+            host,
+            cls,
+            args=args,
+            kwargs=kwargs,
+            ensure_hosts=ensure_hosts,
+            apply_failed_hosts=apply_failed_hosts,
+        )
+
+    return _get_fact(
+        state,
+        host,
+        cls,
+        args,
+        kwargs,
+        ensure_hosts,
+        apply_failed_hosts,
+    )
+
+
+def _get_fact(
+    state: "State",
+    host: "Host",
+    cls: type[FactBase],
+    args: Optional[list] = None,
+    kwargs: Optional[dict] = None,
+    ensure_hosts: Optional[Any] = None,
+    apply_failed_hosts: bool = True,
+) -> Any:
+    fact = cls()
+    name = fact.name
+
+    fact_kwargs, global_kwargs = _handle_fact_kwargs(state, host, cls, args, kwargs)
+
+    kwargs_str = get_kwargs_str(fact_kwargs)
+    logger.debug(
+        "Getting fact: %s (%s) (ensure_hosts: %r)",
+        name,
+        kwargs_str,
+        ensure_hosts,
+    )
+
+    if not host.connected:
+        host.connect(
+            reason=f"to load fact: {name} ({kwargs_str})",
+            raise_exceptions=True,
+        )
+
+    # Facts can override the shell (winrm powershell vs cmd support)
+    if fact.shell_executable:
+        global_kwargs["_shell_executable"] = fact.shell_executable
+
+    command = _make_command(fact.command, fact_kwargs)
+    requires_command = _make_command(fact.requires_command, fact_kwargs)
+    if requires_command:
+        command = StringCommand(
+            # Command doesn't exist, return 0 *or* run & return fact command
+            "!",
+            "command",
+            "-v",
+            requires_command,
+            ">/dev/null",
+            "||",
+            command,
+        )
+
+    status = False
+    output = CommandOutput([])
+
+    executor_kwargs = {
+        key: value for key, value in global_kwargs.items() if key in CONNECTOR_ARGUMENT_KEYS
+    }
 
-    # Expecting the fact as a return value
+    try:
+        status, output = host.run_shell_command(
+            command,
+            print_output=state.print_fact_output,
+            print_input=state.print_fact_input,
+            **executor_kwargs,
+        )
+    except (timeout_error, socket_error, SSHException) as e:
+        log_host_command_error(
+            host,
+            e,
+            timeout=global_kwargs.get("_timeout"),
+        )
+
+    stdout_lines, stderr_lines = output.stdout_lines, output.stderr_lines
+
+    data = fact.default()
+
+    if status:
+        if stdout_lines:
+            data = fact.process(stdout_lines)
+    elif stderr_lines:
+        # If we have error output and that error is sudo or su stating the user
+        # does not exist, do not fail but instead return the default fact value.
+        # This allows for users that don't currently but may be created during
+        # other operations.
+        first_line = stderr_lines[0]
+        if executor_kwargs["_sudo_user"] and re.match(SUDO_REGEX, first_line):
+            status = True
+        if executor_kwargs["_su_user"] and any(re.match(regex, first_line) for regex in SU_REGEXES):
+            status = True
+
+    if status:
+        log_message = "{0}{1}".format(
+            host.print_prefix,
+            "Loaded fact {0}{1}".format(
+                click.style(name, bold=True),
+                f" ({get_kwargs_str(kwargs)})" if kwargs else "",
+            ),
+        )
+        if state.print_fact_info:
+            logger.info(log_message)
+        else:
+            logger.debug(log_message)
     else:
-        # Get the fact
-        fact_data = get_facts(state, name, ensure_hosts=(host,))
+        if not state.print_fact_output:
+            print_host_combined_output(host, output)
+
+        log_error_or_warning(
+            host,
+            global_kwargs["_ignore_errors"],
+            description=("could not load fact: {0} {1}").format(name, get_kwargs_str(fact_kwargs)),
+        )
+
+    # Check we've not failed
+    if apply_failed_hosts and not status and not global_kwargs["_ignore_errors"]:
+        state.fail_hosts({host})
 
-        return fact_data.get(host)
+    return data