Add MregCliClient to encapsulate CLI client methods

mreg_cli/client.py

@@ -6,24 +6,263 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Literal, overload
 
 from mreg_api import MregClient
 from mreg_api.cache import CacheConfig
+from mreg_api.events import Event, EventKind, EventLevel
+from mreg_api.models import Atom, Host, Network, Role
+from typing_extensions import override
 
 from mreg_cli.__about__ import __version__
+from mreg_cli.exceptions import EntityNotFound as CliEntityNotFound
+from mreg_cli.exceptions import TooManyResults
 
 if TYPE_CHECKING:
     from mreg_cli.config import MregCliConfig
 
-_client: MregClient | None = None
+_client: MregCliClient | None = None
 
 
-def init_client(config: MregCliConfig) -> MregClient:
-    """Initialize the global client instance given the config."""
+def _fail_on_truncate(event: Event) -> None:
+    """Event hook that raises TooManyResults if the event is a truncation event."""
+    if event.kind == EventKind.TRUNCATION:
+        raise TooManyResults(f"{event.message} Refine your search.")
+
+
+def _record_events(event: Event) -> None:
+    """Event hook for MregClient to record events in the OutputManager."""
+    from mreg_cli.outputmanager import OutputManager  # noqa: PLC0415
+
+    om = OutputManager()
+    if om.is_suppressing_events():
+        return
+
+    if event.level >= EventLevel.WARNING:
+        om.add_warning(event.message)
+    elif event.kind == EventKind.MUTATION:
+        om.add_ok(event.message)
+    else:
+        om.add_line(event.message)
+
+
+class _AutoListenersMeta(type):
+    """Metaclass that automatically adds mreg-api event listeners to MregCliClient instances.
+
+    Ensures some key objectives:
+    1. Listeners are added automatically without requiring consumers to call a separate method.
+    2. Event listeners are added only once per instance.
+    3. Removes the need to override __init__ in MregCliClient.
+    """
+
+    @override
+    def __call__(cls, *args: object, **kwargs: object) -> MregCliClient:
+        obj: MregCliClient = super().__call__(*args, **kwargs)
+        obj.events.subscribe(_fail_on_truncate)  # fail fast - don't double print
+        obj.events.subscribe(_record_events)
+        return obj
+
+
+class MregCliClient(MregClient, metaclass=_AutoListenersMeta):
+    """Subclass of MregClient that adds CLI-specific functionality."""
+
+    @overload
+    def resolve_policy(
+        self,
+        name: str,
+        *,
+        required: Literal[False],
+    ) -> Atom | Role | None: ...
+
+    @overload
+    def resolve_policy(
+        self,
+        name: str,
+        *,
+        required: Literal[True] = ...,
+    ) -> Atom | Role: ...
+
+    @overload
+    def resolve_policy(
+        self,
+        name: str,
+        *,
+        required: bool = ...,
+    ) -> Atom | Role | None: ...
+
+    def resolve_policy(
+        self,
+        name: str,
+        *,
+        required: bool = True,
+    ) -> Atom | Role | None:
+        """Resolve a host policy name to an Atom or Role.
+
+        Resolution order: Atom name → Role name.
+        """
+        atom = self.atom.get_by_name(name, required=False)
+        if atom is not None:
+            return atom
+        role = self.role.get_by_name(name, required=False)
+        if role is not None:
+            return role
+        if required:
+            raise CliEntityNotFound(f"No atom or role named {name!r}.")
+        return None
+
+    @overload
+    def resolve_host(
+        self,
+        identifier: str | int,
+        *,
+        required: Literal[False],
+        inform_if_cname: bool = ...,
+    ) -> Host | None: ...
+
+    @overload
+    def resolve_host(
+        self,
+        identifier: str | int,
+        *,
+        required: Literal[True] = ...,
+        inform_if_cname: bool = ...,
+    ) -> Host: ...
+
+    @overload
+    def resolve_host(
+        self,
+        identifier: str | int,
+        *,
+        required: bool = ...,
+        inform_if_cname: bool = ...,
+    ) -> Host | None: ...
+
+    def resolve_host(
+        self,
+        identifier: str | int,
+        *,
+        required: bool = True,
+        inform_if_cname: bool = False,
+    ) -> Host | None:
+        """Resolve a host by id, IP, MAC, name, or CNAME target.
+
+        Resolution order: numeric id → IP → MAC → name → CNAME target.
+        """
+        from mreg_cli.outputmanager import OutputManager  # noqa: PLC0415
+
+        # We got passed an integer, assume host ID
+        if isinstance(identifier, int):
+            return self.host.get_by_id(identifier, required=required)
+
+        getters: list[Callable[[str], Host | None]] = []
+        ident_str = str(identifier)
+
+        # Try by ID first if the identifier is a decimal number
+        if ident_str.isdecimal():
+            getters.append(lambda s: self.host.get_by_id(int(s), required=False))
+
+        # Try by IP → MAC → name
+        getters.extend(
+            [
+                lambda s: self.host.get_by_ip(s, required=False),
+                lambda s: self.host.get_by_mac(s, required=False),
+                lambda s: self.host.get_by_name(s, required=False),
+            ]
+        )
+
+        for getter in getters:
+            try:
+                host = getter(ident_str)
+                if host is not None:
+                    return host
+            except Exception:
+                pass
+
+        # Fall back to CNAME
+        cname = self.cname.get_by_name(ident_str, required=False)
+        if cname is not None:
+            host = self.host.get_by_id(cname.host, required=False)
+            # NOTE: should it be an error if CNAME has host ID that doesn't exist? Probably.
+            # At the very least produce a warning if CNAME host ID cannot be resolved.
+            if host is not None:
+                if inform_if_cname:
+                    OutputManager().add_line(f"{ident_str!r} is a CNAME for {host.name!r}")
+                return host
+
+        if required:
+            raise CliEntityNotFound(f"Host {identifier!r} not found.")
+
+        return None
+
+    @overload
+    def resolve_network(
+        self,
+        identifier: str,
+        *,
+        required: Literal[False],
+    ) -> Network | None: ...
+
+    @overload
+    def resolve_network(
+        self,
+        identifier: str,
+        *,
+        required: Literal[True] = ...,
+    ) -> Network: ...
+
+    @overload
+    def resolve_network(
+        self,
+        identifier: str,
+        *,
+        required: bool = ...,
+    ) -> Network | None: ...
+
+    def resolve_network(
+        self,
+        identifier: str,
+        *,
+        required: bool = True,
+    ) -> Network | None:
+        """Resolve a network by IP address, CIDR notation, or numeric id.
+
+        Resolution order: IP address → CIDR / network address → numeric id.
+        """
+        network: Network | None = None
+        try:
+            # Try as IP first
+            network = self.network.get_by_ip(identifier, required=False)
+        except Exception:
+            pass
+
+        if network is None:
+            try:
+                # Try as CIDR / network address
+                network = self.network.get(identifier, required=False)
+            except Exception:
+                pass
+
+        if network is None and identifier.isdecimal():
+            try:
+                network = self.network.first(id=int(identifier), required=False)
+                # network = self.network.get(int(identifier), required=False)
+            except Exception:
+                pass
+
+        if network is None and required:
+            raise CliEntityNotFound(f"Network {identifier!r} not found.")
+        return network
+
+
+def init_client(config: MregCliConfig) -> MregCliClient:
+    """Initialize the global client instance given the config.
+
+    Idempotent: calling this multiple times will return the same client instance.
+    """
     global _client
     if _client is None:
-        _client = MregClient(
+        _client = MregCliClient(
             url=config.url,
             domain=config.domain,
             timeout=config.http_timeout,
@@ -46,11 +285,11 @@ def _reset_client() -> None:  # pyright: ignore[reportUnusedFunction] # in tests
     _client = None
 
 
-def get_client() -> MregClient:
+def get_client() -> MregCliClient:
     """Return the global client instance.
 
     :raises RuntimeError: If init_client() has not been called yet.
     """
     if _client is None:
-        raise RuntimeError("MregClient not initialized — call init_client() first.")
+        raise RuntimeError("MregCliClient not initialized — call init_client() first.")
     return _client

mreg_cli/commands/dhcp.py

@@ -22,7 +22,6 @@
 )
 from mreg_cli.outputmanager import OutputManager
 from mreg_cli.types import Flag
-from mreg_cli.utilities.resolution import resolve_host
 
 logger = logging.getLogger(__name__)
 
@@ -167,7 +166,7 @@ def assoc(args: argparse.Namespace) -> None:
     # Parse Name as an IP address first, if not found, resolve as a host name
     ipaddress = ipaddress_from_ip_arg(name, client)
     if not ipaddress:
-        host = resolve_host(client, name)
+        host = client.resolve_host(name)
         ipaddress = get_associable_ip(host, client)
 
     client.ipaddress.associate_mac(ipaddress, mac, force=force)
@@ -200,7 +199,7 @@ def disassoc(args: argparse.Namespace) -> None:
 
     # Name is not an IP -> resolve as host
     if not ipaddress:
-        host = resolve_host(client, name)
+        host = client.resolve_host(name)
         # Check if name itself looks like a MAC address and find the matching IP.
         # This replicates the old host.has_ip_with_mac(mac) logic.
         if mac := MacAddress.parse(name):

mreg_cli/commands/group.py

@@ -13,7 +13,6 @@
 from mreg_cli.output.history import output_hostgroup_history
 from mreg_cli.outputmanager import OutputManager
 from mreg_cli.types import Flag
-from mreg_cli.utilities.resolution import resolve_host
 
 command_registry = CommandRegistry()
 
@@ -226,7 +225,7 @@ def host_add(args: argparse.Namespace) -> None:
     hostgroup = client.hostgroup.get_by_name(args.group)
 
     for name in args.hosts:
-        host = resolve_host(client, name)
+        host = client.resolve_host(name)
         fqname = host.name
         client.hostgroup.add_host(hostgroup, fqname)
         OutputManager().add_ok(f"Added host {fqname!r} to {args.group!r}")
@@ -251,7 +250,7 @@ def host_remove(args: argparse.Namespace) -> None:
 
     to_remove: set[str] = set()
     for name in args.hosts:
-        host = resolve_host(client, name)
+        host = client.resolve_host(name)
         fqname = host.name
         if not hostgroup.has_host(fqname):
             raise EntityNotFound(f"Host {name!r} ({fqname!r}) not a member in {args.group!r}")
@@ -282,7 +281,7 @@ def host_list(args: argparse.Namespace) -> None:
     :param args: argparse.Namespace (host, traverse-hostgroups)
     """
     client = get_client()
-    host = resolve_host(client, args.host)
+    host = client.resolve_host(args.host)
     hostgroups = client.hostgroup.list_by_host(host, traverse=args.traverse_hostgroups)
     output_hostgroups(hostgroups, multiline=True)