Store

class PreservingLinker(Linker[StatementEntity]):
    """A linker that hands back the canonical id a statement already carries.

    Every backend writer stamps `canonical_id` from the store's linker
    (`nomenklatura.store.sql.SQLWriter.add_statement` and its siblings), so a
    store whose linker knows nothing collapses a resolved statement back onto
    its entity id. Point the store at this linker instead and
    [`preserve`][ftmq.store.base.PreservingLinker.preserve] a statement before
    handing it to the writer: the stamping then reproduces what the statement
    already says.

    Any other id maps to itself - the memory and leveldb writers also resolve
    entity-typed *values* through the linker for their inverted index, and
    those must not be answered with the subject's canonical id.
    """

    def __init__(self) -> None:
        super().__init__({})
        self._stmt: Statement | None = None

    def preserve(self, stmt: Statement) -> None:
        """Resolve this statement's entity id to the canonical id it carries."""
        self._stmt = stmt

    def get_canonical(self, entity_id: str) -> str:
        if self._stmt is not None and entity_id == self._stmt.entity_id:
            return self._stmt.canonical_id
        return entity_id

class Store(nk.Store[Dataset, StatementEntity]):
    """
    Feature add-ons to `nomenklatura.store.Store`
    """

    def __init__(
        self,
        dataset: Dataset | str | None = None,
        linker: Linker | None = None,
        cast_types: bool = True,
        **kwargs,
    ) -> None:
        """
        Initialize a store. This should be called via
        [`get_store`][ftmq.store.get_store]

        Args:
            dataset: A `followthemoney.Dataset` instance to limit the scope to
            linker: A `nomenklatura.Linker` instance with linked / deduped data
            cast_types: Normalize statement values on write (see
                [`ftmq.statements`][ftmq.statements])
        """
        # An unscoped store (no explicit `dataset`) implicitly spans every
        # dataset present in the backend. nomenklatura scopes a view to
        # `dataset.leaf_names`, so without this the store would only surface
        # entities literally tagged `dataset="default"`. Resolved lazily (see
        # `scope`) so opening a store never queries the backend.
        self._implicit_scope = dataset is None
        self.cast_types = cast_types
        linker = linker or get_resolver(kwargs.get("uri"))
        super().__init__(dataset=ensure_dataset(dataset), linker=linker, **kwargs)

    def writer(self, *args: Any, **kwargs: Any) -> Writer:
        """The backend writer, normalizing statement values on the way in.

        Values are cast into the canonical format of their property type (see
        [`ftmq.statements`][ftmq.statements]); values that don't parse are
        passed through unchanged (`ftmq statements cast-types --drop-invalid`
        cleans those out of an existing dump). Disable with the store's
        `cast_types=False`.
        """
        return self.casting_writer(super().writer(*args, **kwargs))

    def casting_writer(self, writer: Writer) -> Writer:
        """Rebless a backend writer so it casts statement values on write (see
        [`writer`][ftmq.store.base.Store.writer]). A store that builds its
        writer itself has to route it through here."""
        if self.cast_types:
            cls: type[Any] = type(writer)
            writer.__class__ = _casting_writer(cls)
        return writer

    def get_scope(self) -> Dataset:
        """
        Return implicit `Dataset` computed from current datasets in store
        """
        raise NotImplementedError

    @property
    def scope(self) -> Dataset:
        """The effective read scope: the store's explicit `dataset`, or all
        datasets present in the backend when it was opened without one."""
        return self.get_scope() if self._implicit_scope else self.dataset

    def view(self, scope: Dataset | None = None, external: bool = False) -> "View":
        raise NotImplementedError

    def default_view(self, external: bool = False) -> "View":
        return self.view(self.scope, external)

    def statements(self, dataset: str | Dataset | None = None) -> Statements:
        """
        Iterate the raw statements in this store, as they are stored.

        Unlike [`iterate`][ftmq.store.base.Store.iterate], which reads
        *entities*, this yields the stored rows unchanged - assembling an
        entity rewrites entity-typed values to their canonical ids and
        synthesizes its own `id` statement, so a dump taken that way would
        carry statement ids that aren't in the store and would load back as new
        rows instead of updating the existing ones. External statements are
        included; the order is unspecified.

        Only the SQL family of backends implements this.

        Args:
            dataset: `Dataset` instance or name to limit scope to

        Yields:
            Generator of `followthemoney.Statement`
        """
        raise NotImplementedError

    def iterate(self, dataset: str | Dataset | None = None) -> StatementEntities:
        """
        Iterate all the entities, optional filter for a dataset.

        Args:
            dataset: `Dataset` instance or name to limit scope to

        Yields:
            Generator of `nomenklatura.entity.CompositeEntity`
        """
        if dataset is not None:
            view = self.view(ensure_dataset(dataset))
        else:
            view = self.default_view()
        yield from view.entities()

class View(nk.View[Dataset, StatementEntity]):
    """
    Feature add-ons to `nomenklatura.store.base.View`
    """

    def query(self, query: Query | None = None) -> StatementEntities:
        """
        Get the entities of a store, optionally filtered by a
        [`Query`][ftmq.Query] object.

        Args:
            query: The Query filter object

        Yields:
            Generator of `followthemoney.StatementEntity`
        """
        view = self.store.view(self.scope)
        if query:
            yield from query.apply_iter(view.entities())
        else:
            yield from view.entities()

    def get_adjacents(
        self, proxies: Iterable[StatementEntity], inverted: bool | None = False
    ) -> set[StatementEntity]:
        seen: set[StatementEntity] = set()
        for proxy in proxies:
            for _, adjacent in self.get_adjacent(proxy, inverted=bool(inverted)):
                if adjacent.id not in seen:
                    seen.add(adjacent)
        return seen

    def stats(self, query: Query | None = None) -> DatasetStats:
        c = Collector()
        cov = c.collect_many(self.query(query))
        return cov

    def count(self, query: Query | None = None) -> int:
        return self.stats(query).entity_count or 0

    def aggregations(self, query: Query) -> AggregatorResult | None:
        if not query.aggregations:
            return
        _ = [x for x in self.query(query)]
        if query.aggregator:
            res = dict(query.aggregator.result)
            return res