nomenklatura-mpt 4.1.9__py3-none-any.whl

nomenklatura/store/sql.py

@@ -0,0 +1,219 @@
+from typing import Any, Generator, List, Optional, Set, Tuple
+
+from followthemoney import DS, SE, Property, Schema, Statement
+from sqlalchemy import Table, delete, func, select
+from sqlalchemy.engine import Engine, Transaction, create_engine
+from sqlalchemy.dialects.postgresql import insert as psql_insert
+from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+from sqlalchemy.sql.selectable import Select
+
+from nomenklatura import settings
+from nomenklatura.db import get_metadata
+from nomenklatura.resolver import Linker, Identifier
+from nomenklatura.db import make_statement_table
+from nomenklatura.store import Store, View, Writer
+
+
+class SQLStore(Store[DS, SE]):
+    def __init__(
+        self,
+        dataset: DS,
+        linker: Linker[SE],
+        uri: str = settings.DB_URL,
+        **engine_kwargs: Any,
+    ):
+        super().__init__(dataset, linker)
+        if "pool_size" not in engine_kwargs:
+            engine_kwargs["pool_size"] = settings.DB_POOL_SIZE
+        # if uri.lower().startswith("sqlite"):
+        #     engine_kwargs.pop("pool_size", None)
+        metadata = get_metadata()
+        self.engine: Engine = create_engine(uri, **engine_kwargs)
+        self.table = make_statement_table(metadata)
+        metadata.create_all(self.engine, tables=[self.table], checkfirst=True)
+
+    def writer(self) -> Writer[DS, SE]:
+        return SQLWriter(self)
+
+    def view(self, scope: DS, external: bool = False) -> View[DS, SE]:
+        return SQLView(self, scope, external=external)
+
+    def _execute(
+        self, q: Select[Any], stream: bool = True
+    ) -> Generator[Any, None, None]:
+        # execute any read query against sql backend
+        with self.engine.connect() as conn:
+            if stream:
+                conn = conn.execution_options(stream_results=True)
+            cursor = conn.execute(q)
+            while rows := cursor.fetchmany(10_000):
+                yield from rows
+
+    def _iterate_stmts(
+        self, q: Select[Any], stream: bool = True
+    ) -> Generator[Statement, None, None]:
+        for row in self._execute(q, stream=stream):
+            yield Statement.from_db_row(row)
+
+    def _iterate(
+        self, q: Select[Any], stream: bool = True
+    ) -> Generator[SE, None, None]:
+        current_id = None
+        current_stmts: list[Statement] = []
+        for stmt in self._iterate_stmts(q, stream=stream):
+            entity_id = stmt.entity_id
+            if current_id is None:
+                current_id = entity_id
+            if current_id != entity_id:
+                proxy = self.assemble(current_stmts)
+                if proxy is not None:
+                    yield proxy
+                current_id = entity_id
+                current_stmts = []
+            current_stmts.append(stmt)
+        if len(current_stmts):
+            proxy = self.assemble(current_stmts)
+            if proxy is not None:
+                yield proxy
+
+
+class SQLWriter(Writer[DS, SE]):
+    BATCH_STATEMENTS = 10_000
+
+    def __init__(self, store: SQLStore[DS, SE]):
+        self.store: SQLStore[DS, SE] = store
+        self.batch: Set[Statement] = set()
+        self.conn = self.store.engine.connect()
+        self.tx: Optional[Transaction] = None
+
+    def _upsert_batch(self) -> None:
+        if not len(self.batch):
+            return
+        values = [s.to_db_row() for s in self.batch]
+        if self.tx is None:
+            self.tx = self.conn.begin()
+        if self.store.engine.dialect.name == "sqlite":
+            ilstmt = sqlite_insert(self.store.table).values(values)
+            lstmt = ilstmt.on_conflict_do_update(
+                index_elements=["id"],
+                set_=dict(
+                    canonical_id=ilstmt.excluded.canonical_id,
+                    schema=ilstmt.excluded.schema,
+                    prop_type=ilstmt.excluded.prop_type,
+                    lang=ilstmt.excluded.lang,
+                    original_value=ilstmt.excluded.original_value,
+                    last_seen=ilstmt.excluded.last_seen,
+                ),
+            )
+            self.conn.execute(lstmt)
+        elif self.store.engine.dialect.name in ("postgresql", "postgres"):
+            ipstmt = psql_insert(self.store.table).values(values)
+            pstmt = ipstmt.on_conflict_do_update(
+                index_elements=["id"],
+                set_=dict(
+                    canonical_id=ipstmt.excluded.canonical_id,
+                    schema=ipstmt.excluded.schema,
+                    prop_type=ipstmt.excluded.prop_type,
+                    lang=ipstmt.excluded.lang,
+                    original_value=ipstmt.excluded.original_value,
+                    last_seen=ipstmt.excluded.last_seen,
+                ),
+            )
+            self.conn.execute(pstmt)
+        else:
+            msg = f"Upsert not implemented for dialect {self.store.engine.dialect.name}"
+            raise NotImplementedError(msg)
+        self.batch = set()
+
+    def flush(self) -> None:
+        if len(self.batch):
+            self._upsert_batch()
+        if self.tx is not None:
+            self.tx.commit()
+            self.tx = None
+
+    def add_statement(self, stmt: Statement) -> None:
+        if stmt.entity_id is None:
+            return
+        canonical_id = self.store.linker.get_canonical(stmt.entity_id)
+        stmt.canonical_id = canonical_id
+        self.batch.add(stmt)
+        if len(self.batch) >= self.BATCH_STATEMENTS:
+            self._upsert_batch()
+
+    def pop(self, entity_id: str) -> List[Statement]:
+        if self.tx is None:
+            self.tx = self.conn.begin()
+
+        table = self.store.table
+        q = select(table)
+        q = q.where(table.c.canonical_id == entity_id)
+        statements: List[Statement] = []
+        cursor = self.conn.execute(q)
+        for row in cursor.fetchall():
+            statements.append(Statement.from_db_row(row))
+
+        q_delete = delete(table)
+        q_delete = q_delete.where(table.c.canonical_id == entity_id)
+        self.conn.execute(q_delete)
+        return statements
+
+
+class SQLView(View[DS, SE]):
+    def __init__(
+        self, store: SQLStore[DS, SE], scope: DS, external: bool = False
+    ) -> None:
+        super().__init__(store, scope, external=external)
+        self.store: SQLStore[DS, SE] = store
+
+    def get_entity(self, id: str) -> Optional[SE]:
+        table = self.store.table
+        q = select(table)
+        q = q.where(table.c.canonical_id == id)
+        q = q.where(table.c.dataset.in_(self.dataset_names))
+        for proxy in self.store._iterate(q, stream=False):
+            return proxy
+        return None
+
+    def has_entity(self, id: str) -> bool:
+        table = self.store.table
+        q = select(func.count(table.c.id))
+        q = q.where(table.c.canonical_id == id)
+        q = q.where(table.c.dataset.in_(self.dataset_names))
+        with self.store.engine.connect() as conn:
+            cursor = conn.execute(q)
+            count = cursor.scalar()
+            if count is not None and count > 0:
+                return True
+            else:
+                return False
+
+    def get_inverted(self, id: str) -> Generator[Tuple[Property, SE], None, None]:
+        table = self.store.table
+        id_ = Identifier.get(id)
+        ids = [i.id for i in self.store.linker.connected(id_)]
+        q = select(table.c.canonical_id)
+        q = q.where(table.c.prop_type == "entity")
+        q = q.where(table.c.value.in_(ids))
+        q = q.where(table.c.dataset.in_(self.dataset_names))
+        q = q.group_by(table.c.canonical_id)
+        with self.store.engine.connect() as conn:
+            cursor = conn.execute(q)
+            for (canonical_id,) in cursor.fetchall():
+                if canonical_id is None:
+                    continue
+                entity = self.get_entity(canonical_id)
+                if entity is not None:
+                    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]:
+        table: Table = self.store.table
+        q = select(table)
+        q = q.where(table.c.dataset.in_(self.dataset_names))
+        q = q.order_by(table.c.canonical_id)
+        for entity in self.store._iterate(q, stream=True):
+            if include_schemata is not None and entity.schema not in include_schemata:
+                continue
+            yield entity

nomenklatura/store/util.py

@@ -0,0 +1,48 @@
+import orjson
+
+from followthemoney import Statement
+from followthemoney.statement.util import pack_prop, unpack_prop
+
+
+def pack_statement(stmt: Statement) -> bytes:
+    values = (
+        stmt.id,
+        stmt.entity_id,
+        stmt.dataset,
+        pack_prop(stmt.schema, stmt.prop),
+        stmt.value,
+        stmt.lang,
+        stmt.original_value,
+        stmt.first_seen,
+        stmt.last_seen,
+    )
+    return orjson.dumps(values)
+
+
+def unpack_statement(data: bytes, canonical_id: str, external: bool) -> Statement:
+    (
+        id,
+        entity_id,
+        dataset,
+        prop_id,
+        value,
+        lang,
+        original_value,
+        first_seen,
+        last_seen,
+    ) = orjson.loads(data)
+    schema, _, prop = unpack_prop(prop_id)
+    return Statement(
+        id=id,
+        entity_id=entity_id,
+        prop=prop,
+        schema=schema,
+        value=value,
+        lang=lang,
+        dataset=dataset,
+        original_value=original_value,
+        first_seen=first_seen,
+        last_seen=last_seen,
+        canonical_id=canonical_id,
+        external=external,
+    )

nomenklatura/store/versioned.py

@@ -0,0 +1,371 @@
+import orjson
+import logging
+from redis.client import Redis
+from typing import Generator, List, Optional, Set, Tuple, Dict
+from followthemoney import DS, SE, Schema, registry, Property, Statement
+from followthemoney.statement.util import pack_prop, unpack_prop
+
+from nomenklatura.kv import b, bv, get_redis, close_redis
+from nomenklatura.versions import Version
+from nomenklatura.resolver import Linker, Identifier, StrIdent
+from nomenklatura.store.base import Store, View, Writer
+
+log = logging.getLogger(__name__)
+
+
+def _pack_statement(stmt: Statement) -> bytes:
+    values = (
+        stmt.id,
+        stmt.entity_id,
+        stmt.dataset,
+        pack_prop(stmt.schema, stmt.prop),
+        stmt.value,
+        stmt.lang or 0,
+        stmt.original_value or 0,
+        stmt.first_seen,
+        stmt.last_seen,
+        1 if stmt.external else 0,
+    )
+    return orjson.dumps(values)
+
+
+def _unpack_statement(data: bytes, canonical_id: Optional[str] = None) -> Statement:
+    (
+        id,
+        entity_id,
+        dataset,
+        prop_id,
+        value,
+        lang,
+        original_value,
+        first_seen,
+        last_seen,
+        external,
+    ) = orjson.loads(data)
+    schema, _, prop = unpack_prop(prop_id)
+    return Statement(
+        id=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,
+        first_seen=first_seen,
+        last_seen=last_seen,
+        canonical_id=canonical_id or entity_id,
+        external=external == 1,
+    )
+
+
+class VersionedRedisStore(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,
+        dataset: Optional[DS] = None,
+        version: Optional[str] = None,
+        timestamps: bool = False,
+    ) -> "VersionedRedisWriter[DS, SE]":
+        if version is None:
+            version = Version.new().id
+        dataset = dataset or self.dataset
+        return VersionedRedisWriter(
+            self,
+            dataset=dataset,
+            version=version,
+            timestamps=timestamps,
+        )
+
+    def view(
+        self, scope: DS, external: bool = False, versions: Dict[str, str] = {}
+    ) -> "VersionedRedisView[DS, SE]":
+        return VersionedRedisView(self, scope, external=external, versions=versions)
+
+    def update(self, id: StrIdent) -> None:
+        # Noop because the VersionedStore is not resolved.
+        return
+
+    def get_latest(self, dataset: str) -> Optional[str]:
+        """Get the latest version of a dataset in the store."""
+        val = self.db.get(b(f"ds:{dataset}:latest"))
+        return val.decode("utf-8") if val is not None else None
+
+    def get_history(self, dataset: str) -> List[str]:
+        """List all versions of a dataset present in the store."""
+        values = self.db.lrange(f"ds:{dataset}:history", 0, -1)
+        return [v.decode("utf-8") for v in values]
+
+    def has_version(self, dataset: str, version: str) -> bool:
+        """Check if a specific version of a dataset exists in the store."""
+        return self.db.exists(f"ents:{dataset}:{version}") > 0
+
+    def release_version(self, dataset: str, version: str) -> None:
+        """Release the given version of the dataset (i.e. tag it as the latest
+        version in the relevant lookup key)."""
+        history_key = b(f"ds:{dataset}:history")
+        idx = self.db.lpos(history_key, b(version))
+        if idx is None:
+            self.db.lpush(history_key, b(version))
+        latest = self.db.lindex(history_key, 0)
+        if latest is not None:
+            self.db.set(b(f"ds:{dataset}:latest"), latest)
+        log.info("Released store version: %s (%s)", dataset, version)
+
+    def drop_version(self, dataset: str, version: str) -> None:
+        """Delete all data associated with a specific version of a dataset."""
+        pipeline = self.db.pipeline()
+        cmds = 0
+        for prefix in ["stmt", "ents", "inv"]:
+            query = f"{prefix}:{dataset}:{version}*"
+            for key in self.db.scan_iter(query):
+                pipeline.delete(key)
+                cmds += 1
+                if cmds > 1_000:
+                    pipeline.execute()
+                    pipeline = self.db.pipeline()
+                    cmds = 0
+        if cmds > 0:
+            pipeline.execute()
+
+        # TODO: do we even want to remove the version from the history list?
+        self.db.lrem(f"ds:{dataset}:history", 0, b(version))
+        latest_key = f"ds:{dataset}:latest"
+        if b(version) == self.db.get(latest_key):
+            previous = self.db.lindex(b(f"ds:{dataset}:history"), 0)
+            if previous is not None:
+                self.db.set(latest_key, previous)
+            else:
+                self.db.delete(latest_key)
+        log.info("Dropped store version: %s (%s)", dataset, version)
+
+    def close(self) -> None:
+        close_redis()
+
+
+class VersionedRedisWriter(Writer[DS, SE]):
+    BATCH_STATEMENTS = 2_000
+
+    def __init__(
+        self,
+        store: VersionedRedisStore[DS, SE],
+        dataset: DS,
+        version: str,
+        timestamps: bool = False,
+    ):
+        self.version = version
+        self.dataset = dataset
+        self.timestamps = timestamps
+        self.ver = f"{dataset.name}:{version}"
+        self.store: VersionedRedisStore[DS, SE] = store
+        self.prev = store.get_latest(dataset.name)
+        self.buffer: List[Statement] = []
+
+    def __enter__(self) -> "VersionedRedisWriter[DS, SE]":
+        return self
+
+    def flush(self) -> None:
+        db = self.store.db
+        pipeline = db.pipeline()
+
+        statements: Dict[str, Set[Statement]] = {}
+        for stmt in self.buffer:
+            if stmt.entity_id not in statements:
+                statements[stmt.entity_id] = set()
+            statements[stmt.entity_id].add(stmt)
+
+        if len(statements) == 0:
+            return
+
+        # Merge with previous version to get accurate first_seen timestamps
+        if self.timestamps and self.prev:
+            keys = [b(f"stmt:{self.prev}:{e}") for e in statements.keys()]
+            for v in db.sunion(keys):
+                pstmt = _unpack_statement(bv(v))
+                for stmt in self.buffer:
+                    if pstmt.id == stmt.id:
+                        stmt.first_seen = pstmt.first_seen
+                        break
+
+        for entity_id, stmts in statements.items():
+            b_entity_id = b(entity_id)
+            pipeline.sadd(b(f"ents:{self.ver}"), b_entity_id)
+            values = [_pack_statement(s) for s in stmts]
+            pipeline.sadd(f"stmt:{self.ver}:{entity_id}", *values)
+
+            for stmt in stmts:
+                if stmt.prop_type == registry.entity.name:
+                    pipeline.sadd(b(f"inv:{self.ver}:{stmt.value}"), b_entity_id)
+
+        pipeline.execute()
+        self.buffer = []
+
+    def release(self) -> None:
+        """Release the current version of the dataset (i.e. tag it as the latest
+        version in the relevant lookup key)."""
+        self.store.release_version(self.dataset.name, self.version)
+
+    def close(self) -> None:
+        self.release()
+        self.store.close()
+
+    def add_statement(self, stmt: Statement) -> None:
+        if stmt.entity_id is None:
+            return
+        self.buffer.append(stmt)
+        if len(self.buffer) >= self.BATCH_STATEMENTS:
+            self.flush()
+
+    def pop(self, entity_id: str) -> List[Statement]:
+        raise NotImplementedError()
+
+
+class VersionedRedisView(View[DS, SE]):
+    def __init__(
+        self,
+        store: VersionedRedisStore[DS, SE],
+        scope: DS,
+        external: bool = False,
+        versions: Dict[str, str] = {},
+    ) -> None:
+        super().__init__(store, scope, external=external)
+        self.store: VersionedRedisStore[DS, SE] = store
+
+        # Get the latest version for each dataset in the scope
+        self.vers: List[Tuple[str, str]] = []
+        for ds in scope.leaf_names:
+            version = versions.get(ds, self.store.get_latest(ds))
+            if version is not None:
+                self.vers.append((ds, version))
+
+    def _get_stmt_keys(self, entity_id: str) -> List[str]:
+        keys: List[str] = []
+        ident = Identifier.get(entity_id)
+        for id in self.store.linker.connected(ident):
+            keys.extend([f"stmt:{d}:{v}:{id}" for d, v in self.vers])
+        return keys
+
+    def has_entity(self, id: str) -> bool:
+        # FIXME: this implementation does not account for the `external` flag
+        # correctly because it does not check the `stmt.external` field for
+        # each statement.
+        return self.store.db.exists(*self._get_stmt_keys(id)) > 0
+
+    def _get_statements(self, id: str) -> Generator[Statement, None, None]:
+        keys = self._get_stmt_keys(id)
+        if len(keys) == 0:
+            return None
+        elif len(keys) == 1:
+            stmts = self.store.db.smembers(keys[0])
+        else:
+            stmts = {bv(s) for s in self.store.db.sunion(keys)}
+        for v in stmts:
+            stmt = _unpack_statement(bv(v), id)
+            yield stmt
+
+    def get_timestamps(self, id: str) -> Dict[str, str]:
+        """Get the first seen timestamps associated with all statements of an entity.
+
+        Returns a dictionary mapping statement IDs to their first seen timestamps.
+        This can be used by an ETL to generate continuous entity histories.
+        """
+        timestamps: Dict[str, str] = {}
+        for stmt in self._get_statements(id):
+            if stmt.id is not None and stmt.first_seen is not None:
+                timestamps[stmt.id] = stmt.first_seen
+        return timestamps
+
+    def get_entity(self, id: str) -> Optional[SE]:
+        statements: List[Statement] = []
+        for stmt in self._get_statements(id):
+            if not stmt.external or self.external:
+                stmt.canonical_id = self.store.linker.get_canonical(stmt.entity_id)
+                if stmt.prop_type == registry.entity.name:
+                    stmt.value = self.store.linker.get_canonical(stmt.value)
+                statements.append(stmt)
+        return self.store.assemble(statements)
+
+    def get_inverted(self, id: str) -> Generator[Tuple[Property, SE], None, None]:
+        keys: List[str] = []
+        ident = Identifier.get(id)
+        for ent_id in self.store.linker.connected(ident):
+            keys.extend([f"inv:{d}:{v}:{ent_id}" for d, v in self.vers])
+        refs = (
+            {bv(v) for v in self.store.db.sunion(keys)}
+            if len(keys) > 0
+            else self.store.db.smembers(keys[0])
+        )
+        entities: Set[str] = set()
+        for v in refs:
+            entity_id = v.decode("utf-8")
+            entities.add(self.store.linker.get_canonical(entity_id))
+        for entity_id in entities:
+            entity = self.get_entity(entity_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 statements(self, resolve: bool = False) -> Generator[Statement, None, None]:
+        """Iterate over all statements in the view. If `resolve` is set to `True`,
+        canonical IDs are applied to the statement and its value.
+
+        NOTE: The `external` flag of the view will be used to filter statements, too.
+        """
+        for ds, ver in self.vers:
+            for id in self.store.db.sscan_iter(b(f"ents:{ds}:{ver}")):
+                entity_id = id.decode("utf-8")
+                stmt_key = f"stmt:{ds}:{ver}:{entity_id}"
+                for stmt_text in self.store.db.smembers(b(stmt_key)):
+                    stmt = _unpack_statement(stmt_text, entity_id)
+                    if stmt.external and not self.external:
+                        continue
+                    if resolve:
+                        stmt = self.store.linker.apply_statement(stmt)
+                    yield stmt
+
+    def entities(self, include_schemata: Optional[List[Schema]] = None) -> Generator[SE, None, None]:
+        if len(self.vers) == 0:
+            return
+        if len(self.vers) == 1:
+            scope_name = b(f"ents:{self.vers[0][0]}:{self.vers[0][1]}")
+        else:
+            version = Version.new().id + ":iter"
+            scope_name = b(f"ents:{self.scope.name}:{version}")
+            parts = [b(f"ents:{d}:{v}") for d, v in self.vers]
+            self.store.db.sunionstore(scope_name, parts)
+
+        # Keep track of canonical entities to avoid yielding the same
+        # de-duplicated entity multiple times. This intrinsically leaks
+        # memory, so we're being careful to only record entity IDs
+        # that are part of a cluster with more than one ID.
+        try:
+            seen: Set[str] = set()
+            for id in self.store.db.sscan_iter(scope_name):
+                entity_id = id.decode("utf-8")
+                ident = Identifier.get(entity_id)
+                connected = self.store.linker.connected(ident)
+                if len(connected) > 1:
+                    canonical_id = max(connected).id
+                    if canonical_id in seen:
+                        continue
+                    seen.add(canonical_id)
+                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
+        finally:
+            if len(self.vers) > 1:
+                self.store.db.delete(scope_name)

nomenklatura/tui/__init__.py

@@ -0,0 +1,17 @@
+from typing import Optional
+from followthemoney import DS, SE
+
+from nomenklatura.store import Store
+
+from nomenklatura.tui.app import DedupeApp, DedupeState
+from nomenklatura.resolver import Resolver
+
+__all__ = ["dedupe_ui"]
+
+
+def dedupe_ui(
+    resolver: Resolver[SE], store: Store[DS, SE], url_base: Optional[str] = None
+) -> None:
+    app = DedupeApp()
+    app.dedupe = DedupeState(resolver, store, url_base=url_base)
+    app.run()