aci-sim/aci_sim/mit/store.py

"""MITStore — DN-keyed in-memory object store for ACI MOs."""

from __future__ import annotations
import copy

from aci_sim.mit.mo import MO
from aci_sim.mit.dn import parent_dn


class MITStore:
    """Maintains a ``dn -> MO`` map and a ``parent_dn -> [child dn]`` index.

    Iteration order is deterministic (insertion order of DNs).
    The store is deep-copyable for baseline/snapshot workflows.
    """

    def __init__(self) -> None:
        self._mos: dict[str, MO] = {}               # dn → MO
        self._children: dict[str, list[str]] = {}   # parent_dn → [child dn]

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _register(self, dn: str) -> None:
        """Register *dn* AND ensure its full ancestor chain is linked (idempotent).

        Walking the whole chain (not just the immediate parent) lets :meth:`subtree`
        traverse through structural intermediates that were never added as MOs — e.g.
        a deep ``epmMacEp`` at ``…/sys/ctx-[…]/bd-[…]/db-ep/mac-…`` whose ctx/bd/db-ep
        containers have no MO of their own. Without this, ``subtree(node/sys)`` could
        not reach it.
        """
        self._children.setdefault(dn, [])
        cur = dn
        while True:
            pdn = parent_dn(cur)
            if not pdn or pdn == cur:
                break
            siblings = self._children.setdefault(pdn, [])
            if cur not in siblings:
                siblings.append(cur)
            cur = pdn

    def _remove_subtree(self, dn: str) -> None:
        """Remove *dn* and all its descendants from both maps."""
        for child_dn in list(self._children.get(dn, [])):
            self._remove_subtree(child_dn)
        self._mos.pop(dn, None)
        self._children.pop(dn, None)
        pdn = parent_dn(dn)
        siblings = self._children.get(pdn)
        if siblings is not None:
            try:
                siblings.remove(dn)
            except ValueError:
                pass

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def add(self, mo: MO) -> None:
        """Insert *mo* (and recursively its ``.children``) without merging.

        Stores a *childless* copy — hierarchy lives only in the parent index
        (same invariant as :meth:`upsert`). Otherwise a stored tree MO would
        carry a second copy of its subtree via ``MO.to_imdata()``, leaking
        ``children`` onto non-subtree queries and duplicating descendants on
        subtree queries.
        """
        dn = mo.dn
        self._mos[dn] = MO(mo.class_name, **mo.attrs)
        self._register(dn)
        for child in mo.children:
            self.add(child)

    def upsert(self, mo: MO) -> None:
        """Merge *mo*'s attrs into an existing entry or insert it.

        Recurses into ``mo.children``.
        If ``mo.attrs["status"] == "deleted"``, removes the DN and its entire
        subtree instead of merging.
        """
        dn = mo.dn
        if mo.attrs.get("status") == "deleted":
            self._remove_subtree(dn)
            return

        existing = self._mos.get(dn)
        if existing is None:
            new_mo = MO(mo.class_name, **mo.attrs)
            self._mos[dn] = new_mo
            self._register(dn)
        else:
            existing.attrs.update(mo.attrs)

        for child in mo.children:
            self.upsert(child)

    def get(self, dn: str) -> MO | None:
        """Return the MO at *dn*, or ``None`` if not present."""
        return self._mos.get(dn)

    def delete(self, dn: str) -> None:
        """Remove *dn* and its entire subtree from the store.

        Idempotent: deleting a *dn* that is not present (or was already
        removed) is a silent no-op, matching ``upsert``'s ``status="deleted"``
        branch (also unconditional — no existence check) and real APIC's
        DELETE /api/mo/{dn}.json used by cisco.aci's ``state=absent`` path
        (PR-9). Public wrapper around the same ``_remove_subtree`` helper
        ``upsert`` already uses, so callers outside this module (the DELETE
        HTTP route) don't need to reach into a private method.
        """
        self._remove_subtree(dn)

    def children(self, dn: str, classes: set[str] | None = None) -> list[MO]:
        """Return direct children of *dn*, optionally filtered by class names."""
        result: list[MO] = []
        for cdn in self._children.get(dn, []):
            mo = self._mos.get(cdn)
            if mo is None:
                continue
            if classes is None or mo.class_name in classes:
                result.append(mo)
        return result

    def child_dns(self, dn: str) -> list[str]:
        """Return the raw child-DN list of *dn* in insertion order.

        Includes structural intermediates that have no MO of their own (see
        :meth:`_register`). Used by the query engine to reconstruct the nested
        subtree for ``rsp-subtree=full`` without flattening the hierarchy.
        """
        return list(self._children.get(dn, []))

    def subtree(self, dn: str, classes: set[str] | None = None) -> list[MO]:
        """Return all descendants of *dn* (not *dn* itself), depth-first.

        If *classes* is given, only include MOs whose class_name is in the set;
        but always recurse into ALL children regardless (mirrors APIC behaviour
        where a class filter on rsp-subtree-class traverses the full tree).
        """
        acc: list[MO] = []
        self._collect_subtree(dn, classes, acc)
        return acc

    def _collect_subtree(self, dn: str, classes: set[str] | None, acc: list[MO]) -> None:
        for cdn in self._children.get(dn, []):
            mo = self._mos.get(cdn)
            if mo is not None and (classes is None or mo.class_name in classes):
                acc.append(mo)
            # Recurse even through MO-less structural intermediates so deeply
            # nested descendants (e.g. epm* under uninstantiated containers) are reached.
            self._collect_subtree(cdn, classes, acc)

    def by_class(self, cls: str) -> list[MO]:
        """Return all MOs whose ``class_name == cls``, in insertion order."""
        return [mo for mo in self._mos.values() if mo.class_name == cls]

    def all(self) -> list[MO]:
        """Return all MOs in insertion order."""
        return list(self._mos.values())

    # ------------------------------------------------------------------
    # Deep-copy support (for baseline / snapshot)
    # ------------------------------------------------------------------

    def __deepcopy__(self, memo: dict) -> MITStore:
        new = MITStore.__new__(MITStore)
        new._mos = copy.deepcopy(self._mos, memo)
        new._children = copy.deepcopy(self._children, memo)
        return new