nomenklatura-mpt 4.1.9__py3-none-any.whl

nomenklatura/store/base.py

@@ -0,0 +1,130 @@
+from types import TracebackType
+from typing import Optional, Generator, List, Tuple, Generic, Type, cast
+from followthemoney import Schema, registry, Property, DS, Statement
+from followthemoney import StatementEntity, SE
+from followthemoney.statement.util import get_prop_type
+
+from nomenklatura.resolver import Linker, StrIdent
+
+
+class Store(Generic[DS, SE]):
+    """A data storage and retrieval mechanism for statement-based entity data.
+    Essentially, this is a triple store which can be implemented using various
+    backends."""
+
+    def __init__(self, dataset: DS, linker: Linker[SE]):
+        self.dataset = dataset
+        self.linker = linker
+        self.entity_class = cast(Type[SE], StatementEntity)
+
+    def writer(self) -> "Writer[DS, SE]":
+        raise NotImplementedError()
+
+    def view(self, scope: DS, external: bool = False) -> "View[DS, SE]":
+        raise NotImplementedError()
+
+    def default_view(self, external: bool = False) -> "View[DS, SE]":
+        return self.view(self.dataset, external=external)
+
+    def assemble(self, statements: List[Statement]) -> Optional[SE]:
+        if not len(statements):
+            return None
+        for stmt in statements:
+            if get_prop_type(stmt.schema, stmt.prop) == registry.entity.name:
+                stmt.value = self.linker.get_canonical(stmt.value)
+        entity = self.entity_class.from_statements(self.dataset, statements)
+        if entity.id is not None:
+            entity.extra_referents.update(self.linker.get_referents(entity.id))
+        return entity
+
+    def update(self, id: StrIdent) -> None:
+        canonical_id = self.linker.get_canonical(id)
+        with self.writer() as writer:
+            for referent in self.linker.get_referents(canonical_id):
+                for stmt in writer.pop(referent):
+                    stmt.canonical_id = canonical_id
+                    writer.add_statement(stmt)
+
+    def close(self) -> None:
+        pass
+
+    def __repr__(self) -> str:
+        return f"<{type(self).__name__}({self.dataset.name!r})>"
+
+
+class Writer(Generic[DS, SE]):
+    """Bulk writing operations."""
+
+    def __init__(self, store: Store[DS, SE]):
+        self.store = store
+
+    def add_statement(self, stmt: Statement) -> None:
+        raise NotImplementedError()
+
+    def add_entity(self, entity: SE) -> None:
+        for stmt in entity.statements:
+            self.add_statement(stmt)
+
+    def pop(self, entity_id: str) -> List[Statement]:
+        raise NotImplementedError()
+
+    def flush(self) -> None:
+        pass
+
+    def close(self) -> None:
+        self.store.close()
+
+    def __enter__(self) -> "Writer[DS, SE]":
+        return self
+
+    def __exit__(
+        self,
+        type: Optional[Type[BaseException]],
+        value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
+        self.flush()
+
+    def __repr__(self) -> str:
+        return f"<{type(self).__name__}({self.store!r})>"
+
+
+class View(Generic[DS, SE]):
+    def __init__(self, store: Store[DS, SE], scope: DS, external: bool = False):
+        self.store = store
+        self.scope = scope
+        self.dataset_names = scope.leaf_names
+        self.external = external
+
+    def has_entity(self, id: str) -> bool:
+        raise NotImplementedError()
+
+    def get_entity(self, id: str) -> Optional[SE]:
+        raise NotImplementedError()
+
+    def get_inverted(self, id: str) -> Generator[Tuple[Property, SE], None, None]:
+        raise NotImplementedError()
+
+    def get_adjacent(
+        self, entity: SE, inverted: bool = True
+    ) -> Generator[Tuple[Property, SE], None, None]:
+        for prop, value in entity.itervalues():
+            if prop.type == registry.entity:
+                child = self.get_entity(value)
+                if child is not None:
+                    yield prop, child
+
+        if inverted and entity.id is not None:
+            for prop, adjacent in self.get_inverted(entity.id):
+                yield prop, adjacent
+
+    def entities(self, include_schemata: List[Schema] = []) -> Generator[SE, None, None]:
+        """Iterate over all entities in the view.
+
+        If `include_schemata` is provided, only entities of the provided schemata will be returned.
+        Note that `schemata` will not be expanded via "is_a" relationships."""
+
+        raise NotImplementedError()
+
+    def __repr__(self) -> str:
+        return f"<{type(self).__name__}({self.scope.name!r})>"

nomenklatura/store/level.py

@@ -0,0 +1,272 @@
+#
+# LevelDB-based store for Nomenklatura.
+# A lot of the code in this module is extremely performance-sensitive, so it is unrolled and
+# doesn't use helper functions in some places where it would otherwise be more readable.
+#
+# Specific examples:
+# * Not calling a helper to byte-encode values.
+# * Not having a helper method for building entities.
+import gc
+import orjson
+import logging
+from pathlib import Path
+from typing import Any, Generator, List, Optional, Set, Tuple
+from rigour.env import ENCODING as E
+
+import plyvel  # type: ignore
+from followthemoney import model, DS, SE, Schema, registry, Property, Statement
+from followthemoney.exc import InvalidData
+from followthemoney.statement.util import get_prop_type
+
+from nomenklatura.resolver import Linker
+from nomenklatura.store.base import Store, View, Writer
+
+log = logging.getLogger(__name__)
+MAX_OPEN_FILES = 1000
+
+
+def unpack_statement(
+    keys: List[str],
+    data: bytes,
+) -> Statement:
+    _, canonical_id, ext, dataset, schema, stmt_id = keys
+    (
+        entity_id,
+        prop,
+        value,
+        lang,
+        original_value,
+        origin,
+        first_seen,
+        last_seen,
+    ) = orjson.loads(data)
+    return Statement(
+        id=stmt_id,
+        entity_id=entity_id,
+        prop=prop,
+        schema=schema,
+        value=value,
+        lang=None if lang == 0 else lang,
+        dataset=dataset,
+        original_value=None if original_value == 0 else original_value,
+        origin=None if origin == 0 else origin,
+        first_seen=first_seen,
+        last_seen=last_seen,
+        canonical_id=canonical_id,
+        external=ext == "x",
+    )
+
+
+class LevelDBStore(Store[DS, SE]):
+    def __init__(self, dataset: DS, linker: Linker[SE], path: Path):
+        super().__init__(dataset, linker)
+        self.path = path
+        self.db = plyvel.DB(
+            path.as_posix(),
+            create_if_missing=True,
+            max_open_files=MAX_OPEN_FILES,
+        )
+
+    def optimize(self) -> None:
+        """Optimize the database by compacting it."""
+        self.db.compact_range()
+        self.db.close()
+        gc.collect()
+        self.db = plyvel.DB(
+            self.path.as_posix(),
+            create_if_missing=False,
+            max_open_files=MAX_OPEN_FILES,
+        )
+        log.info("Optimized LevelDB at %s", self.path)
+
+    def writer(self) -> Writer[DS, SE]:
+        return LevelDBWriter(self)
+
+    def view(self, scope: DS, external: bool = False) -> View[DS, SE]:
+        return LevelDBView(self, scope, external=external)
+
+    def close(self) -> None:
+        self.db.close()
+
+
+class LevelDBWriter(Writer[DS, SE]):
+    BATCH_STATEMENTS = 100_000
+
+    def __init__(self, store: LevelDBStore[DS, SE]):
+        self.store: LevelDBStore[DS, SE] = store
+        self.batch: Optional[Any] = None
+        self.batch_size = 0
+
+    def flush(self) -> None:
+        if self.batch is not None:
+            self.batch.write()
+        self.batch = None
+        self.batch_size = 0
+
+    def add_statement(self, stmt: Statement) -> None:
+        if stmt.entity_id is None:
+            return
+        if self.batch_size >= self.BATCH_STATEMENTS:
+            self.flush()
+        if self.batch is None:
+            self.batch = self.store.db.write_batch()
+        canonical_id = self.store.linker.get_canonical(stmt.entity_id)
+        stmt.canonical_id = canonical_id
+
+        ext = "x" if stmt.external else ""
+        key = f"s:{canonical_id}:{ext}:{stmt.dataset}:{stmt.schema}:{stmt.id}".encode(E)
+        values = (
+            stmt.entity_id,
+            stmt.prop,
+            stmt.value,
+            stmt.lang or 0,
+            stmt.original_value or 0,
+            stmt.origin or 0,
+            stmt.first_seen,
+            stmt.last_seen,
+        )
+        data = orjson.dumps(values)
+        self.batch.put(key, data)
+        if get_prop_type(stmt.schema, stmt.prop) == registry.entity.name:
+            vc = self.store.linker.get_canonical(stmt.value)
+            key = f"i:{vc}:{stmt.canonical_id}".encode(E)
+            self.batch.put(key, b"")
+
+        self.batch_size += 1
+
+    def pop(self, entity_id: str) -> List[Statement]:
+        if self.batch_size >= self.BATCH_STATEMENTS:
+            self.flush()
+        if self.batch is None:
+            self.batch = self.store.db.write_batch()
+        statements: List[Statement] = []
+        datasets: Set[str] = set()
+        prefix = f"s:{entity_id}:".encode(E)
+        with self.store.db.iterator(prefix=prefix) as it:
+            for k, v in it:
+                self.batch.delete(k)
+                stmt = unpack_statement(k.decode(E).split(":"), v)
+                statements.append(stmt)
+                datasets.add(stmt.dataset)
+
+                if stmt.prop_type == registry.entity.name:
+                    vc = self.store.linker.get_canonical(stmt.value)
+                    self.batch.delete(f"i:{vc}:{entity_id}".encode(E))
+        return list(statements)
+
+
+class LevelDBView(View[DS, SE]):
+    def __init__(
+        self, store: LevelDBStore[DS, SE], scope: DS, external: bool = False
+    ) -> None:
+        super().__init__(store, scope, external=external)
+        self.store: LevelDBStore[DS, SE] = store
+        self.dataset_names: Set[str] = set(scope.dataset_names)
+
+    def has_entity(self, id: str) -> bool:
+        prefix = f"s:{id}:".encode(E)
+        with self.store.db.iterator(prefix=prefix, include_value=False) as it:
+            for v in it:
+                _, _, ext, dataset, _, _ = v.decode(E).split(":")
+                if dataset not in self.dataset_names:
+                    continue
+                if ext == "x" and not self.external:
+                    continue
+                return True
+        return False
+
+    def get_entity(self, id: str) -> Optional[SE]:
+        statements: List[Statement] = []
+        prefix = f"s:{id}:".encode(E)
+        with self.store.db.iterator(prefix=prefix) as it:
+            for k, v in it:
+                keys = k.decode(E).split(":")
+                _, _, ext, dataset, _, _ = keys
+                if dataset not in self.dataset_names:
+                    continue
+                if ext == "x" and not self.external:
+                    continue
+                statements.append(unpack_statement(keys, v))
+        return self.store.assemble(statements)
+
+    def get_inverted(self, id: str) -> Generator[Tuple[Property, SE], None, None]:
+        prefix = f"i:{id}:".encode(E)
+        with self.store.db.iterator(prefix=prefix, include_value=False) as it:
+            for k in it:
+                _, _, ref = k.decode(E).split(":")
+                entity = self.get_entity(ref)
+                if entity is None:
+                    continue
+                for prop, value in entity.itervalues():
+                    if value == id and prop.reverse is not None:
+                        yield prop.reverse, entity
+
+    def entities(
+        self, include_schemata: Optional[List[Schema]] = None
+    ) -> Generator[SE, None, None]:
+        with self.store.db.iterator(prefix=b"s:", fill_cache=False) as it:
+            current_id: Optional[str] = None
+            current_schema: Optional[Schema] = None
+            current_fail: bool = False
+            statements: List[Statement] = []
+            for k, v in it:
+                keys = k.decode(E).split(":")
+                _, canonical_id, ext, dataset, schema, _ = keys
+
+                if ext == "x" and not self.external:
+                    continue
+                if dataset not in self.dataset_names:
+                    continue
+
+                # If we're seeing a new canonical ID, yield the previous entity
+                if canonical_id != current_id:
+                    if (
+                        include_schemata is not None
+                        and current_schema not in include_schemata
+                    ):
+                        statements = []
+                    if len(statements) > 0 and not current_fail:
+                        entity = self.store.assemble(statements)
+                        if entity is not None:
+                            yield entity
+                    current_id = canonical_id
+                    current_schema = None
+                    current_fail = False
+                    statements = []
+
+                # If we're not filtering on schemata, we can skip the expensive-ish schema building here
+                # The checking is done by store.assemble() anyway
+                if include_schemata is not None:
+                    if current_schema is None:
+                        current_schema = model.get(schema)
+                        # If the statement is of an unknown schema
+                        if current_schema is None:
+                            log.error("Unknown schema %r: %s", (schema, current_id))
+                            # Mark the entity as failed, but we need to iterate through the rest of the statements
+                            current_fail = True
+                            continue
+                    # If the schema of the statement does not exactly match the schema of the current entity,
+                    # find a common parent schema.
+                    elif current_schema.name != schema:
+                        try:
+                            current_schema = model.common_schema(current_schema, schema)
+                        except InvalidData as inv:
+                            msg = "Invalid schema %s for %r: %s" % (
+                                schema,
+                                current_id,
+                                inv,
+                            )
+                            log.error(msg)
+                            # Mark the entity as failed, but we need to iterate through the rest of the statements
+                            current_fail = True
+                            continue
+
+                statements.append(unpack_statement(keys, v))
+
+            # Handle the last entity at the end of the iterator
+            if include_schemata is not None and current_schema not in include_schemata:
+                statements = []
+            if len(statements) > 0 and not current_fail:
+                entity = self.store.assemble(statements)
+                if entity is not None:
+                    yield entity

nomenklatura/store/memory.py

@@ -0,0 +1,102 @@
+from typing import Dict, Set, List, Optional, Generator, Tuple
+from followthemoney import DS, SE, Schema, registry, Property, Statement
+
+from nomenklatura.store.base import Store, View, Writer
+from nomenklatura.resolver import Linker
+
+
+class MemoryStore(Store[DS, SE]):
+    def __init__(self, dataset: DS, linker: Linker[SE]):
+        super().__init__(dataset, linker)
+        self.stmts: Dict[str, Set[Statement]] = {}
+        self.inverted: Dict[str, Set[str]] = {}
+        self.entities: Dict[str, Set[str]] = {}
+
+    def writer(self) -> Writer[DS, SE]:
+        return MemoryWriter(self)
+
+    def view(self, scope: DS, external: bool = False) -> View[DS, SE]:
+        return MemoryView(self, scope, external=external)
+
+
+class MemoryWriter(Writer[DS, SE]):
+    def __init__(self, store: MemoryStore[DS, SE]):
+        self.store: MemoryStore[DS, SE] = store
+
+    def add_statement(self, stmt: Statement) -> None:
+        if stmt.entity_id is None:
+            return
+        canonical_id = stmt.canonical_id or self.store.linker.get_canonical(
+            stmt.entity_id
+        )
+        if canonical_id not in self.store.stmts:
+            self.store.stmts[canonical_id] = set()
+        self.store.stmts[canonical_id].add(stmt)
+
+        if stmt.dataset not in self.store.entities:
+            self.store.entities[stmt.dataset] = set()
+        self.store.entities[stmt.dataset].add(canonical_id)
+
+        if stmt.prop_type == registry.entity.name:
+            inverted_id = self.store.linker.get_canonical(stmt.value)
+            if inverted_id not in self.store.inverted:
+                self.store.inverted[inverted_id] = set()
+            self.store.inverted[inverted_id].add(canonical_id)
+
+    def pop(self, entity_id: str) -> List[Statement]:
+        statements = self.store.stmts.pop(entity_id, set())
+        for stmt in statements:
+            if stmt.dataset in self.store.entities:
+                self.store.entities[stmt.dataset].discard(entity_id)
+
+            if stmt.prop_type == registry.entity.name:
+                inverted_id = self.store.linker.get_canonical(stmt.value)
+                if inverted_id in self.store.inverted:
+                    self.store.inverted[inverted_id].discard(entity_id)
+
+        return list(statements)
+
+
+class MemoryView(View[DS, SE]):
+    def __init__(
+        self, store: MemoryStore[DS, SE], scope: DS, external: bool = False
+    ) -> None:
+        super().__init__(store, scope, external=external)
+        self.store: MemoryStore[DS, SE] = store
+
+    def has_entity(self, id: str) -> bool:
+        for stmt in self.store.stmts.get(id, []):
+            if self.external is False and stmt.external:
+                continue
+            return True
+        return False
+
+    def get_entity(self, id: str) -> Optional[SE]:
+        if id not in self.store.stmts:
+            return None
+        stmts: List[Statement] = []
+        for stmt in self.store.stmts[id]:
+            if self.external is False and stmt.external:
+                continue
+            stmts.append(stmt)
+        return self.store.assemble(stmts)
+
+    def get_inverted(self, id: str) -> Generator[Tuple[Property, SE], None, None]:
+        for inverted_id in self.store.inverted.get(id, []):
+            entity = self.get_entity(inverted_id)
+            if entity is None:
+                continue
+            for prop, value in entity.itervalues():
+                if value == id and prop.reverse is not None:
+                    yield prop.reverse, entity
+
+    def entities(self, include_schemata: Optional[List[Schema]] = None) -> Generator[SE, None, None]:
+        entity_ids: Set[str] = set()
+        for scope in self.dataset_names:
+            entity_ids.update(self.store.entities.get(scope, []))
+        for entity_id in entity_ids:
+            entity = self.get_entity(entity_id)
+            if entity is not None:
+                if include_schemata is not None and entity.schema not in include_schemata:
+                    continue
+                yield entity

nomenklatura/store/redis_.py

@@ -0,0 +1,131 @@
+from redis.client import Redis, Pipeline
+from typing import Generator, List, Optional, Set, Tuple
+from followthemoney import DS, SE, Schema, registry, Property, Statement
+
+from nomenklatura.kv import get_redis, close_redis, b
+from nomenklatura.resolver import Linker
+from nomenklatura.store.base import Store, View, Writer
+from nomenklatura.store.util import pack_statement, unpack_statement
+
+
+class RedisStore(Store[DS, SE]):
+    def __init__(
+        self,
+        dataset: DS,
+        linker: Linker[SE],
+        db: Optional["Redis[bytes]"] = None,
+    ):
+        super().__init__(dataset, linker)
+        if db is None:
+            db = get_redis()
+        self.db = db
+
+    def writer(self) -> Writer[DS, SE]:
+        return RedisWriter(self)
+
+    def view(self, scope: DS, external: bool = False) -> View[DS, SE]:
+        return RedisView(self, scope, external=external)
+
+    def close(self) -> None:
+        close_redis()
+
+
+class RedisWriter(Writer[DS, SE]):
+    BATCH_STATEMENTS = 100_000
+
+    def __init__(self, store: RedisStore[DS, SE]):
+        self.store: RedisStore[DS, SE] = store
+        self.pipeline: Optional["Pipeline[bytes]"] = None
+        self.batch_size = 0
+
+    def flush(self) -> None:
+        if self.pipeline is not None:
+            self.pipeline.execute()
+        self.pipeline = None
+        self.batch_size = 0
+
+    def add_statement(self, stmt: Statement) -> None:
+        if stmt.entity_id is None:
+            return
+        if self.batch_size >= self.BATCH_STATEMENTS:
+            self.flush()
+        if self.pipeline is None:
+            self.pipeline = self.store.db.pipeline()
+        canonical_id = self.store.linker.get_canonical(stmt.entity_id)
+        stmt.canonical_id = canonical_id
+
+        self.pipeline.sadd(b(f"ds:{stmt.dataset}"), b(canonical_id))
+        key = f"x:{canonical_id}" if stmt.external else f"s:{canonical_id}"
+        self.pipeline.sadd(b(key), pack_statement(stmt))
+        if stmt.prop_type == registry.entity.name:
+            vc = self.store.linker.get_canonical(stmt.value)
+            self.pipeline.sadd(b(f"i:{vc}"), b(canonical_id))
+
+        self.batch_size += 1
+
+    def pop(self, entity_id: str) -> List[Statement]:
+        if self.batch_size >= self.BATCH_STATEMENTS:
+            self.flush()
+        if self.pipeline is None:
+            self.pipeline = self.store.db.pipeline()
+        statements: List[Statement] = []
+        datasets: Set[str] = set()
+        keys = (f"s:{entity_id}", f"x:{entity_id}")
+        for v in self.store.db.sunion(keys):
+            stmt = unpack_statement(v, entity_id, False)  # type: ignore
+            statements.append(stmt)
+            datasets.add(stmt.dataset)
+
+            if stmt.prop_type == registry.entity.name:
+                vc = self.store.linker.get_canonical(stmt.value)
+                self.pipeline.srem(b(f"i:{vc}"), b(entity_id))
+
+        for dataset in datasets:
+            self.pipeline.srem(b(f"ds:{dataset}"), b(entity_id))
+
+        return list(statements)
+
+
+class RedisView(View[DS, SE]):
+    def __init__(
+        self, store: RedisStore[DS, SE], scope: DS, external: bool = False
+    ) -> None:
+        super().__init__(store, scope, external=external)
+        self.store: RedisStore[DS, SE] = store
+
+    def has_entity(self, id: str) -> bool:
+        keys = [b(f"s:{id}")]
+        if self.external:
+            keys.append(b(f"x:{id}"))
+        return self.store.db.exists(*keys) > 0
+
+    def get_entity(self, id: str) -> Optional[SE]:
+        statements: List[Statement] = []
+        keys = [b(f"s:{id}")]
+        if self.external:
+            keys.append(b(f"x:{id}"))
+        for v in self.store.db.sunion(keys):
+            statements.append(unpack_statement(v, id, False))  # type: ignore
+        return self.store.assemble(statements)
+
+    def get_inverted(self, id: str) -> Generator[Tuple[Property, SE], None, None]:
+        for v in self.store.db.smembers(b(f"i:{id}")):
+            entity = self.get_entity(v.decode("utf-8"))
+            if entity is None:
+                continue
+            for prop, value in entity.itervalues():
+                if value == id and prop.reverse is not None:
+                    yield prop.reverse, entity
+
+    def entities(self, include_schemata: Optional[List[Schema]] = None) -> Generator[SE, None, None]:
+        scope_name = b(f"ds:{self.scope.name}")
+        if self.scope.is_collection:
+            parts = [b(f"ds:{d}") for d in self.scope.leaf_names]
+            self.store.db.sunionstore(scope_name, parts)
+
+        for id in self.store.db.sscan_iter(scope_name):
+            entity = self.get_entity(id.decode("utf-8"))
+            if entity is not None:
+                if include_schemata is not None and entity.schema not in include_schemata:
+                    continue
+                yield entity