shell_utils

drunc.controller.interface.shell_utils

Classes

StatusTableUpdater(ctx, refresh_per_second=2, *args, **kwargs)

Bases: Progress

Initialize the status table updater.

Parameters:

Name	Type	Description	Default
ctx	ControllerContext | UnifiedShellContext	The controller context or unified shell context.	required
refresh_per_second	float	How often to refresh the status table.	2
*args	object	Additional positional arguments for the Progress superclass.	()
**kwargs	object	Additional keyword arguments for the Progress superclass.	{}

Returns:

Type	Description
None	None.

Source code in drunc/controller/interface/shell_utils.py

def __init__(
    self,
    ctx: ControllerContext | UnifiedShellContext,
    refresh_per_second: float = 2,
    *args: object,
    **kwargs: object,
) -> None:
    """
    Initialize the status table updater.

    Args:
        ctx: The controller context or unified shell context.
        refresh_per_second: How often to refresh the status table.
        *args: Additional positional arguments for the Progress superclass.
        **kwargs: Additional keyword arguments for the Progress superclass.

    Returns:
        None.

    Raises:
        None.
    """
    # Store the controller context for later use.
    self.ctx = ctx

    # Obtain the shared Rich console for consistent output styling.
    shared_console = get_shared_rich_console(self.ctx.log)
    if shared_console:
        kwargs["console"] = shared_console

    # Detect if standard output is a real terminal screen. This is most common in
    # integration tests, where the output is typically captured and not displayed on
    # a real terminal. This prevents the cascade of status tables.
    # Note - "TERM" == "dumb" indicates a non-interactive terminal. See the UNIX
    # docs: https://invisible-island.net/ncurses/terminfo.src.html#tic-dumb
    self.is_interactive = sys.stdout.isatty() and os.environ.get("TERM") != "dumb"

    # Completely disable the animation loop if writing to a file or CI pipe
    if not self.is_interactive:
        kwargs["disable"] = True

    # Remove default Live parameters so our required implemnentation holds.
    kwargs.pop("vertical_overflow", None)

    # Remove the 'fit' parameter to ensure our custom layout handling is used.
    kwargs.pop("fit", None)

    super().__init__(*args, refresh_per_second=refresh_per_second, **kwargs)

    # Only configure Live parameters if we are actually animating
    if hasattr(self, "live") and self.is_interactive:
        self.live.vertical_overflow = "visible"

    self.update_table()

Methods:

get_renderable()

Get the renderable object for the status table, wrapped in a Group.

Returns:

Type	Description
	A Rich renderable object representing the status table.

Source code in drunc/controller/interface/shell_utils.py

def get_renderable(self):
    """
    Get the renderable object for the status table, wrapped in a Group.

    Args:
        None.

    Returns:
        A Rich renderable object representing the status table.

    Raises:
        None.
    """
    # Retrieve the raw status table, if it exists. If not, return a Group containing
    # only the other renderables.
    raw_table = getattr(self, "table", None)
    if not raw_table:
        return Group(*self.get_renderables())

    # Wrap the raw table in a WrappingSafeTable to ensure it is displayed correctly.
    protected_table = WrappingSafeTable(raw_table)
    return Group(*self.get_renderables(), protected_table)

stop()

Stop the live display and print the final state if animation was disabled.

Returns:

Type	Description
None	None.

Source code in drunc/controller/interface/shell_utils.py

def stop(self) -> None:
    """
    Stop the live display and print the final state if animation was disabled.

    Args:
        None.

    Returns:
        None.

    Raises:
        None.
    """
    super().stop()

    # If animation was disabled for a file, print the final state once at the end.
    if getattr(self, "disable", False):
        self.update_table()

        # Print the task description
        for task in self.tasks:
            self.console.print(task.description)

        # Print the protected table
        raw_table = getattr(self, "table", None)
        if raw_table:
            self.console.print(WrappingSafeTable(raw_table))

update_table()

Update the status table with the latest information from the controller context.

Returns:

Type	Description
None	None.

Source code in drunc/controller/interface/shell_utils.py

def update_table(self) -> None:
    """
    Update the status table with the latest information from the controller context.

    Args:
        None.

    Returns:
        None.

    Raises:
        None.
    """
    self.table = render_status_table(self.ctx)

WrappingSafeTable(renderable)

Renders the table that is safe for the updater, with integtest passes.

When running the StatusTableUpdater, this class ensures that tables are rendered at their full width to prevent ellipses and maintain proper cursor tracking. This is necessary for integration tests to be independent of terminal width limitations.

Initialize the WrappingSafeTable with the given renderable.

Parameters:

Name	Type	Description	Default
renderable	Table | Group	The table or group to be rendered safely.	required

Returns:

Name	Type	Description
None		This initializer does not return anything.

Raises:

Type	Description
TypeError	If the provided renderable is not a Table or Group.

Source code in drunc/controller/interface/shell_utils.py

def __init__(self, renderable: Table | Group):
    """
    Initialize the WrappingSafeTable with the given renderable.

    Args:
        renderable (Table | Group): The table or group to be rendered safely.

    Returns:
        None: This initializer does not return anything.

    Raises:
        TypeError: If the provided renderable is not a Table or Group.
    """

    def lock_table(obj):
        """
        Lock the table to prevent automatic expansion and enable proper wrapping.

        Args:
            obj (Table | Group): The table or group to lock.

        Returns:
            None: This function modifies the object in place.
        """
        # If the object is a Group, recursively lock its renderables, otherwise lock
        # the table itself.
        if isinstance(obj, Group):
            for item in obj.renderables:
                lock_table(item)
        elif hasattr(obj, "columns"):
            obj.expand = False
            obj.width = None
            for col in obj.columns:
                col.no_wrap = True

    lock_table(renderable)
    self.renderable = renderable

Methods:

__rich_console__(console, options)

Render the table safely for the console, with wrapping and cursor tracking. This dunder method is automatically called by the Rich module when necessary.

Parameters:

Name	Type	Description	Default
console	Console	The Rich console to render to.	required
options	ConsoleOptions	The console options, including width constraints.	required

Yields:

Name	Type	Description
Text	RenderResult	Lines of the rendered table, with hard wrapping applied.

Source code in drunc/controller/interface/shell_utils.py

def __rich_console__(
    self, console: Console, options: ConsoleOptions
) -> RenderResult:
    """
    Render the table safely for the console, with wrapping and cursor tracking.
    This dunder method is automatically called by the Rich module when necessary.

    Args:
        console: The Rich console to render to.
        options: The console options, including width constraints.

    Yields:
        Text: Lines of the rendered table, with hard wrapping applied.
    """
    # Set the table width to a very large value to prevent automatic wrapping
    unconstrained = options.update(max_width=10000)
    lines = console.render_lines(self.renderable, unconstrained)

    # For each line of this table render it with hard wrapping, so no contents are
    # wrapped automatically by the console's width, or concatenated which introduces
    # integration test failure
    for line_segments in lines:
        t = Text(overflow="fold")
        for seg in line_segments:
            t.append(seg.text, style=seg.style)

        # Strip trailing padding to prevent phantom blank lines
        t.rstrip()
        yield t

__rich_measure__(console, options)

Measure the renderable for layout purposes. This dunder method is automatically called by the Rich module when necessary.

Parameters:

Name	Type	Description	Default
console	Console	The Rich console to measure against.	required
options	ConsoleOptions	The console options, including width constraints.	required

Returns:

Name	Type	Description
Measurement	Measurement	The minimum and maximum width of the renderable.

Source code in drunc/controller/interface/shell_utils.py

def __rich_measure__(
    self, console: Console, options: ConsoleOptions
) -> Measurement:
    """
    Measure the renderable for layout purposes.
    This dunder method is automatically called by the Rich module when necessary.

    Args:
        console: The Rich console to measure against.
        options: The console options, including width constraints.

    Returns:
        Measurement: The minimum and maximum width of the renderable.
    """
    return Measurement(console.width, console.width)

Functions:

generate_fsm_command(ctx, transition, controller_name)

Generate a click command for a given FSM transition.

Parameters:

Name	Type	Description	Default
ctx		UnifiedShellContext	required
transition	FSMCommandDescription	FSMCommandDescription of the transition to generate the command for	required
controller_name	str	Name of the controller to run the command on	required

Returns:

Type	Description
	A click command that can be added to the CLI

Raises:

Type	Description
Exception	If the argument type is unhandled

Source code in drunc/controller/interface/shell_utils.py

def generate_fsm_command(ctx, transition: FSMCommandDescription, controller_name: str):
    """
    Generate a click command for a given FSM transition.

    Args:
        ctx: UnifiedShellContext
        transition: FSMCommandDescription of the transition to generate the command for
        controller_name: Name of the controller to run the command on

    Returns:
        A click command that can be added to the CLI

    Raises:
        Exception: If the argument type is unhandled
    """

    # Construct the partial command executing the defined FSM command with click options
    cmd = functools.partial(
        run_one_fsm_command,
        controller_name=controller_name,
        transition_name=transition.name,
    )
    cmd = click.pass_obj(cmd)
    cmd = click.option(
        "--target",
        type=str,
        help="The target to address",
        default="",
    )(cmd)

    # Define the mapping of gRPC argument types to click types
    ClickValueType: TypeAlias = type[str] | type[int] | type[float] | type[bool]
    type_map: dict[int, ClickValueType] = {
        Argument.Type.STRING: str,
        Argument.Type.INT: int,
        Argument.Type.FLOAT: float,
        Argument.Type.BOOL: bool,
    }

    # Define the mapping of gRPC argument types to their corresponding protobuf message
    # types for default value unpacking. The ProtobufScalarMsgType is a type alias for
    # the union of the protobuf message types used for scalar values, and is used by
    # mypy to check the types of the unpacked default values. The msg_map is the
    # dictionary that maps the gRPC argument types to their corresponding protobuf
    # message types, which are later handed off to click to define the click command
    # options.
    ProtobufScalarMsgType: TypeAlias = (
        type[string_msg] | type[int_msg] | type[float_msg] | type[bool_msg]
    )
    msg_map: dict[ClickValueType, ProtobufScalarMsgType] = {
        str: string_msg,
        int: int_msg,
        float: float_msg,
        bool: bool_msg,
    }

    # Iterate over the Arguments of the Transitions, and add them as click options to
    # the click command
    for argument in transition.arguments:  # type: Argument
        # Map the gRPC argument type to a click type, raise an exception if the type is
        # unhandled
        atype = type_map.get(argument.type)
        if not atype:
            raise Exception(f"Unhandled argument type '{argument.type}'")

        # Unpack the default value of the argument if it exists, and convert it to the
        # appropriate type
        raw_default: int | float | str | bool | None = None
        if argument.HasField("default_value"):
            unpacked = unpack_any(argument.default_value, msg_map[atype])
            raw_default = atype(unpacked.value)

        # Check for default values defined in the environment variables
        argument_name_cli: str = argument.name.lower().replace("_", "-")
        env_var: str = f"DRUNC_{argument.name.upper()}_DEFAULT"
        env_val: str | None = os.getenv(env_var)

        # Assign the default value if it is present
        if env_val is not None:
            log.info(f"Env override for {argument_name_cli}: {env_val}")
            default_value = atype(env_val)
        else:
            default_value = raw_default

        # Add the argument to the click command
        cmd = click.option(
            f"--{argument_name_cli}",
            type=atype,
            default=default_value,
            show_default=True,
            required=(
                (argument.presence == Argument.Presence.MANDATORY)
                and (default_value is None)
            ),
            help=argument.help,
        )(cmd)

    # Construct the click command
    cmd_name: str = format_name_for_cli(transition.name)
    cmd = click.command(
        name=cmd_name,
        help=f"Execute the transition {transition.name} on the controller {controller_name}",
    )(cmd)

    return cmd, cmd_name

get_hostname_smart(ip_or_host, timeout_seconds=0.2) cached

Resolves an IP or hostname to a human-readable hostname, with optimizations: 1. Caches all results. 2. Replaces localhost/loopback with the machine's actual hostname. 3. Uses a short timeout for reverse DNS lookups on other IPs.

Source code in drunc/controller/interface/shell_utils.py

@functools.lru_cache(maxsize=4096)
def get_hostname_smart(ip_or_host: str, timeout_seconds: float = 0.2) -> str:
    """
    Resolves an IP or hostname to a human-readable hostname, with optimizations:
    1. Caches all results.
    2. Replaces localhost/loopback with the machine's actual hostname.
    3. Uses a short timeout for reverse DNS lookups on other IPs.
    """

    if not ip_or_host:
        return ""

    if ip_or_host == "localhost":
        return socket.getfqdn()

    try:
        ip_address = ipaddress.ip_address(ip_or_host)
    except ValueError:
        fqdn = socket.getfqdn(ip_or_host)
        return fqdn if fqdn else ip_or_host

    if ip_address.is_loopback:
        return socket.getfqdn()

    original_timeout = socket.getdefaulttimeout()
    try:
        socket.setdefaulttimeout(timeout_seconds)
        try:
            hostname, _, _ = socket.gethostbyaddr(str(ip_address))
            return hostname
        except (socket.herror, socket.gaierror, socket.timeout, OSError):
            return ip_or_host

    finally:
        socket.setdefaulttimeout(original_timeout)

run_one_fsm_command(obj, controller_name, transition_name, target, **kwargs)

Run one FSM command on the controller

Parameters:

Name	Type	Description	Default
controller_name	str	Name of the controller	required
transition_name	str	Name of the transition to run	required
obj	UnifiedShellContext	Unified shell context	required
target	str	Target to run the command on	required
**kwargs		Arguments to the command	{}

Returns:

Type	Description
None	None

Raises:

Type	Description
ArgumentException	If there is an issue with the arguments
ServerTimeout	If the server times out

Source code in drunc/controller/interface/shell_utils.py

def run_one_fsm_command(
    obj: UnifiedShellContext,
    controller_name: str,
    transition_name: str,
    target: str,
    **kwargs,
) -> None:
    """
    Run one FSM command on the controller

    Args:
        controller_name (str): Name of the controller
        transition_name (str): Name of the transition to run
        obj (UnifiedShellContext): Unified shell context
        target (str): Target to run the command on
        **kwargs: Arguments to the command

    Returns:
        None

    Raises:
        ArgumentException: If there is an issue with the arguments
        ServerTimeout: If the server times out
    """
    log.info(
        f"Running transition '{transition_name}' on controller '{controller_name}', targeting: '{target if target else controller_name}'"
    )

    # If running in batch or semibatch mode, and error is detected, exit
    if (
        obj.running_mode in [UnifiedShellMode.BATCH, UnifiedShellMode.SEMIBATCH]
        and obj.get_driver("controller").status().status.in_error
        and not obj.no_stop_error_batch_mode
    ):
        obj.get_driver("controller").status()
        log.error(
            "Running in batch mode, and because error state is detected, exiting."
        )
        sys.exit(1)

    execute_along_path = False
    execute_on_all_subsequent_children_in_path = True

    execute_on_root_controller = False
    if target == "":
        execute_on_root_controller = True
    elif target == controller_name:
        execute_on_root_controller = True

    if execute_on_root_controller:
        fsm_description = (
            obj.get_driver("controller")
            .describe_fsm(
                target=controller_name,
                execute_along_path=True,
                execute_on_all_subsequent_children_in_path=False,
            )
            .description
        )

        command_desc = search_fsm_command(transition_name, fsm_description.commands)

        if command_desc is None:
            log.error(
                f'Command "{transition_name}" does not exist, or is not accessible right now'
            )
            return
    else:

        class DummyCommand:
            arguments: list[Argument]

        command_desc = DummyCommand()
        command_desc.arguments = []

    try:
        formated_args = validate_and_format_fsm_arguments(
            kwargs, command_desc.arguments
        )
        data = FSMCommand(
            command_name=transition_name,
            arguments=formated_args,
        )

        resolved_target = target or controller_name
        pre_status = obj.get_driver("controller").status(
            target=resolved_target,
            execute_along_path=execute_along_path,
            execute_on_all_subsequent_children_in_path=execute_on_all_subsequent_children_in_path,
        )
        not_ready = collect_not_ready(pre_status)
        if not_ready:
            log.warning(
                f"The following nodes could not be reached and will not execute '{transition_name}': {not_ready}. "
                f"If this is expected, consider excluding them with the 'exclude' command before retrying."
            )
            if obj.running_mode in [UnifiedShellMode.BATCH, UnifiedShellMode.SEMIBATCH]:
                sys.exit(1)
            return

        timeout = 60
        time_start = time.time()
        result = None

        with ThreadPoolExecutor() as executor:
            future = executor.submit(
                obj.get_driver("controller").execute_fsm_command,
                command=data,
                target=target,
                execute_along_path=execute_along_path,
                execute_on_all_subsequent_children_in_path=execute_on_all_subsequent_children_in_path,
                timeout=timeout,
            )

            with StatusTableUpdater(obj) as updater:
                task = updater.add_task(
                    f"Waiting for [yellow]{transition_name}[/yellow] to complete...",
                    total=timeout,
                )
                while time.time() - time_start < timeout and not future.done():
                    updater.update(task, completed=time.time() - time_start)
                    updater.update_table()
                    time.sleep(0.5)

                updater.update_table()

            result = future.result(timeout=1)

    except ArgumentException as ae:
        log.exception(
            str(ae)
        )  # TODO: Manually raise exception, see if the str declaration is needed with rich handling
        return
    except ServerTimeout:
        log.error(
            "The command timed out, unfortunately this means the server is in undefined state, and [red]your best option at this stage is to [bold]terminate[/bold] and [bold]boot[/bold][/]."
        )
        # The following line is outdated, but in the future when error states and their
        # recovery are better defined, we can provide better options to the user.
        # log.error(
        #     "Alternatively, if you are patient, you can try to wait a bit longer and send [yellow]'status'[/yellow] to check if the command ends up being executed (you may want to check the logs of the controller and application with the [yellow]'logs'[/yellow] command)."
        # )

        # Mark the controller as in error state, so that if the user tries to run
        # another command, it will be prevented, and they will be encouraged to check
        # the error application logs
        status_response = obj.get_driver("controller").status()
        apps_that_timed_out = get_all_apps_with_named_substate(
            status_response, "executing_cmd"
        )
        apps_that_timed_out_str = ", ".join(apps_that_timed_out)
        err_str = (
            "The session did not complete the stateful transition in the specified "
            f"time of {timeout} seconds. To investigate the cause, [yellow]check the "
            f"logs of {apps_that_timed_out_str} with the logs command[/] as:"
        )
        log.error(err_str)
        for app in apps_that_timed_out:
            log.error(f"\t[yellow]logs -n {app}[/]")
        obj.get_driver("controller").send_log(
            err_str, severity="ERROR", logger=LoggerTarget.MAIN
        )
        obj.get_driver("controller").to_error(
            execute_on_all_subsequent_children_in_path=False
        )

        statuses = obj.get_driver("controller").status()
        descriptions = obj.get_driver("controller").describe()
        t = get_status_table(statuses, descriptions)
        obj.print(t)
        obj.print_status_summary()

        return

    if not result:
        return

    t = Table(title=f"{transition_name} execution report")
    t.add_column("Name")
    t.add_column("Command execution")
    t.add_column("FSM transition")

    def bool_to_success(
        flag_message: int, message_type: type[ResponseFlag] | type[FSMResponseFlag]
    ) -> str:
        flag = message_type.Name(flag_message).replace("_", " ").title()
        success = False

        if (
            message_type == FSMResponseFlag
            and flag_message == FSMResponseFlag.FSM_EXECUTED_SUCCESSFULLY
        ):
            success = True
        if (
            message_type == ResponseFlag
            and flag_message == ResponseFlag.EXECUTED_SUCCESSFULLY
        ):
            success = True

        return f"[dark_green]{flag}[/]" if success else f"[red]{flag}[/]"

    def add_to_table(
        table: Table, response: FSMExecutionResponseLike, prefix: str = ""
    ) -> None:
        executed_command = response.flag == ResponseFlag.EXECUTED_SUCCESSFULLY

        table.add_row(
            prefix + response.name,
            bool_to_success(response.flag, message_type=ResponseFlag),
            (
                bool_to_success(response.fsm_flag, message_type=FSMResponseFlag)
                if executed_command
                else "[red]NA[/]"
            ),
        )
        for child_response in sorted(response.children, key=lambda c: c.name):
            add_to_table(table, child_response, "  " + prefix)

    add_to_table(t, result)

    execution_report_table = format_table_width(obj, t, False)
    obj.print(execution_report_table, soft_wrap=True)

    status_table = format_table_width(obj, render_status_table(obj), False)
    obj.print(status_table, soft_wrap=True)

    obj.print_status_summary()

validate_and_format_fsm_arguments(arguments, command_arguments)

Validates and formats the arguments passed to an FSM command based on the command's argument descriptions.

Parameters:

Name	Type	Description	Default
arguments	dict	A dictionary of argument names and their values passed to the command.	required
command_arguments	list	A list of Argument descriptions for the command.	required

Returns:

Name	Type	Description
dict	dict[str, Any]	A dictionary of argument names and their formatted values, ready to be sent to the controller.

Raises:

Type	Description
ArgumentException	If there is an issue with the arguments (missing, duplicate, invalid type, or unhandled type)

Source code in drunc/controller/interface/shell_utils.py

def validate_and_format_fsm_arguments(
    arguments: dict[str, int | bool | str | float | None] | None,
    command_arguments: list[Argument],
) -> dict[str, any_pb2.Any]:
    """
    Validates and formats the arguments passed to an FSM command based on the command's
    argument descriptions.

    Args:
        arguments (dict): A dictionary of argument names and their values passed to the command.
        command_arguments (list): A list of Argument descriptions for the command.

    Returns:
        dict: A dictionary of argument names and their formatted values, ready to be sent to the controller.

    Raises:
        ArgumentException: If there is an issue with the arguments (missing, duplicate, invalid type, or unhandled type)
    """
    # If the argument dict is empty, don't bother trying to read it
    if not arguments:
        return {}

    # Define the output dict that will be sent to the controller, with argument names
    # and their formatted values
    out_dict: dict[str, any_pb2.Any] = {}

    # Strip out any arguments that are None, as they are considered not passed, and will
    # be set to default values if they exist, or raise an error if they are mandatory
    # without default value
    arguments: dict[str, int | bool | str | float] = {
        k: v for k, v in arguments.items() if v is not None
    }

    # Iterate over the command's argument descriptions, validate the passed arguments,
    # and format them to be sent to the controller
    for argument_desc in command_arguments:  #  type: Argument
        aname: str = argument_desc.name
        atype: str = Argument.Type.Name(argument_desc.type)

        # Check for duplicate arguments
        if aname in out_dict:
            raise DuplicateArgument(aname)

        # Check for missing mandatory arguments
        if (
            argument_desc.presence == Argument.Presence.MANDATORY
            and aname not in arguments
        ):
            raise MissingArgument(aname, atype)

        # If the argument is not passed, and it has a default value, use the default value
        value: str | int | float | bool | None = arguments.get(aname)
        if value is None:
            if argument_desc.HasField("default_value"):
                out_dict[aname] = argument_desc.default_value
            continue

        # Convert the argument value to the appropriate type based on the argument
        # description, and format it to be sent to the controller
        match argument_desc.type:
            case Argument.Type.INT:
                try:
                    value = int(value)
                except Exception as e:
                    raise InvalidArgumentType(aname, value, atype) from e
                value = int_msg(value=value)
            case Argument.Type.FLOAT:
                try:
                    value = float(value)
                except Exception as e:
                    raise InvalidArgumentType(aname, value, atype) from e
                value = float_msg(value=value)
            case Argument.Type.STRING:
                value = string_msg(value=value)
            case Argument.Type.BOOL:
                bvalue = value  # .lower() in ['true', '1', 't', 'y', 'yes', 'yeah', 'yup', 'certainly']
                try:
                    value = bool_msg(value=bvalue)
                except Exception as e:
                    raise InvalidArgumentType(aname, value, atype) from e
            case _:
                try:
                    pretty_type = Argument.Type.Name(argument_desc.type)
                except:
                    pretty_type = argument_desc.type
                raise UnhandledArgumentType(argument_desc.name, pretty_type)
        out_dict[aname] = pack_to_any(value)

    return out_dict