shell_utils

drunc.utils.shell_utils

Shell utilities for DRUNC.

Classes

InterruptedCommand(message='An error occurred in Drunc.', grpc_error_code=None, details=None, reason=None, domain=None, **detail_kwargs)

Bases: DruncShellException

Exception thrown to interrupt a shell command without a full stack trace.

Source code in drunc/exceptions.py

def __init__(
    self,
    message: str | None = "An error occurred in Drunc.",
    grpc_error_code: int | None = None,
    details: str | None = None,
    reason: str | None = None,
    domain: str | None = None,
    **detail_kwargs: object,
) -> None:
    super().__init__(message)

    self.message: str = (
        message if message is not None else "An error occurred in Drunc."
    )

    self.grpc_error_code: int = (
        grpc_error_code
        if grpc_error_code is not None
        else int(getattr(self.__class__, "grpc_error_code", code_pb2.INTERNAL))
    )

    self.reason: str = (
        reason
        if reason is not None
        else str(getattr(self.__class__, "reason", self.__class__.__name__))
    )

    self.domain: str = (
        domain
        if domain is not None
        else str(getattr(self.__class__, "domain", "drunc"))
    )

    self.details: str | None = details
    self.detail_kwargs: dict[str, object] = detail_kwargs

    error_metadata: dict[str, str] = {"message": self.message}
    for key, value in self.detail_kwargs.items():
        error_metadata[key] = str(value)

    self.base_error_info = error_details_pb2.ErrorInfo(
        reason=self.reason, domain=self.domain, metadata=error_metadata
    )

ShellContext(*args, **kwargs)

Base class for shell contexts.

Initialize the shell context.

Parameters:

Name	Type	Description	Default
*args	object	Additional positional arguments.	()
**kwargs	object	Additional keyword arguments.	{}

Source code in drunc/utils/shell_utils.py

def __init__(self, *args: object, **kwargs: object) -> None:
    """Initialize the shell context.

    Args:
        *args: Additional positional arguments.
        **kwargs: Additional keyword arguments.
    """
    log = get_logger("utils.ShellContext")
    self.dynamic_commands: set[str] = set()
    try:
        self.reset(*args, **kwargs)
    except Exception as e:
        log.exception(e)
        exit(1)

Methods:

create_drivers(**kwargs) abstractmethod

Create drivers for the context.

Parameters:

Name	Type	Description	Default
**kwargs	object	Additional keyword arguments.	{}

Returns:

Type	Description
MutableMapping[str, object]	MutableMapping[str, object]: A mapping of driver names to driver objects.

Source code in drunc/utils/shell_utils.py

@abc.abstractmethod
def create_drivers(self, **kwargs: object) -> MutableMapping[str, object]:
    """Create drivers for the context.

    Args:
        **kwargs: Additional keyword arguments.

    Returns:
        MutableMapping[str, object]: A mapping of driver names to driver objects.
    """
    pass

create_token(**kwargs) abstractmethod

Create a token for the context.

Parameters:

Name	Type	Description	Default
**kwargs	object	Additional keyword arguments.	{}

Returns:

Name	Type	Description
Token	Token	A token object.

Source code in drunc/utils/shell_utils.py

@abc.abstractmethod
def create_token(self, **kwargs: object) -> Token:
    """Create a token for the context.

    Args:
        **kwargs: Additional keyword arguments.

    Returns:
        Token: A token object.
    """
    pass

delete_driver(name)

Delete a driver from the context.

Parameters:

Name	Type	Description	Default
name	str	The name of the driver to delete.	required

Source code in drunc/utils/shell_utils.py

def delete_driver(self, name: str) -> None:
    """Delete a driver from the context.

    Args:
        name: The name of the driver to delete.
    """
    log = get_logger("utils.ShellContext")
    if name in self._drivers:
        log.info(f"You will not be able to issue commands to the {name} anymore.")
        del self._drivers[name]
        log.info(f"{name.capitalize()} driver has been deleted.")

get_controller_driver(quiet_fail=False)

Get the root controller driver from the context.

Parameters:

Name	Type	Description	Default
quiet_fail	bool	If True, return None on failure instead of raising an exception.	False

Returns:

Name	Type	Description
ControllerDriver	ControllerDriver	The process manager driver.

Raises:

Type	Description
RuntimeError	If the process manager driver is not initialized.

Source code in drunc/utils/shell_utils.py

def get_controller_driver(self, quiet_fail: bool = False) -> ControllerDriver:
    """
    Get the root controller driver from the context.

    Args:
        quiet_fail: If True, return None on failure instead of raising an exception.

    Returns:
        ControllerDriver: The process manager driver.

    Raises:
        RuntimeError: If the process manager driver is not initialized.
    """
    ctrld = self.get_driver("controller", quiet_fail=quiet_fail)

    if not isinstance(ctrld, ControllerDriver):
        raise RuntimeError("ControllerDriver is not loaded!")

    return ctrld

get_driver(name=None, quiet_fail=False)

Get a driver from the context.

Parameters:

Name	Type	Description	Default
name	str | None	The name of the driver. If None, returns the only driver if there is exactly one.	None
quiet_fail	bool	If True, return None on failure instead of raising an exception.	False

Returns:

Name	Type	Description
object	object	The driver object, or None if quiet_fail is True and the driver is not found.

Raises:

Type	Description
DruncShellException	If there are multiple drivers and no name is specified.
SystemExit	If the driver is not found and quiet_fail is False.

Source code in drunc/utils/shell_utils.py

def get_driver(self, name: str | None = None, quiet_fail: bool = False) -> object:
    """Get a driver from the context.

    Args:
        name: The name of the driver. If None, returns the only driver if there is exactly one.
        quiet_fail: If True, return None on failure instead of raising an exception.

    Returns:
        object: The driver object, or None if quiet_fail is True and the driver is not found.

    Raises:
        DruncShellException: If there are multiple drivers and no name is specified.
        SystemExit: If the driver is not found and quiet_fail is False.
    """
    try:
        if name:
            return self._drivers[name]
        elif len(self._drivers) > 1:
            raise DruncShellException("More than one driver in this context")
        return list(self._drivers.values())[0]
    except KeyError:
        if quiet_fail:
            return None
        log = get_logger("utils.ShellContext")
        log.exception(
            "Controller-specific commands cannot be sent until the session is booted"
        )
        log.debug(f"Drivers available are {self._drivers.keys()}")
        raise SystemExit(
            1
        )  # used to avoid having to catch multiple Attribute errors when this function gets called

get_pm_driver(quiet_fail=False)

Get the process manager driver from the context.

Parameters:

Name	Type	Description	Default
quiet_fail	bool	If True, return None on failure instead of raising an exception.	False

Returns:

Name	Type	Description
ProcessManagerDriver	ProcessManagerDriver	The process manager driver.

Raises:

Type	Description
RuntimeError	If the process manager driver is not initialized.

Source code in drunc/utils/shell_utils.py

def get_pm_driver(self, quiet_fail: bool = False) -> ProcessManagerDriver:
    """
    Get the process manager driver from the context.

    Args:
        quiet_fail: If True, return None on failure instead of raising an exception.

    Returns:
        ProcessManagerDriver: The process manager driver.

    Raises:
        RuntimeError: If the process manager driver is not initialized.
    """
    pmd = self.get_driver("process_manager", quiet_fail=quiet_fail)

    # Check the type for mypy and runtime safety.
    if not isinstance(pmd, ProcessManagerDriver):
        raise RuntimeError("ProcessManagerDriver is not loaded!")

    return pmd

get_shell_id()

Get the shell ID.

This is primarily used by the superclasses of the ShellContext to identify the type of shell (e.g., PM shell, Unified shell) for logging purposes.

Returns:

Type	Description
str | None	str | None: The shell ID, or None if not set.

Source code in drunc/utils/shell_utils.py

def get_shell_id(self) -> str | None:
    """
    Get the shell ID.

    This is primarily used by the superclasses of the ShellContext to identify the
    type of shell (e.g., PM shell, Unified shell) for logging purposes.

    Args:
        None

    Returns:
        str | None: The shell ID, or None if not set.

    Raises:
        None
    """
    return self.shell_id

get_token()

Get the token from the context.

Returns:

Name	Type	Description
Token	Token	The token object.

Source code in drunc/utils/shell_utils.py

def get_token(self) -> Token:
    """Get the token from the context.

    Returns:
        Token: The token object.
    """
    return self._token

has_driver(name)

Check if a driver exists in the context.

Parameters:

Name	Type	Description	Default
name	str	The name of the driver.	required

Returns:

Name	Type	Description
bool	bool	True if the driver exists, False otherwise.

Source code in drunc/utils/shell_utils.py

def has_driver(self, name: str) -> bool:
    """Check if a driver exists in the context.

    Args:
        name: The name of the driver.

    Returns:
        bool: True if the driver exists, False otherwise.
    """
    return name in self._drivers

print(*args, **kwargs)

Print to the console.

Parameters:

Name	Type	Description	Default
*args	object	Positional arguments to pass to the console.	()
**kwargs	object	Keyword arguments to pass to the console.	{}

Source code in drunc/utils/shell_utils.py

def print(self, *args: object, **kwargs: object) -> None:
    """Print to the console.

    Args:
        *args: Positional arguments to pass to the console.
        **kwargs: Keyword arguments to pass to the console.
    """
    self._console.print(*args, **kwargs)  # type: ignore[arg-type]

print_status_summary()

Print a summary of the FSM status and available transitions.

Source code in drunc/utils/shell_utils.py

def print_status_summary(self) -> None:
    """Print a summary of the FSM status and available transitions."""
    log = get_logger("utils.ShellContext")
    controller = self.get_controller_driver()
    status = controller.status().status
    describe_fsm = controller.describe_fsm().description
    current_state = status.state
    if status.in_error:
        log.error(
            f"[red] FSM is in error ({status})[/red], not currently accepting new commands."
        )
    else:
        available_actions = [
            command.name.replace("_", "-") for command in describe_fsm.commands
        ]
        available_sequences = [
            seq.id.replace("_", "-") for seq in describe_fsm.sequences
        ]

        log.info(
            f"Current FSM status is [green]{current_state}[/green]. Available transitions are [green]{'[/green], [green]'.join(available_actions)}[/green]. Available sequence commands are [green]{'[/green], [green]'.join(available_sequences)}[/green]."
        )

reset(*args, **kwargs) abstractmethod

Reset the shell context.

Parameters:

Name	Type	Description	Default
*args	object	Additional positional arguments.	()
**kwargs	object	Additional keyword arguments.	{}

Source code in drunc/utils/shell_utils.py

@abc.abstractmethod
def reset(self, *args: object, **kwargs: object) -> None:
    """Reset the shell context.

    Args:
        *args: Additional positional arguments.
        **kwargs: Additional keyword arguments.
    """
    pass

rule(*args, **kwargs)

Print a rule to the console.

Parameters:

Name	Type	Description	Default
*args	object	Positional arguments to pass to the console.	()
**kwargs	object	Keyword arguments to pass to the console.	{}

Source code in drunc/utils/shell_utils.py

def rule(self, *args: object, **kwargs: object) -> None:
    """Print a rule to the console.

    Args:
        *args: Positional arguments to pass to the console.
        **kwargs: Keyword arguments to pass to the console.
    """
    self._console.rule(*args, **kwargs)  # type: ignore[arg-type]

set_driver(name, driver)

Set a driver in the context.

Parameters:

Name	Type	Description	Default
name	str	The name of the driver.	required
driver	object	The driver object.	required

Raises:

Type	Description
DruncShellException	If a driver with the same name already exists.

Source code in drunc/utils/shell_utils.py

def set_driver(self, name: str, driver: object) -> None:
    """Set a driver in the context.

    Args:
        name: The name of the driver.
        driver: The driver object.

    Raises:
        DruncShellException: If a driver with the same name already exists.
    """
    if name in self._drivers:
        raise DruncShellException(f"Driver {name} already present in this context")
    self._drivers[name] = driver

terminate() abstractmethod

Terminate the shell context.

Source code in drunc/utils/shell_utils.py

@abc.abstractmethod
def terminate(self) -> None:
    """Terminate the shell context."""
    pass

Functions:

create_dummy_token_from_uname()

Create a dummy token from the current username.

Returns:

Name	Type	Description
Token	Token	A dummy token with the current username.

Source code in drunc/utils/shell_utils.py

def create_dummy_token_from_uname() -> Token:
    """Create a dummy token from the current username.

    Returns:
        Token: A dummy token with the current username.
    """
    user = getpass.getuser()
    return (
        Token(  # fake token, but should be figured out from the environment/authoriser
            token=f"{user}-token", user_name=user
        )
    )

log_pm_cmd(obj)

Log a process-manager shell command with only explicitly provided arguments.

The current Click command context is inspected and only parameters whose source is COMMANDLINE are included in the log message. This keeps defaulted values out of the message while still recording the command name, optional session name, and shell identity.

These are sent over via so that it can be displayed in the process manager shell

Parameters:

Name	Type	Description	Default
obj	ShellContext	Active shell context used to send the log message.	required

Source code in drunc/utils/shell_utils.py

def log_pm_cmd(obj: ShellContext) -> None:
    """Log a process-manager shell command with only explicitly provided arguments.

    The current Click command context is inspected and only parameters whose source is
    ``COMMANDLINE`` are included in the log message. This keeps defaulted values out
    of the message while still recording the command name, optional session name, and
    shell identity.

    These are sent over via  so that it can be displayed in the process manager
    shell

    Args:
        obj (ShellContext): Active shell context used to send the log message.
    """

    ctx_cmd = click.get_current_context(silent=True)
    cmd_name = ctx_cmd.command.name if ctx_cmd and ctx_cmd.command else None
    parms_dict: dict[str, str] = {}
    if ctx_cmd and ctx_cmd.command:
        for param in ctx_cmd.command.params:
            name = param.name
            if name is None:
                continue
            if (
                ctx_cmd.get_parameter_source(name)
                == click.core.ParameterSource.COMMANDLINE
            ):
                parms_dict[name] = f"{ctx_cmd.params[name]!r}"

    args = f" with arguments {parms_dict}" if parms_dict else ""
    session = f" for session {obj.session_name}" if hasattr(obj, "session_name") else ""
    msg = f"{getpass.getuser()} sent {cmd_name}{args}{session} via {obj.get_shell_id()}"
    pm_driver = obj.get_pm_driver()
    pm_driver.log_on_server(msg)