diff --git a/pyproject.toml b/pyproject.toml index 33b1fb7c..233eb8dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,7 @@ requires-python = ">=3.11,<3.14" dependencies = [ "click>=8.1.7", "cryptography>=44.0.2", + "duckdb-engine>=0.13.6", "duckdb>=1.1.1", "email-validator>=2.3.0", "faker>=36.1.1", diff --git a/src/matchbox/client/adapters/__init__.py b/src/matchbox/client/adapters/__init__.py new file mode 100644 index 00000000..27918c94 --- /dev/null +++ b/src/matchbox/client/adapters/__init__.py @@ -0,0 +1,5 @@ +"""Local (client-side) database adapter implementations.""" + +from matchbox.client.adapters.duckdb import MatchboxLocalDuckDB + +__all__ = ["MatchboxLocalDuckDB"] diff --git a/src/matchbox/client/adapters/duckdb/__init__.py b/src/matchbox/client/adapters/duckdb/__init__.py new file mode 100644 index 00000000..c3f27741 --- /dev/null +++ b/src/matchbox/client/adapters/duckdb/__init__.py @@ -0,0 +1,5 @@ +"""DuckDB implementation of the local (client-side) Matchbox cluster store.""" + +from matchbox.client.adapters.duckdb.adapter import MatchboxLocalDuckDB + +__all__ = ["MatchboxLocalDuckDB"] diff --git a/src/matchbox/client/adapters/duckdb/adapter.py b/src/matchbox/client/adapters/duckdb/adapter.py new file mode 100644 index 00000000..55f12167 --- /dev/null +++ b/src/matchbox/client/adapters/duckdb/adapter.py @@ -0,0 +1,743 @@ +"""DuckDB implementation of the local (client-side) Matchbox cluster store. + +Shares SQL builders from matchbox.common.adapters.sql with the Postgres +backend. Engine-specific code lives in db.py and orm.py. +""" + +import base64 +from collections.abc import Callable +from pathlib import Path + +import pyarrow as pa +from pyarrow import Table as ArrowTable +from sqlalchemy import ( + BigInteger, + column, + delete, + exists, + func, + insert, + literal, + literal_column, + select, + text, +) +from sqlalchemy import ( + table as sa_table, +) +from sqlalchemy.engine import Engine +from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.sql.selectable import Select + +from matchbox.client.adapters.duckdb import db, orm +from matchbox.client.base import MatchboxLocalDBAdapter +from matchbox.common.adapters.protocol import MatchboxLocalBackends, MatchboxSnapshot +from matchbox.common.adapters.sql import tables +from matchbox.common.adapters.sql.insert import ( + select_cluster_map, + select_contains_pairs, + select_key_expansion, + select_new_cluster_hashes, + select_resolver_membership, +) +from matchbox.common.adapters.sql.query import ( + assemble_matches, + build_matching_leaves_cte, + build_target_cluster_cte, + build_unified_query, + resolver_membership_subquery, +) +from matchbox.common.adapters.sql.resolver import ( + build_expanded_leaves_subquery, + build_leaf_hash_groups_query, + hash_resolver_parents, +) +from matchbox.common.adapters.sql.snapshot import dump_tables, restore_tables +from matchbox.common.arrow import ( + SCHEMA_CLUSTERS, + SCHEMA_MODEL_EDGES, + SCHEMA_QUERY, + SCHEMA_QUERY_WITH_LEAVES, +) +from matchbox.common.dtos import ( + Match, + ModelStepPath, + ResolverStepPath, + SourceStepPath, + Step, + StepPath, + StepType, +) +from matchbox.common.exceptions import ( + MatchboxDataNotFound, + MatchboxDeletionNotConfirmed, +) + + +def _translate_lineage( + session: Session, lineage_names: list[str] +) -> list[tuple[int, int | None]]: + """Translate step names from self.lineage() into (step_id, source_config_id). + + Model steps are dropped: they have no source_config_id and aren't + queryable, unlike source and resolver steps. + """ + translated: list[tuple[int, int | None]] = [] + for name in lineage_names: + step = orm.Steps.from_name(session, name) + if step.type == StepType.MODEL.value: + continue + source_config_id = ( + step.source_config.source_config_id if step.source_config else None + ) + translated.append((step.step_id, source_config_id)) + return translated + + +_RAW_DATA_TABLE_PREFIX = "mb_raw_data__" +_QUERY_CACHE_TABLE_PREFIX = "mb_query_cache__" + + +def _raw_data_table_name(source_step_id: int) -> str: + """Physical table name for a source step's raw data.""" + return f"{_RAW_DATA_TABLE_PREFIX}{source_step_id}" + + +def _query_cache_table_name(cache_id: int) -> str: + """Physical table name for a cached query result.""" + return f"{_QUERY_CACHE_TABLE_PREFIX}{cache_id}" + + +def _materialise( + session: Session, physical_name: str, view_name: str, table: ArrowTable +) -> None: + """Zero-copy materialise an ArrowTable as a physical DuckDB table.""" + raw_conn = session.connection().connection.driver_connection + raw_conn.register(view_name, table) + ddl = f'CREATE OR REPLACE TABLE "{physical_name}" AS SELECT * FROM "{view_name}"' + session.execute(text(ddl)) + + +def _arrow_query(session: Session, stmt: Select) -> ArrowTable: + """Execute a read-only select via DuckDB's native Arrow export.""" + connection = session.connection() + schema_translate_map = connection.get_execution_options().get( + "schema_translate_map" + ) + compiled = stmt.compile( + dialect=connection.engine.dialect, + compile_kwargs={"literal_binds": True}, + schema_translate_map=schema_translate_map, + render_schema_translate=True, + ) + raw_conn = connection.connection.driver_connection + return raw_conn.sql(str(compiled)).to_arrow_table() + + +def _dump_dynamic_table(session: Session, name: str) -> list[dict]: + """Fetch a physical table's rows as JSON-safe dicts, bytes wrapped as base64. + + Matches dump_tables' convention, so the snapshot format stays uniform. + """ + raw_conn = session.connection().connection.driver_connection + cursor = raw_conn.sql(f'SELECT * FROM "{name}"') + columns = [d[0] for d in cursor.description] + rows = cursor.fetchall() + return [ + { + col: {"base64": base64.b64encode(val).decode("ascii")} + if isinstance(val, bytes) + else val + for col, val in zip(columns, row, strict=True) + } + for row in rows + ] + + +def _restore_dynamic_table(session: Session, name: str, rows: list[dict]) -> None: + """Rebuild a physical table from _dump_dynamic_table's output.""" + decoded = [ + { + col: base64.b64decode(val["base64"]) if isinstance(val, dict) else val + for col, val in row.items() + } + for row in rows + ] + _materialise(session, name, "restore_tmp", pa.Table.from_pylist(decoded)) + + +def _count_source_clusters(session: Session) -> int: + """Count distinct clusters that have at least one source key.""" + stmt = ( + select(func.count(func.distinct(tables.Clusters.c.cluster_id))) + .select_from(tables.Clusters) + .join( + tables.ClusterSourceKey, + tables.ClusterSourceKey.c.cluster_id == tables.Clusters.c.cluster_id, + ) + ) + return session.execute(stmt).scalar_one() + + +def _count_model_clusters(session: Session) -> int: + """Count distinct clusters proposed by at least one resolver.""" + stmt = ( + select(func.count(func.distinct(tables.Clusters.c.cluster_id))) + .select_from(tables.Clusters) + .join( + tables.ResolverClusters, + tables.ResolverClusters.c.cluster_id == tables.Clusters.c.cluster_id, + ) + ) + return session.execute(stmt).scalar_one() + + +class _Countable: + """Binds a session-taking counter to this adapter's session factory. + + Local has no global singleton session like Postgres's MBDB - a + process may hold several stores - so each count needs its own engine. + """ + + def __init__( + self, session_factory: sessionmaker, counter: Callable[[Session], int] + ) -> None: + self._session_factory = session_factory + self._counter = counter + + def count(self) -> int: + """Counts the number of rows matching the statement.""" + with self._session_factory() as session: + return self._counter(session) + + +class MatchboxLocalDuckDBQueryMixin: + """Query mixin for the local DuckDB adapter.""" + + def query( # noqa: D102 + self, + source: SourceStepPath, + resolver: ResolverStepPath | None = None, + return_leaf_id: bool = False, + limit: int | None = None, + ) -> ArrowTable: + with self._session() as session: + source_step = orm.Steps.from_name(session, source.name, StepType.SOURCE) + # Fail loudly if the source has no data, rather than return nothing. + if source_step.source_config is None: + raise MatchboxDataNotFound(table="source_configs", data=[source.name]) + + self_step = ( + source_step + if resolver is None + else orm.Steps.from_name(session, resolver.name, StepType.RESOLVER) + ) + + lineage = _translate_lineage( + session, + self.lineage(self_step.name, sources=[source_step.name]), + ) + + query_stmt = build_unified_query( + lineage=lineage, level="key", include_source_config_id=False + ).order_by( + literal_column("root_id"), + literal_column("leaf_id"), + tables.ClusterSourceKey.c.key, + ) + if limit is not None: + query_stmt = query_stmt.limit(limit) + + arrow_table = _arrow_query(session, query_stmt) + + if return_leaf_id: + return ( + arrow_table.select(["root_id", "key", "leaf_id"]) + .rename_columns(["id", "key", "leaf_id"]) + .cast(SCHEMA_QUERY_WITH_LEAVES) + ) + return ( + arrow_table.select(["root_id", "key"]) + .rename_columns(["id", "key"]) + .cast(SCHEMA_QUERY) + ) + + def match( # noqa: D102 + self, + key: str, + source: SourceStepPath, + targets: list[SourceStepPath], + resolver: ResolverStepPath, + ) -> list[Match]: + with self._session() as session: + resolver_step = orm.Steps.from_name( + session, resolver.name, StepType.RESOLVER + ) + source_step = orm.Steps.from_name(session, source.name, StepType.SOURCE) + if source_step.source_config is None: + raise MatchboxDataNotFound(table="source_configs", data=[source.name]) + source_config_id = source_step.source_config.source_config_id + + target_config_ids = [] + for target in targets: + target_step = orm.Steps.from_name(session, target.name, StepType.SOURCE) + if target_step.source_config is None: + raise MatchboxDataNotFound( + table="source_configs", data=[target.name] + ) + target_config_ids.append(target_step.source_config.source_config_id) + + source_and_target_ids = [source_config_id, *target_config_ids] + + lineage = _translate_lineage(session, self.lineage(resolver_step.name)) + + target_cluster_cte = build_target_cluster_cte( + key=key, source_config_id=source_config_id, lineage=lineage + ) + matching_leaves_cte = build_matching_leaves_cte( + source_and_target_ids=source_and_target_ids, + lineage=lineage, + target_cluster_cte=target_cluster_cte, + ) + + matched_rows = session.execute( + select( + matching_leaves_cte.c.cluster_id, + matching_leaves_cte.c.source_config_id, + matching_leaves_cte.c.key, + ) + ).all() + + return assemble_matches( + matched_rows=matched_rows, + source=source, + source_config_id=source_config_id, + targets=targets, + target_config_ids=target_config_ids, + ) + + +class MatchboxLocalDuckDBDataMixin: + """Data mixin for the local DuckDB adapter.""" + + def create_step(self, step: Step, path: StepPath) -> None: # noqa: D102 + with self._session() as session: + existing = session.execute( + select(orm.Steps).where(orm.Steps.name == path.name) + ).scalar_one_or_none() + + if existing is None: + new_step = orm.Steps( + name=path.name, + type=step.step_type.value, + fingerprint=step.fingerprint, + ) + session.add(new_step) + session.flush() + + if step.step_type == StepType.SOURCE: + session.add(orm.SourceConfigs(step_id=new_step.step_id)) + else: + existing.type = step.step_type.value + existing.fingerprint = step.fingerprint + + session.commit() + + def insert_source_data( # noqa: D102 + self, path: SourceStepPath, data_hashes: ArrowTable + ) -> None: + with self._session() as session: + step = orm.Steps.from_name(session, path.name, StepType.SOURCE) + self._cascade_invalidate(session, step) + if step.source_config is None: + raise MatchboxDataNotFound(table="source_configs", data=[path.name]) + source_config_id = step.source_config.source_config_id + + if data_hashes.num_rows == 0: + session.commit() + return + + raw_conn = session.connection().connection.driver_connection + raw_conn.register("incoming_hashes", data_hashes.select(["hash", "keys"])) + incoming = sa_table("incoming_hashes", column("hash"), column("keys")) + + session.execute( + insert(tables.Clusters).from_select( + ["cluster_hash"], select_new_cluster_hashes(incoming) + ) + ) + session.execute( + insert(tables.ClusterSourceKey).from_select( + ["cluster_id", "source_config_id", "key"], + select_key_expansion(incoming, source_config_id), + ) + ) + session.commit() + + def insert_model_data( # noqa: D102 + self, path: ModelStepPath, results: ArrowTable + ) -> None: + with self._session() as session: + step = orm.Steps.from_name(session, path.name, StepType.MODEL) + self._cascade_invalidate(session, step) + + if results.num_rows == 0: + session.commit() + return + + raw_conn = session.connection().connection.driver_connection + raw_conn.register( + "incoming_edges", results.select(["left_id", "right_id", "score"]) + ) + incoming = sa_table( + "incoming_edges", column("left_id"), column("right_id"), column("score") + ) + + session.execute( + insert(tables.ModelEdges).from_select( + ["step_id", "left_id", "right_id", "score"], + select( + literal(step.step_id, BigInteger).label("step_id"), + incoming.c.left_id, + incoming.c.right_id, + incoming.c.score, + ), + ) + ) + session.commit() + + def insert_resolver_data( # noqa: D102 + self, path: ResolverStepPath, data: ArrowTable + ) -> None: + with self._session() as session: + step = orm.Steps.from_name(session, path.name, StepType.RESOLVER) + self._cascade_invalidate(session, step) + + if data.num_rows == 0: + session.commit() + return + + raw_conn = session.connection().connection.driver_connection + raw_conn.register( + "incoming_resolver_assignments", + data.select(["parent_id", "child_id"]), + ) + incoming = sa_table( + "incoming_resolver_assignments", + column("parent_id"), + column("child_id"), + ) + + expanded_leaves = build_expanded_leaves_subquery(incoming) + hash_rows = session.execute( + build_leaf_hash_groups_query(expanded_leaves) + ).all() + hash_table = hash_resolver_parents(hash_rows) + + raw_conn.register("resolver_hashes", hash_table) + resolver_hashes = sa_table( + "resolver_hashes", column("parent_id"), column("cluster_hash") + ) + + # 1) Bulk-insert new clusters: anti-join, ID omitted (sequence-assigned) + session.execute( + insert(tables.Clusters).from_select( + ["cluster_hash"], + select_new_cluster_hashes(resolver_hashes, hash_col="cluster_hash"), + ) + ) + + # 2) Map parent_id to canonical cluster_id, now hashes exist + cluster_map = select_cluster_map(resolver_hashes) + + # 3) ResolverClusters: distinct (step_id, cluster_id) + session.execute( + insert(tables.ResolverClusters).from_select( + ["step_id", "cluster_id"], + select_resolver_membership(step.step_id, cluster_map), + ) + ) + + # 4) Contains: new (root, leaf) pairs, anti-joined against existing + candidate_contains = select_contains_pairs( + expanded_leaves, cluster_map + ).subquery() + session.execute( + insert(tables.Contains).from_select( + ["root", "leaf"], + select(candidate_contains.c.root, candidate_contains.c.leaf).where( + ~exists( + select(1).where( + tables.Contains.c.root == candidate_contains.c.root, + tables.Contains.c.leaf == candidate_contains.c.leaf, + ) + ) + ), + ) + ) + + session.commit() + + def get_model_data(self, path: ModelStepPath) -> ArrowTable: # noqa: D102 + with self._session() as session: + step = orm.Steps.from_name(session, path.name, StepType.MODEL) + arrow_table = _arrow_query( + session, + select( + tables.ModelEdges.c.left_id, + tables.ModelEdges.c.right_id, + tables.ModelEdges.c.score, + ).where(tables.ModelEdges.c.step_id == step.step_id), + ) + return arrow_table.cast(SCHEMA_MODEL_EDGES) + + def get_resolver_data(self, path: ResolverStepPath) -> ArrowTable: # noqa: D102 + with self._session() as session: + step = orm.Steps.from_name(session, path.name, StepType.RESOLVER) + membership = resolver_membership_subquery( + step_id=step.step_id, alias="assignments" + ) + stmt = select( + membership.c.root_id.label("parent_id"), + membership.c.leaf_id.label("child_id"), + ).order_by(membership.c.root_id, membership.c.leaf_id) + arrow_table = _arrow_query(session, stmt) + + return arrow_table.cast(SCHEMA_CLUSTERS) + + def dump(self) -> MatchboxSnapshot: # noqa: D102 + with self._session() as session: + data = dump_tables(session, [*orm.LOCAL_TABLES, *orm.SHARED_TABLES]) + physical_names = [ + _raw_data_table_name(row["source_step_id"]) for row in data["raw_data"] + ] + [ + _query_cache_table_name(row["cache_id"]) for row in data["query_cache"] + ] + # dump_tables only covers catalog rows - physical tables need + # dumping separately, same row format. + data["__dynamic_cache_tables__"] = { + name: _dump_dynamic_table(session, name) for name in physical_names + } + return MatchboxSnapshot(backend_type=MatchboxLocalBackends.DUCKDB, data=data) + + def restore(self, snapshot: MatchboxSnapshot) -> None: # noqa: D102 + if snapshot.backend_type != MatchboxLocalBackends.DUCKDB: + raise TypeError( + f"Cannot restore {snapshot.backend_type} snapshot to duckdb backend" + ) + + self.clear(certain=True) + + with self._session() as session: + restore_tables( + session, [*orm.LOCAL_TABLES, *orm.SHARED_TABLES], snapshot.data + ) + for name, rows in snapshot.data.get("__dynamic_cache_tables__", {}).items(): + _restore_dynamic_table(session, name, rows) + session.commit() + + def clear(self, certain: bool) -> None: # noqa: D102 + if not certain: + raise MatchboxDeletionNotConfirmed( + "This operation will drop all rows in the database but not the " + "tables themselves. Rerun with certain=True to continue." + ) + with self._session() as session: + for tbl in [*orm.SHARED_TABLES, *orm.LOCAL_TABLES]: + session.execute(delete(tbl)) + session.commit() + db.drop_dynamic_cache_tables( + self._engine, (_RAW_DATA_TABLE_PREFIX, _QUERY_CACHE_TABLE_PREFIX) + ) + + def drop(self, certain: bool) -> None: # noqa: D102 + if not certain: + raise MatchboxDeletionNotConfirmed( + "This operation will drop the entire database and recreate it. " + "Rerun with certain=True to continue." + ) + db.drop_dynamic_cache_tables( + self._engine, (_RAW_DATA_TABLE_PREFIX, _QUERY_CACHE_TABLE_PREFIX) + ) + db.drop_db(self._engine, [*orm.SHARED_TABLES, *orm.LOCAL_TABLES]) + db.create_db(self._engine, orm.LOCAL_TABLES) + db.create_db(self._engine, orm.SHARED_TABLES) + + +class MatchboxLocalDuckDBCacheMixin: + """Local-only mixin: raw data, query cache, cascade invalidation.""" + + def insert_raw_data(self, path: SourceStepPath, table: ArrowTable) -> None: # noqa: D102 + with self._session() as session: + step = orm.Steps.from_name(session, path.name, StepType.SOURCE) + physical_name = _raw_data_table_name(step.step_id) + _materialise(session, physical_name, "incoming_raw_data", table) + session.execute( + delete(orm.RawData).where(orm.RawData.source_step_id == step.step_id) + ) + session.add(orm.RawData(source_step_id=step.step_id)) + session.commit() + + def get_raw_data( # noqa: D102 + self, path: SourceStepPath, keys: list[str] | None = None + ) -> ArrowTable: + with self._session() as session: + step = orm.Steps.from_name(session, path.name, StepType.SOURCE) + row = session.execute( + select(orm.RawData).where(orm.RawData.source_step_id == step.step_id) + ).scalar_one_or_none() + + if row is None: + raise MatchboxDataNotFound(table="raw_data", data=[path.name]) + + physical_name = _raw_data_table_name(row.source_step_id) + raw_conn = session.connection().connection.driver_connection + if keys is None: + return raw_conn.sql(f'SELECT * FROM "{physical_name}"').to_arrow_table() + raw_conn.register("key_filter", pa.table({"key": keys})) + return raw_conn.sql( + f'SELECT t.* FROM "{physical_name}" t ' + 'SEMI JOIN "key_filter" f ON t.key = f.key' + ).to_arrow_table() + + def cache_query( # noqa: D102 + self, key: str, table: ArrowTable, depends_on: list[StepPath] + ) -> None: + with self._session() as session: + existing = session.execute( + select(orm.QueryCache).where(orm.QueryCache.cache_key == key) + ).scalar_one_or_none() + if existing is not None: + old_name = _query_cache_table_name(existing.cache_id) + session.execute(text(f'DROP TABLE IF EXISTS "{old_name}"')) + session.execute( + delete(orm.QueryCacheStep).where( + orm.QueryCacheStep.cache_id == existing.cache_id + ) + ) + session.execute( + delete(orm.QueryCache).where( + orm.QueryCache.cache_id == existing.cache_id + ) + ) + + row = orm.QueryCache(cache_key=key) + session.add(row) + session.flush() # assigns cache_id from its duckdb sequence default + + physical_name = _query_cache_table_name(row.cache_id) + _materialise(session, physical_name, "incoming_query_cache", table) + + step_ids = { + orm.Steps.from_name(session, p.name).step_id for p in depends_on + } + session.add_all( + orm.QueryCacheStep(cache_id=row.cache_id, step_id=step_id) + for step_id in step_ids + ) + session.commit() + + def get_cached_query(self, key: str) -> ArrowTable | None: # noqa: D102 + with self._session() as session: + row = session.execute( + select(orm.QueryCache).where(orm.QueryCache.cache_key == key) + ).scalar_one_or_none() + if row is None: + return None + physical_name = _query_cache_table_name(row.cache_id) + raw_conn = session.connection().connection.driver_connection + return raw_conn.sql(f'SELECT * FROM "{physical_name}"').to_arrow_table() + + def drop_step_data(self, path: StepPath) -> None: # noqa: D102 + with self._session() as session: + step = orm.Steps.from_name(session, path.name) + self._cascade_invalidate(session, step) + session.commit() + + def _cascade_invalidate(self, session: Session, step: "orm.Steps") -> None: + """Drop a step's data and its descendants', and their query cache. + + Responsibility for committing lies with the caller. + + Deletes are unconditional and type-agnostic: harmless no-ops for + tables a step doesn't apply to. RawData and Clusters/Contains are + left alone - canonical state and content-addressed rows, not + caches. Query cache invalidation is selective, via QueryCacheStep. + """ + step_ids = [step.step_id] + for name in self.descendants(step.name): + descendant = orm.Steps.from_name(session, name) + step_ids.append(descendant.step_id) + + session.execute( + delete(tables.ModelEdges).where(tables.ModelEdges.c.step_id.in_(step_ids)) + ) + session.execute( + delete(tables.ResolverClusters).where( + tables.ResolverClusters.c.step_id.in_(step_ids) + ) + ) + session.execute( + delete(tables.ClusterSourceKey).where( + tables.ClusterSourceKey.c.source_config_id.in_( + select(orm.SourceConfigs.source_config_id).where( + orm.SourceConfigs.step_id.in_(step_ids) + ) + ) + ) + ) + + affected_ids = ( + session.execute( + select(orm.QueryCacheStep.cache_id) + .where(orm.QueryCacheStep.step_id.in_(step_ids)) + .distinct() + ) + .scalars() + .all() + ) + for cache_id in affected_ids: + session.execute( + text(f'DROP TABLE IF EXISTS "{_query_cache_table_name(cache_id)}"') + ) + session.execute( + delete(orm.QueryCacheStep).where( + orm.QueryCacheStep.cache_id.in_(affected_ids) + ) + ) + session.execute( + delete(orm.QueryCache).where(orm.QueryCache.cache_id.in_(affected_ids)) + ) + + +class MatchboxLocalDuckDB( + MatchboxLocalDuckDBQueryMixin, + MatchboxLocalDuckDBDataMixin, + MatchboxLocalDuckDBCacheMixin, + MatchboxLocalDBAdapter, +): + """A DuckDB adapter for the local (client-side) Matchbox cluster store.""" + + def __init__(self, path: Path | None = None) -> None: + """Initialise the DuckDB adapter. + + Args: + path: Path to the duckdb file. None means in-memory and + ephemeral, dying with the process. + """ + self._engine: Engine = db.create_local_engine(path) + self._session_factory: sessionmaker = sessionmaker(bind=self._engine) + + db.create_db(self._engine, orm.LOCAL_TABLES) + db.create_db(self._engine, orm.SHARED_TABLES) + + self.all_clusters = _Countable(self._session_factory, orm.Clusters.count) + self.source_clusters = _Countable(self._session_factory, _count_source_clusters) + self.model_clusters = _Countable(self._session_factory, _count_model_clusters) + self.creates = _Countable(self._session_factory, orm.ResolverClusters.count) + self.merges = _Countable(self._session_factory, orm.Contains.count) + self.proposes = _Countable(self._session_factory, orm.ModelEdges.count) + + def _session(self) -> Session: + """Return a new session bound to this adapter's engine.""" + return self._session_factory() diff --git a/src/matchbox/client/adapters/duckdb/db.py b/src/matchbox/client/adapters/duckdb/db.py new file mode 100644 index 00000000..3bb85267 --- /dev/null +++ b/src/matchbox/client/adapters/duckdb/db.py @@ -0,0 +1,81 @@ +"""Engine, session, and DDL management for the local DuckDB adapter. + +Engine and session state live on the adapter instance, not a module-level +global, since a process can hold several independent local stores at once. +""" + +from pathlib import Path + +from sqlalchemy import Table, create_engine +from sqlalchemy.engine import Engine +from sqlalchemy.schema import CreateTable + +from matchbox.client.adapters.duckdb import orm + +# (table, PK column) pairs needing a duckdb sequence - no SERIAL/BIGSERIAL. +# Composite-PK and foreign-keyed-PK tables need no sequence, so are absent. +SEQUENCE_COLUMNS: list[tuple[Table, str]] = [ + (orm.Steps.__table__, "step_id"), + (orm.SourceConfigs.__table__, "source_config_id"), + (orm.Clusters.__table__, "cluster_id"), + (orm.ClusterSourceKey.__table__, "key_id"), + (orm.ModelEdges.__table__, "result_id"), + (orm.QueryCache.__table__, "cache_id"), +] + + +def create_local_engine(path: Path | None) -> Engine: + """Create a duckdb engine, in-memory and ephemeral if no path is given.""" + url = f"duckdb:///{path}" if path is not None else "duckdb:///:memory:" + return create_engine(url).execution_options(schema_translate_map={"mb": None}) + + +def create_db(engine: Engine, order: list[Table]) -> None: + """Create tables in order, without FK constraints. + + Duckdb has no ON DELETE CASCADE, so the adapter enforces integrity + itself and FK constraints are stripped rather than rendered. Sequence + defaults are created first, then attached once the table exists. + """ + with engine.begin() as conn: + for tbl, col_name in SEQUENCE_COLUMNS: + if tbl in order: + conn.exec_driver_sql( + f'CREATE SEQUENCE IF NOT EXISTS "{tbl.name}_{col_name}_seq"' + ) + + for tbl in order: + ddl = str( + CreateTable(tbl, include_foreign_key_constraints=[]).compile( + dialect=conn.dialect + ) + ) + conn.exec_driver_sql(ddl) + + for tbl, col_name in SEQUENCE_COLUMNS: + if tbl in order: + conn.exec_driver_sql( + f'ALTER TABLE "{tbl.name}" ALTER COLUMN "{col_name}" ' + f"SET DEFAULT nextval('{tbl.name}_{col_name}_seq')" + ) + + +def drop_db(engine: Engine, tables: list[Table]) -> None: + """Drop the given tables, if they exist.""" + with engine.begin() as conn: + for tbl in tables: + conn.exec_driver_sql(f'DROP TABLE IF EXISTS "{tbl.name}"') + + +def drop_dynamic_cache_tables(engine: Engine, prefixes: tuple[str, ...]) -> None: + """Drop tables named with one of prefixes - RawData/QueryCache's dynamic ones.""" + with engine.begin() as conn: + names = [ + row[0] + for row in conn.exec_driver_sql( + "SELECT table_name FROM duckdb_tables()" + ).fetchall() + if row[0].startswith(prefixes) + ] + for name in names: + conn.exec_driver_sql(f'DROP TABLE IF EXISTS "{name}"') diff --git a/src/matchbox/client/adapters/duckdb/orm.py b/src/matchbox/client/adapters/duckdb/orm.py new file mode 100644 index 00000000..88b3b785 --- /dev/null +++ b/src/matchbox/client/adapters/duckdb/orm.py @@ -0,0 +1,228 @@ +"""ORM classes for the local DuckDB backend. + +"steps" and "source_configs" are local-only (no run_id/collection/ +upload_stage - one file is one run). + +The five shared cluster tables are copied onto this module's own +LOCAL_METADATA via Table.to_metadata() (schema=None), rather than reusing +tables.METADATA directly, so that: + +- table names never collide with the Postgres ORM's tables in a process + that imports both. +- all of local's tables live in one MetaData object, so FK references + between them resolve without a cross-metadata lookup failure. + +Query building still references the original tables.Clusters etc. +objects (schema "mb") - SQL generation only needs matching table and +column names, not object identity. +""" + +from typing import Optional + +from sqlalchemy import ( + BigInteger, + FetchedValue, + ForeignKey, + LargeBinary, + MetaData, + Text, + func, + select, +) +from sqlalchemy.orm import ( + Mapped, + Session, + declarative_base, + mapped_column, + relationship, +) + +from matchbox.common.adapters.sql import tables +from matchbox.common.dtos import StepType +from matchbox.common.exceptions import MatchboxStepNotFoundError, MatchboxStepTypeError + +LOCAL_METADATA = MetaData() +LocalBase = declarative_base(metadata=LOCAL_METADATA) + + +class CountMixin: + """Adds `.count(session)` to an ORM class mapped to a whole table. + + Unlike Postgres's CountMixin, count() takes an explicit session: local + has no global singleton session to pull one from, since a process can + hold several independent MatchboxLocalDuckDB instances at once. + """ + + @classmethod + def count(cls, session: Session) -> int: + """Counts the number of rows in the table.""" + return session.execute(select(func.count()).select_from(cls)).scalar_one() + + +class Steps(LocalBase): + """Local step registry. No run_id/collection/upload_stage: one file is one run.""" + + __tablename__ = "steps" + + step_id: Mapped[int] = mapped_column( + BigInteger, + primary_key=True, + autoincrement=False, + server_default=FetchedValue(), + ) + name: Mapped[str] = mapped_column(Text, unique=True) + type: Mapped[str] = mapped_column(Text) + fingerprint: Mapped[bytes] = mapped_column(LargeBinary) + + source_config: Mapped[Optional["SourceConfigs"]] = relationship( + back_populates="step", uselist=False + ) + + @classmethod + def from_name( + cls, session: Session, name: str, expected_type: StepType | None = None + ) -> "Steps": + """Resolve a step by name, optionally validating its type. + + Args: + session: Database session. + name: The step's name. + expected_type: If given, raise if the step isn't of this type. + + Raises: + MatchboxStepNotFoundError: If no step has this name. + MatchboxStepTypeError: If expected_type is given and doesn't match. + """ + step = session.execute(select(cls).where(cls.name == name)).scalar_one_or_none() + if step is None: + raise MatchboxStepNotFoundError(name=name) + if expected_type is not None and step.type != expected_type.value: + raise MatchboxStepTypeError( + step_name=name, + step_type=StepType(step.type), + expected_step_types=[expected_type], + ) + return step + + +class SourceConfigs(LocalBase): + """Degenerate one-row-per-source table. + + Keeps the join shape identical to the server's, without needing the + rest of a real SourceConfig. + """ + + __tablename__ = "source_configs" + + source_config_id: Mapped[int] = mapped_column( + BigInteger, + primary_key=True, + autoincrement=False, + server_default=FetchedValue(), + ) + step_id: Mapped[int] = mapped_column( + BigInteger, ForeignKey("steps.step_id"), unique=True + ) + + step: Mapped["Steps"] = relationship(back_populates="source_config") + + +class RawData(LocalBase): + """Catalog of raw warehouse rows per source step. + + Canonical local state - expensive to re-fetch, so kept until + explicitly replaced. One row per source step, FK'd 1:1 to Steps: the + actual rows live in a physical DuckDB table named by + adapter._raw_data_table_name(source_step_id), never stored here as + data - see adapter.py for why. + """ + + __tablename__ = "raw_data" + + source_step_id: Mapped[int] = mapped_column( + BigInteger, ForeignKey("steps.step_id"), primary_key=True, autoincrement=False + ) + + +class QueryCache(LocalBase): + """Catalog of cleaned query results. + + Disposable: cleared wholesale whenever any step's data changes, since + recomputing is cheap. cache_key is the opaque identity of a query + definition (query config plus upstream fingerprints); cache_id is a + sequence-backed surrogate used only to name the physical table + (adapter._query_cache_table_name(cache_id)) - see adapter.py. + """ + + __tablename__ = "query_cache" + + cache_id: Mapped[int] = mapped_column( + BigInteger, primary_key=True, autoincrement=False, server_default=FetchedValue() + ) + cache_key: Mapped[str] = mapped_column(Text, unique=True) + + +class QueryCacheStep(LocalBase): + """Bridging table: which steps a cached query result depends on. + + A Query isn't itself a DAG step (unlike Source/Model/Resolver), so it + has no step_id of its own to hang cache invalidation off. This table + records the steps (sources + resolver) a cached result was built + from, so _cascade_invalidate can drop only the cache entries that + actually depend on a changed step. + """ + + __tablename__ = "query_cache_steps" + + cache_id: Mapped[int] = mapped_column( + BigInteger, ForeignKey("query_cache.cache_id"), primary_key=True + ) + step_id: Mapped[int] = mapped_column( + BigInteger, ForeignKey("steps.step_id"), primary_key=True + ) + + +class Clusters(CountMixin, LocalBase): + """Table of indexed data and clusters that match it.""" + + __table__ = tables.Clusters.to_metadata(LOCAL_METADATA, schema=None) + + +class ClusterSourceKey(CountMixin, LocalBase): + """Table for storing source primary keys for clusters.""" + + __table__ = tables.ClusterSourceKey.to_metadata(LOCAL_METADATA, schema=None) + + +class Contains(CountMixin, LocalBase): + """Cluster lineage table.""" + + __table__ = tables.Contains.to_metadata(LOCAL_METADATA, schema=None) + + +class ModelEdges(CountMixin, LocalBase): + """Table of results for a model step.""" + + __table__ = tables.ModelEdges.to_metadata(LOCAL_METADATA, schema=None) + + +class ResolverClusters(CountMixin, LocalBase): + """Association table linking resolver steps to cluster IDs.""" + + __table__ = tables.ResolverClusters.to_metadata(LOCAL_METADATA, schema=None) + + +LOCAL_TABLES = [ + Steps.__table__, + SourceConfigs.__table__, + RawData.__table__, + QueryCache.__table__, + QueryCacheStep.__table__, +] +SHARED_TABLES = [ + Clusters.__table__, + ClusterSourceKey.__table__, + Contains.__table__, + ModelEdges.__table__, + ResolverClusters.__table__, +] diff --git a/src/matchbox/client/base.py b/src/matchbox/client/base.py new file mode 100644 index 00000000..8b31bba1 --- /dev/null +++ b/src/matchbox/client/base.py @@ -0,0 +1,209 @@ +"""Base classes and utilities for Matchbox local (client-side) database adapters.""" + +from abc import ABC, abstractmethod +from collections import deque +from pathlib import Path +from typing import Protocol + +from pyarrow import Table +from pydantic import BaseModel, Field + +from matchbox.common.adapters.protocol import ( + MatchboxClusterStoreAdapter, + MatchboxLocalBackends, +) +from matchbox.common.dtos import SourceStepPath, StepName, StepPath + + +class HasGraph(Protocol): + """Structural type for anything exposing step topology as graph.""" + + graph: dict[StepName, list[StepName]] + + +def _ancestor_depths( + graph: dict[StepName, list[StepName]], start: StepName +) -> dict[StepName, int]: + """BFS over parent edges from start. + + Shortest depth per ancestor, self excluded. + """ + depths: dict[StepName, int] = {} + queue: deque[tuple[StepName, int]] = deque( + (parent, 1) for parent in graph.get(start, []) + ) + while queue: + name, depth = queue.popleft() + if name in depths and depths[name] <= depth: + continue + depths[name] = depth + queue.extend((parent, depth + 1) for parent in graph.get(name, [])) + return depths + + +def compute_lineage( + graph: dict[StepName, list[StepName]], + resolver: StepName, + sources: list[StepName] | None = None, +) -> list[StepName]: + """Ordered ancestor step names for resolver, closest first, self first. + + When sources is given, ancestors are restricted to those on a path + to one of sources: kept if it is a source, or is itself downstream + of one. + """ + ancestors = _ancestor_depths(graph, resolver) + + if sources: + children: dict[StepName, list[StepName]] = {} + for name, parents in graph.items(): + for parent in parents: + children.setdefault(parent, []).append(name) + + allowed = set(sources) + for source in sources: + allowed.update(_ancestor_depths(children, source)) + ancestors = { + name: depth for name, depth in ancestors.items() if name in allowed + } + + ordered = sorted(ancestors.items(), key=lambda pair: (pair[1], pair[0])) + return [resolver, *(name for name, _ in ordered)] + + +def compute_descendants( + graph: dict[StepName, list[StepName]], step: StepName +) -> list[StepName]: + """All step names downstream of step - inverts the parent graph, then BFS.""" + children: dict[StepName, list[StepName]] = {} + for name, parents in graph.items(): + for parent in parents: + children.setdefault(parent, []).append(name) + return list(_ancestor_depths(children, step)) + + +class MatchboxLocalSettings(BaseModel): + """Settings for a local (client-side) Matchbox backend.""" + + backend_type: MatchboxLocalBackends = MatchboxLocalBackends.DUCKDB + path: Path | None = Field( + default=None, + description=( + "Path to the local database file. None means in-memory and ephemeral." + ), + ) + + +class MatchboxLocalDBAdapter(MatchboxClusterStoreAdapter, ABC): + """An abstract base class for Matchbox local (client-side) database adapters. + + Extends MatchboxClusterStoreAdapter (the query block, data block, and + cluster counts) with the local-only surface: raw warehouse data, a + disposable query cache, and explicit cascade invalidation. + + Local stores replace-on-rerun (with cascade invalidation of descendant + steps), unlike server backends, which are write-once per step per run. + + The adapter is bound to a live DAG via bind(). lineage() and + descendants() read its graph fresh on every call, so steps added + later are picked up automatically. + """ + + _graph_source: HasGraph | None = None + + def bind(self, graph_source: HasGraph) -> None: + """Bind a live source of step topology. + + Stored by reference, not copied: lineage()/descendants() read + graph_source.graph fresh on every call, so later mutations are + picked up automatically. + """ + self._graph_source = graph_source + + def _graph(self) -> dict[StepName, list[StepName]]: + if self._graph_source is None: + raise RuntimeError( + "Adapter has no bound graph source. Call bind() before " + "lineage() or descendants()." + ) + return self._graph_source.graph + + def lineage( + self, resolver: StepName, sources: list[StepName] | None = None + ) -> list[StepName]: + """Ordered ancestor step names, highest priority first, self first. + + Restricted to paths that lead to sources, when given. Walks the + bound graph fresh on every call. + """ + return compute_lineage(self._graph(), resolver, sources) + + def descendants(self, step: StepName) -> list[StepName]: + """All step names downstream of step (for cascade invalidation).""" + return compute_descendants(self._graph(), step) + + @abstractmethod + def insert_raw_data(self, path: SourceStepPath, table: Table) -> None: + """Insert raw warehouse rows for a source step. Canonical local state. + + Replaces any existing raw data for this step, and cascades: + descendant steps' data is dropped, so stale downstream results can + never be served. + """ + ... + + @abstractmethod + def get_raw_data( + self, path: SourceStepPath, keys: list[str] | None = None + ) -> Table: + """Get raw warehouse rows for a source step, optionally filtered by key.""" + ... + + @abstractmethod + def cache_query(self, key: str, table: Table, depends_on: list[StepPath]) -> None: + """Cache cleaned query data, keyed by query config plus upstream fingerprints. + + Disposable: unlike insert_raw_data, this is just a cache, and can be + recomputed at any time. + + Args: + key: Opaque identity of the query definition this result was + built from. + table: The data to cache. + depends_on: The steps (sources and resolver) this result + depends on, so cascade invalidation can drop it when any + of them change. + """ + ... + + @abstractmethod + def get_cached_query(self, key: str) -> Table | None: + """Get cached cleaned query data, or None if not cached.""" + ... + + @abstractmethod + def drop_step_data(self, path: StepPath) -> None: + """Explicitly invalidate a step's data. Cascades to descendant steps.""" + ... + + +def get_local_backend_class( + backend_type: MatchboxLocalBackends, +) -> type[MatchboxLocalDBAdapter]: + """Get the appropriate local backend class based on the backend type.""" + if backend_type == MatchboxLocalBackends.DUCKDB: + from matchbox.client.adapters.duckdb import ( # noqa: PLC0415 + MatchboxLocalDuckDB, + ) + + return MatchboxLocalDuckDB + else: + raise ValueError(f"Unsupported local backend type: {backend_type}") + + +def settings_to_local_backend( + settings: MatchboxLocalSettings, +) -> MatchboxLocalDBAdapter: + """Create a local backend adapter with injected settings.""" + BackendClass = get_local_backend_class(settings.backend_type) + return BackendClass(path=settings.path) diff --git a/src/matchbox/common/adapters/protocol.py b/src/matchbox/common/adapters/protocol.py index 52524602..ae814137 100644 --- a/src/matchbox/common/adapters/protocol.py +++ b/src/matchbox/common/adapters/protocol.py @@ -3,7 +3,7 @@ import json from abc import ABC, abstractmethod from enum import StrEnum -from typing import Any, Protocol +from typing import Any, Protocol, TypeAlias from pyarrow import Table from pydantic import BaseModel, field_validator @@ -40,12 +40,22 @@ class ListableAndCountable(Countable, Listable): pass -class MatchboxBackends(StrEnum): - """The available backends for Matchbox.""" +class MatchboxServerBackends(StrEnum): + """The available server backends for Matchbox.""" POSTGRES = "postgres" +class MatchboxLocalBackends(StrEnum): + """The available local (client-side) backends for Matchbox.""" + + DUCKDB = "duckdb" + + +MatchboxBackends: TypeAlias = MatchboxServerBackends | MatchboxLocalBackends +"""Any backend, server or local, that can produce a MatchboxSnapshot.""" + + class MatchboxSnapshot(BaseModel): """A snapshot of the Matchbox database.""" diff --git a/src/matchbox/common/adapters/sql/insert.py b/src/matchbox/common/adapters/sql/insert.py new file mode 100644 index 00000000..2967317e --- /dev/null +++ b/src/matchbox/common/adapters/sql/insert.py @@ -0,0 +1,74 @@ +"""Ingest select builders, shared across relational backends. + +literal() values use the generic BigInteger type, never a +dialect-specific one (e.g. postgresql.BIGINT). +""" + +from sqlalchemy import BigInteger, exists, func, literal, select +from sqlalchemy.sql.expression import TableClause +from sqlalchemy.sql.selectable import Select, Subquery + +from matchbox.common.adapters.sql import tables + + +def select_new_cluster_hashes(incoming: TableClause, hash_col: str = "hash") -> Select: + """Distinct hashes in incoming absent from Clusters (anti-join).""" + hash_column = incoming.c[hash_col] + return ( + select(hash_column) + .distinct() + .where(~exists(select(1).where(tables.Clusters.c.cluster_hash == hash_column))) + ) + + +def select_cluster_map(incoming_hashes: TableClause) -> Subquery: + """(parent_id, cluster_id) by joining staged hashes to Clusters.""" + return ( + select(incoming_hashes.c.parent_id, tables.Clusters.c.cluster_id) + .select_from( + incoming_hashes.join( + tables.Clusters, + tables.Clusters.c.cluster_hash == incoming_hashes.c.cluster_hash, + ) + ) + .subquery("cluster_map") + ) + + +def select_key_expansion(incoming: TableClause, source_config_id: int) -> Select: + """(cluster_id, source_config_id, key) with keys unnested.""" + return select( + tables.Clusters.c.cluster_id, + literal(source_config_id, BigInteger).label("source_config_id"), + func.unnest(incoming.c["keys"]).label("key"), + ).select_from( + incoming.join( + tables.Clusters, tables.Clusters.c.cluster_hash == incoming.c.hash + ) + ) + + +def select_contains_pairs(expanded_leaves: Subquery, cluster_map: Subquery) -> Select: + """Distinct (root, leaf) pairs for the Contains insert.""" + return ( + select( + cluster_map.c.cluster_id.label("root"), + expanded_leaves.c.leaf_id.label("leaf"), + ) + .select_from( + expanded_leaves.join( + cluster_map, + expanded_leaves.c.parent_id == cluster_map.c.parent_id, + ) + ) + .where(cluster_map.c.cluster_id != expanded_leaves.c.leaf_id) + .distinct() + ) + + +def select_resolver_membership(step_id: int, cluster_map: Subquery) -> Select: + """Distinct (step_id, cluster_id) for ResolverClusters.""" + return select( + literal(step_id, BigInteger).label("step_id"), + cluster_map.c.cluster_id, + ).distinct() diff --git a/src/matchbox/common/adapters/sql/query.py b/src/matchbox/common/adapters/sql/query.py index 113be656..e6cc2b15 100644 --- a/src/matchbox/common/adapters/sql/query.py +++ b/src/matchbox/common/adapters/sql/query.py @@ -11,11 +11,12 @@ from typing import Literal -from sqlalchemy import and_, func, join, select +from sqlalchemy import Row, and_, func, join, select from sqlalchemy.sql.elements import ColumnElement from sqlalchemy.sql.selectable import CTE, Select, Subquery from matchbox.common.adapters.sql import tables +from matchbox.common.dtos import Match, SourceStepPath def build_unified_query( @@ -197,3 +198,36 @@ def resolver_membership_subquery( # UNION deduplicates in case a root cluster also appears as a leaf return roots_query.union(leaves_query).subquery(alias) + + +def assemble_matches( + matched_rows: list[Row], + source: SourceStepPath, + source_config_id: int, + targets: list[SourceStepPath], + target_config_ids: list[int], +) -> list[Match]: + """Assemble Match DTOs from (cluster_id, source_config_id, key) rows. + + Groups keys by source_config_id and builds one Match per target, + defaulting to an empty set when no keys were found for that target. + """ + cluster: int | None = None + matches_by_source_id: dict[int, set[str]] = {} + for cluster_id, source_config_id_result, key_in_source in matched_rows: + if cluster is None: + cluster = cluster_id + matches_by_source_id.setdefault(source_config_id_result, set()).add( + key_in_source + ) + + return [ + Match( + cluster=cluster, + source=source, + source_id=matches_by_source_id.get(source_config_id, set()), + target=target, + target_id=matches_by_source_id.get(target_config_id, set()), + ) + for target, target_config_id in zip(targets, target_config_ids, strict=False) + ] diff --git a/src/matchbox/common/adapters/sql/resolver.py b/src/matchbox/common/adapters/sql/resolver.py index 42ef9c3e..71f22d5e 100644 --- a/src/matchbox/common/adapters/sql/resolver.py +++ b/src/matchbox/common/adapters/sql/resolver.py @@ -1,10 +1,4 @@ -"""Resolver cluster canonicalisation, shared across relational backends. - -Expanding a resolver's incoming (parent_id, child_id) assignments to their -leaf-level cluster IDs, and deriving a single composite hash per parent -cluster from its sorted leaf hashes, so identical parent clusters dedupe to -one row regardless of which backend computed them. -""" +"""Resolver cluster canonicalisation, shared across relational backends.""" from collections.abc import Iterable @@ -45,8 +39,7 @@ def build_leaf_hash_groups_query(expanded_leaves: Subquery) -> Select: """Build a query grouping leaf hashes per parent cluster. Uses an inner join to the clusters table, so unknown leaf IDs are - silently dropped here. Callers relying on validation should catch this - downstream, for example as FK violations when inserting into contains. + silently dropped. """ return ( select( diff --git a/src/matchbox/common/adapters/sql/snapshot.py b/src/matchbox/common/adapters/sql/snapshot.py new file mode 100644 index 00000000..ce720327 --- /dev/null +++ b/src/matchbox/common/adapters/sql/snapshot.py @@ -0,0 +1,86 @@ +"""Generic table dump/restore, shared across relational backends. + +Handles only data movement: selecting every row and base64-encoding bytes +columns for JSON-safety, or decoding and bulk-inserting them back. +Engine-specific concerns stay with the caller. +""" + +import base64 +from typing import Any + +from sqlalchemy import Table, insert, select +from sqlalchemy.orm import Session + + +def dump_tables( + session: Session, tables: list[Table] +) -> dict[str, list[dict[str, Any]]]: + """Dump every row of the given tables to a JSON-safe dict, keyed by table name. + + Args: + session: Database session. + tables: The tables to dump, in any order. + + Returns: + A dict mapping table name to its rows, each row a dict of column + name to value. Bytes values are wrapped as {"base64": ...}. + """ + data: dict[str, list[dict[str, Any]]] = {} + + for table in tables: + records = session.execute(select(table)).mappings().all() + + table_data = [] + for record in records: + record_dict = dict(record) + for key, value in record_dict.items(): + if isinstance(value, bytes): + record_dict[key] = { + "base64": base64.b64encode(value).decode("ascii") + } + table_data.append(record_dict) + + data[table.name] = table_data + + return data + + +def restore_tables( + session: Session, + tables: list[Table], + data: dict[str, list[dict[str, Any]]], + batch_size: int = 10_000, +) -> None: + """Restore rows into the given tables, in order, from a dump_tables() dict. + + Args: + session: Database session. + tables: The tables to restore, in dependency order (parents first). + data: A dict as produced by dump_tables(). + batch_size: The number of records to insert per batch. + + Raises: + ValueError: If a table isn't present in data. + """ + for table in tables: + if table.name not in data: + raise ValueError(f"Invalid: Table {table.name} not found in snapshot.") + + records = data[table.name] + if not records: + continue + + processed_records = [] + for record in records: + processed_record = {} + for key, value in record.items(): + if isinstance(value, dict) and "base64" in value: + processed_record[key] = base64.b64decode(value["base64"]) + else: + processed_record[key] = value + processed_records.append(processed_record) + + for i in range(0, len(processed_records), batch_size): + batch = processed_records[i : i + batch_size] + session.execute(insert(table), batch) + session.flush() diff --git a/src/matchbox/common/adapters/sql/tables.py b/src/matchbox/common/adapters/sql/tables.py index aae147c6..0ccbbc4c 100644 --- a/src/matchbox/common/adapters/sql/tables.py +++ b/src/matchbox/common/adapters/sql/tables.py @@ -22,10 +22,6 @@ reference tables.Clusters rather than importing names directly. These are Table instances, not classes: the PascalCase name signals the pairing, not that they can be subclassed or instantiated. - -TODO: once a second backend exists, add dedicated tests for this module -and the rest of matchbox.common.adapters.sql, rather than relying solely -on the PostgreSQL adapter test suite. """ from sqlalchemy import ( @@ -33,6 +29,7 @@ BigInteger, CheckConstraint, Column, + FetchedValue, ForeignKey, Index, LargeBinary, @@ -53,7 +50,13 @@ Clusters = Table( "clusters", METADATA, - Column("cluster_id", BigInteger, primary_key=True), + Column( + "cluster_id", + BigInteger, + primary_key=True, + autoincrement=False, + server_default=FetchedValue(), + ), Column("cluster_hash", LargeBinary, nullable=False), UniqueConstraint("cluster_hash", name="clusters_hash_key"), ) @@ -62,7 +65,13 @@ ClusterSourceKey = Table( "cluster_keys", METADATA, - Column("key_id", BigInteger, primary_key=True), + Column( + "key_id", + BigInteger, + primary_key=True, + autoincrement=False, + server_default=FetchedValue(), + ), Column( "cluster_id", BigInteger, @@ -108,7 +117,13 @@ ModelEdges = Table( "model_edges", METADATA, - Column("result_id", BigInteger, primary_key=True, autoincrement=True), + Column( + "result_id", + BigInteger, + primary_key=True, + autoincrement=False, + server_default=FetchedValue(), + ), Column( "step_id", BigInteger, diff --git a/src/matchbox/common/factories/scenarios.py b/src/matchbox/common/factories/scenarios.py index 82b56f54..ba73f764 100644 --- a/src/matchbox/common/factories/scenarios.py +++ b/src/matchbox/common/factories/scenarios.py @@ -25,6 +25,7 @@ from polars.testing import assert_frame_equal from sqlalchemy import Engine +from matchbox.client.base import MatchboxLocalDBAdapter from matchbox.client.queries import Query from matchbox.common.adapters.protocol import ( MatchboxClusterStoreAdapter, @@ -73,7 +74,7 @@ def decorator(func: ScenarioBuilder) -> ScenarioBuilder: def _generate_cache_key( - backend: MatchboxDBAdapter, + backend: MatchboxClusterStoreAdapter, scenario_type: str, warehouse: Engine, n_entities: int = 10, @@ -145,7 +146,10 @@ def create_bare_scenario( Scope: EITHER, since it touches nothing on backend. """ - return TestkitDAG() + dag_testkit = TestkitDAG() + if isinstance(backend, MatchboxLocalDBAdapter): + backend.bind(dag_testkit.dag) + return dag_testkit @register_scenario("admin") @@ -1254,7 +1258,7 @@ def _categorical_feature(name: str, *choices: str, **kwargs: Any) -> FeatureConf @contextmanager def setup_scenario( - backend: MatchboxDBAdapter, + backend: MatchboxClusterStoreAdapter, scenario_type: Literal[ "bare", "admin", @@ -1293,9 +1297,11 @@ def setup_scenario( # our new warehouse object dag_testkit = dag_testkit.model_copy(deep=True) - # Restore backend and write sources to warehouse + # Restore backend, write sources to warehouse, rebind if local backend.restore(snapshot=snapshot) _testkitdag_to_location(warehouse, dag_testkit) + if isinstance(backend, MatchboxLocalDBAdapter): + backend.bind(dag_testkit.dag) else: # Create new TestkitDAG with proper backend integration scenario_builder = SCENARIO_REGISTRY[scenario_type] diff --git a/src/matchbox/server/base.py b/src/matchbox/server/base.py index e9232327..c427deee 100644 --- a/src/matchbox/server/base.py +++ b/src/matchbox/server/base.py @@ -19,8 +19,8 @@ from matchbox.common.adapters.protocol import ( Countable, ListableAndCountable, - MatchboxBackends, MatchboxClusterStoreAdapter, + MatchboxServerBackends, ) from matchbox.common.dtos import ( BackendResourceType, @@ -164,7 +164,7 @@ class MatchboxServerSettings(BaseSettings): ) batch_size: int = Field(default=250_000) - backend_type: MatchboxBackends + backend_type: MatchboxServerBackends datastore: MatchboxDatastoreSettings task_runner: Literal["api", "celery"] redis_uri: str | None @@ -242,10 +242,10 @@ def get_settings(cls) -> "MatchboxServerSettings": def get_backend_settings( - backend_type: MatchboxBackends, + backend_type: MatchboxServerBackends, ) -> type[MatchboxServerSettings]: """Get the appropriate settings class based on the backend type.""" - if backend_type == MatchboxBackends.POSTGRES: + if backend_type == MatchboxServerBackends.POSTGRES: from matchbox.server.postgresql import MatchboxPostgresSettings # noqa: PLC0415 return MatchboxPostgresSettings @@ -254,9 +254,11 @@ def get_backend_settings( raise ValueError(f"Unsupported backend type: {backend_type}") -def get_backend_class(backend_type: MatchboxBackends) -> type["MatchboxDBAdapter"]: +def get_backend_class( + backend_type: MatchboxServerBackends, +) -> type["MatchboxDBAdapter"]: """Get the appropriate backend class based on the backend type.""" - if backend_type == MatchboxBackends.POSTGRES: + if backend_type == MatchboxServerBackends.POSTGRES: from matchbox.server.postgresql import MatchboxPostgres # noqa: PLC0415 return MatchboxPostgres diff --git a/src/matchbox/server/postgresql/adapter/admin.py b/src/matchbox/server/postgresql/adapter/admin.py index a3935843..5f9e9193 100644 --- a/src/matchbox/server/postgresql/adapter/admin.py +++ b/src/matchbox/server/postgresql/adapter/admin.py @@ -26,8 +26,8 @@ from matchbox.server.base import PERMISSION_GRANTS from matchbox.server.postgresql.db import ( MBDB, - MatchboxBackends, MatchboxPostgresSettings, + MatchboxServerBackends, ) from matchbox.server.postgresql.orm import ( Clusters, @@ -365,7 +365,7 @@ def clear(self, certain: bool) -> None: # noqa: D102 ) def restore(self, snapshot: MatchboxSnapshot) -> None: # noqa: D102 - if snapshot.backend_type != MatchboxBackends.POSTGRES: + if snapshot.backend_type != MatchboxServerBackends.POSTGRES: raise TypeError( f"Cannot restore {snapshot.backend_type} snapshot to PostgreSQL backend" ) diff --git a/src/matchbox/server/postgresql/db.py b/src/matchbox/server/postgresql/db.py index ce501273..9f322bce 100644 --- a/src/matchbox/server/postgresql/db.py +++ b/src/matchbox/server/postgresql/db.py @@ -25,7 +25,7 @@ from matchbox.common.adapters.sql.tables import METADATA from matchbox.common.datatypes import require from matchbox.common.logging import logger -from matchbox.server.base import MatchboxBackends, MatchboxServerSettings +from matchbox.server.base import MatchboxServerBackends, MatchboxServerSettings class MatchboxPostgresCoreSettings(BaseModel): @@ -64,7 +64,7 @@ class MatchboxPostgresSettings(MatchboxServerSettings): Inherits the core settings and adds the PostgreSQL-specific settings. """ - backend_type: MatchboxBackends = MatchboxBackends.POSTGRES + backend_type: MatchboxServerBackends = MatchboxServerBackends.POSTGRES postgres: MatchboxPostgresCoreSettings = Field( default_factory=MatchboxPostgresCoreSettings diff --git a/src/matchbox/server/postgresql/utils/db.py b/src/matchbox/server/postgresql/utils/db.py index 8b65f305..fa7b3558 100644 --- a/src/matchbox/server/postgresql/utils/db.py +++ b/src/matchbox/server/postgresql/utils/db.py @@ -18,7 +18,7 @@ from sqlalchemy.sql import Select from sqlalchemy.sql.type_api import TypeEngine -from matchbox.common.adapters.protocol import MatchboxBackends, MatchboxSnapshot +from matchbox.common.adapters.protocol import MatchboxServerBackends, MatchboxSnapshot from matchbox.common.datatypes import require from matchbox.common.dtos import ( BackendResourceType, @@ -64,7 +64,7 @@ def dump() -> MatchboxSnapshot: data[table.name] = table_data - return MatchboxSnapshot(backend_type=MatchboxBackends.POSTGRES, data=data) + return MatchboxSnapshot(backend_type=MatchboxServerBackends.POSTGRES, data=data) def restore(snapshot: MatchboxSnapshot, batch_size: int) -> None: diff --git a/src/matchbox/server/postgresql/utils/insert.py b/src/matchbox/server/postgresql/utils/insert.py index 253bf3dd..24f58107 100644 --- a/src/matchbox/server/postgresql/utils/insert.py +++ b/src/matchbox/server/postgresql/utils/insert.py @@ -1,7 +1,7 @@ """Utilities for inserting data into the PostgreSQL backend.""" import pyarrow as pa -from sqlalchemy import exists, func, join, literal, select +from sqlalchemy import BigInteger, func, join, literal, select from sqlalchemy.dialects.postgresql import ( ARRAY, BIGINT, @@ -13,6 +13,13 @@ from sqlalchemy.orm import Session from sqlalchemy.sql.expression import TableClause +from matchbox.common.adapters.sql.insert import ( + select_cluster_map, + select_contains_pairs, + select_key_expansion, + select_new_cluster_hashes, + select_resolver_membership, +) from matchbox.common.adapters.sql.resolver import ( build_expanded_leaves_subquery, build_leaf_hash_groups_query, @@ -99,17 +106,9 @@ def insert_hashes(path: SourceStepPath, data_hashes: pa.Table, batch_size: int) ): try: # Add clusters - new_hashes = ( - select(incoming.c.hash) - .distinct() - .where( - ~exists(select(1).where(Clusters.cluster_hash == incoming.c.hash)) - ) - ) - stmt_insert_clusters = ( insert(Clusters) - .from_select(["cluster_hash"], new_hashes) + .from_select(["cluster_hash"], select_new_cluster_hashes(incoming)) .on_conflict_do_nothing(index_elements=[Clusters.cluster_hash]) .returning(Clusters.cluster_id) ) @@ -127,19 +126,11 @@ def insert_hashes(path: SourceStepPath, data_hashes: pa.Table, batch_size: int) session.flush() # Add source keys - exploded = select( - Clusters.cluster_id, - literal(source_config_id, BIGINT).label("source_config_id"), - func.unnest(incoming.c["keys"]).label("key"), - ).select_from( - incoming.join(Clusters, Clusters.cluster_hash == incoming.c.hash) - ) - stmt_insert_keys = ( insert(ClusterSourceKey) .from_select( ["cluster_id", "source_config_id", "key"], - exploded, + select_key_expansion(incoming, source_config_id), ) .returning(ClusterSourceKey.key_id) ) @@ -215,17 +206,16 @@ def insert_model_edges( ) as incoming_edges, ): try: - edges_select = select( - literal(step.step_id, BIGINT).label("step_id"), - incoming_edges.c.left_id, - incoming_edges.c.right_id, - incoming_edges.c.score, - ) stmt_insert_edges = ( insert(ModelEdges) .from_select( ["step_id", "left_id", "right_id", "score"], - edges_select, + select( + literal(step.step_id, BigInteger).label("step_id"), + incoming_edges.c.left_id, + incoming_edges.c.right_id, + incoming_edges.c.score, + ), ) .returning(ModelEdges.result_id) ) @@ -382,15 +372,8 @@ def insert_clusters( insert(Clusters) .from_select( ["cluster_hash"], - select(incoming_hashes.c.cluster_hash) - .distinct() - .where( - ~exists( - select(1).where( - Clusters.cluster_hash - == incoming_hashes.c.cluster_hash - ) - ) + select_new_cluster_hashes( + incoming_hashes, hash_col="cluster_hash" ), ) .on_conflict_do_nothing(index_elements=[Clusters.cluster_hash]) @@ -399,19 +382,7 @@ def insert_clusters( # Map each parent_id to its canonical Clusters.cluster_id # by joining hashes back to the now-populated Clusters table - cluster_map = ( - select( - incoming_hashes.c.parent_id, - Clusters.cluster_id, - ) - .select_from( - incoming_hashes.join( - Clusters, - Clusters.cluster_hash == incoming_hashes.c.cluster_hash, - ) - ) - .subquery("cluster_map") - ) + cluster_map = select_cluster_map(incoming_hashes) # Contains # Record which leaves belong to each new resolver cluster @@ -419,17 +390,7 @@ def insert_clusters( insert(Contains) .from_select( ["root", "leaf"], - select( - cluster_map.c.cluster_id, - expanded_leaves.c.leaf_id, - ) - .select_from( - expanded_leaves.join( - cluster_map, - expanded_leaves.c.parent_id == cluster_map.c.parent_id, - ) - ) - .distinct(), + select_contains_pairs(expanded_leaves, cluster_map), ) .on_conflict_do_nothing( index_elements=[Contains.root, Contains.leaf] @@ -442,10 +403,7 @@ def insert_clusters( insert(ResolverClusters) .from_select( ["step_id", "cluster_id"], - select( - literal(step_id, BIGINT).label("step_id"), - cluster_map.c.cluster_id, - ).distinct(), + select_resolver_membership(step_id, cluster_map), ) .on_conflict_do_nothing( index_elements=[ diff --git a/src/matchbox/server/uploads.py b/src/matchbox/server/uploads.py index cecfe13b..4638d42d 100644 --- a/src/matchbox/server/uploads.py +++ b/src/matchbox/server/uploads.py @@ -26,8 +26,8 @@ from matchbox.common.exceptions import MatchboxServerFileError from matchbox.common.logging import configure_celery_logging, log_mem_usage, logger from matchbox.server.base import ( - MatchboxBackends, MatchboxDBAdapter, + MatchboxServerBackends, MatchboxServerSettings, get_backend_settings, settings_to_backend, @@ -274,7 +274,7 @@ def process_upload_celery( # If using Postgres, we must reset the global database connections # to avoid using inherited C-pointers from the parent process (ADBC). - if CELERY_SETTINGS.backend_type == MatchboxBackends.POSTGRES: + if CELERY_SETTINGS.backend_type == MatchboxServerBackends.POSTGRES: from matchbox.server.postgresql.db import MBDB # noqa: PLC0415 MBDB._disconnect_adbc() @@ -314,7 +314,7 @@ def process_upload_celery( log_mem_usage() # Cleanup connections - if CELERY_SETTINGS.backend_type == MatchboxBackends.POSTGRES: + if CELERY_SETTINGS.backend_type == MatchboxServerBackends.POSTGRES: from matchbox.server.postgresql.db import MBDB # noqa: PLC0415 MBDB._disconnect() diff --git a/test/adapters/__init__.py b/test/adapters/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/server/adapter/test_adapter_query.py b/test/adapters/test_adapter_query.py similarity index 96% rename from test/server/adapter/test_adapter_query.py rename to test/adapters/test_adapter_query.py index a9a6697c..4ddf2c96 100644 --- a/test/server/adapter/test_adapter_query.py +++ b/test/adapters/test_adapter_query.py @@ -6,23 +6,24 @@ import pyarrow.compute as pc import pytest from sqlalchemy import Engine -from test.fixtures.db import SERVER_BACKENDS +from matchbox.common.adapters.protocol import MatchboxClusterStoreAdapter from matchbox.common.arrow import SCHEMA_QUERY, SCHEMA_QUERY_WITH_LEAVES from matchbox.common.dtos import Match from matchbox.common.factories.entities import SourceEntity from matchbox.common.factories.scenarios import setup_scenario -from matchbox.server.base import MatchboxDBAdapter +from test.fixtures.db import CLUSTER_STORES -@pytest.mark.parametrize("backend", SERVER_BACKENDS) -@pytest.mark.docker +@pytest.mark.parametrize("backend", CLUSTER_STORES) class TestMatchboxQueryBackend: @pytest.fixture(autouse=True) def setup( - self, backend_instance: MatchboxDBAdapter, sqla_sqlite_warehouse: Engine + self, + backend_instance: MatchboxClusterStoreAdapter, + sqla_sqlite_warehouse: Engine, ) -> None: - self.backend: MatchboxDBAdapter = backend_instance + self.backend: MatchboxClusterStoreAdapter = backend_instance self.scenario = partial(setup_scenario, warehouse=sqla_sqlite_warehouse) def test_query_only_source(self) -> None: diff --git a/test/client/adapters/__init__.py b/test/client/adapters/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/client/adapters/duckdb/__init__.py b/test/client/adapters/duckdb/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/client/adapters/duckdb/test_duckdb_core.py b/test/client/adapters/duckdb/test_duckdb_core.py new file mode 100644 index 00000000..eec5b076 --- /dev/null +++ b/test/client/adapters/duckdb/test_duckdb_core.py @@ -0,0 +1,57 @@ +import pyarrow as pa +from sqlalchemy import Engine + +from matchbox.client.adapters.duckdb import MatchboxLocalDuckDB +from matchbox.common.factories.scenarios import setup_scenario + + +def test_drop_leaves_no_dynamic_tables( + matchbox_local_duckdb: MatchboxLocalDuckDB, sqla_sqlite_warehouse: Engine +) -> None: + """drop() sweeps RawData/QueryCache's dynamically named physical tables.""" + with setup_scenario( + matchbox_local_duckdb, "index", warehouse=sqla_sqlite_warehouse + ) as dag_testkit: + crn = dag_testkit.sources["crn"].path + matchbox_local_duckdb.insert_raw_data( + crn, pa.table({"key": ["k1"], "name": ["Alice"]}) + ) + matchbox_local_duckdb.cache_query( + "key1", pa.table({"id": [1]}), depends_on=[crn] + ) + + matchbox_local_duckdb.drop(certain=True) + + with matchbox_local_duckdb._engine.connect() as conn: + names = { + row[0] + for row in conn.exec_driver_sql( + "SELECT table_name FROM duckdb_tables()" + ).fetchall() + } + prefixes = ("mb_raw_data__", "mb_query_cache__") + assert not any(n.startswith(prefixes) for n in names) + + +def test_dump_restore_round_trips_dynamic_tables( + matchbox_local_duckdb: MatchboxLocalDuckDB, sqla_sqlite_warehouse: Engine +) -> None: + """dump()/restore() carry RawData/QueryCache row content, not just pointers.""" + with setup_scenario( + matchbox_local_duckdb, "index", warehouse=sqla_sqlite_warehouse + ) as dag_testkit: + crn = dag_testkit.sources["crn"].path + matchbox_local_duckdb.insert_raw_data( + crn, pa.table({"key": ["k1"], "name": ["Alice"]}) + ) + matchbox_local_duckdb.cache_query( + "key1", pa.table({"id": [1]}), depends_on=[crn] + ) + snapshot = matchbox_local_duckdb.dump() + + matchbox_local_duckdb.restore(snapshot) + + assert matchbox_local_duckdb.get_raw_data(crn).to_pylist() == [ + {"key": "k1", "name": "Alice"} + ] + assert matchbox_local_duckdb.get_cached_query("key1").to_pylist() == [{"id": 1}] diff --git a/test/client/local/__init__.py b/test/client/local/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/client/local/test_cache_mixin.py b/test/client/local/test_cache_mixin.py new file mode 100644 index 00000000..00793284 --- /dev/null +++ b/test/client/local/test_cache_mixin.py @@ -0,0 +1,192 @@ +"""Tests for MatchboxLocalDuckDBCacheMixin: raw data, query cache, cascade.""" + +from functools import partial + +import pyarrow as pa +import pytest +from sqlalchemy import Engine + +from matchbox.client.base import MatchboxLocalDBAdapter +from matchbox.common.dtos import CollectionName, RunID, StepPath +from matchbox.common.exceptions import MatchboxDataNotFound, MatchboxStepNotFoundError +from matchbox.common.factories.scenarios import setup_scenario +from test.fixtures.db import LOCAL_BACKENDS + +# Placeholder collection/run: local ignores both, only the step name matters. +_NONEXISTENT_STEP = StepPath( + collection=CollectionName("local_test"), run=RunID(1), name="nonexistent" +) + + +@pytest.mark.parametrize("backend", LOCAL_BACKENDS) +class TestRawData: + @pytest.fixture(autouse=True) + def setup( + self, backend_instance: MatchboxLocalDBAdapter, sqla_sqlite_warehouse: Engine + ) -> None: + self.backend: MatchboxLocalDBAdapter = backend_instance + self.scenario = partial(setup_scenario, warehouse=sqla_sqlite_warehouse) + + def test_round_trip(self) -> None: + """Insertion and retrieval of raw data returns that data.""" + with self.scenario(self.backend, "index") as dag_testkit: + crn = dag_testkit.sources["crn"].path + + table = pa.table({"key": ["k1", "k2"], "name": ["Alice", "Bob"]}) + self.backend.insert_raw_data(crn, table) + + fetched = self.backend.get_raw_data(crn) + assert sorted(fetched.to_pylist(), key=lambda r: r["key"]) == sorted( + table.to_pylist(), key=lambda r: r["key"] + ) + + def test_filters_by_keys(self) -> None: + """We can filter retrieved raw data by keys.""" + with self.scenario(self.backend, "index") as dag_testkit: + crn = dag_testkit.sources["crn"].path + + table = pa.table( + {"key": ["k1", "k2", "k3"], "name": ["Alice", "Bob", "Carl"]} + ) + self.backend.insert_raw_data(crn, table) + + fetched = self.backend.get_raw_data(crn, keys=["k2"]) + assert fetched.to_pylist() == [{"key": "k2", "name": "Bob"}] + + def test_filters_by_empty_keys(self) -> None: + """Filtering raw data by empty keys returns nothing.""" + with self.scenario(self.backend, "index") as dag_testkit: + crn = dag_testkit.sources["crn"].path + self.backend.insert_raw_data( + crn, pa.table({"key": ["k1"], "name": ["Alice"]}) + ) + + fetched = self.backend.get_raw_data(crn, keys=[]) + assert fetched.num_rows == 0 + + def test_missing_raises(self) -> None: + """Fetching uncached raw data raises.""" + with self.scenario(self.backend, "index") as dag_testkit: + dh = dag_testkit.sources["dh"].path + + with pytest.raises(MatchboxDataNotFound): + self.backend.get_raw_data(dh) + + def test_unknown_step_raises(self) -> None: + """Fetching raw data from an unregistered step raises.""" + with pytest.raises(MatchboxStepNotFoundError): + self.backend.get_raw_data(_NONEXISTENT_STEP) + + def test_replaces(self) -> None: + """Reinserting replaces raw data.""" + with self.scenario(self.backend, "index") as dag_testkit: + crn = dag_testkit.sources["crn"].path + + self.backend.insert_raw_data( + crn, pa.table({"key": ["k1"], "name": ["Alice"]}) + ) + self.backend.insert_raw_data( + crn, pa.table({"key": ["k2"], "name": ["Bob"]}) + ) + + fetched = self.backend.get_raw_data(crn) + assert fetched.to_pylist() == [{"key": "k2", "name": "Bob"}] + + +@pytest.mark.parametrize("backend", LOCAL_BACKENDS) +class TestQueryCache: + @pytest.fixture(autouse=True) + def setup( + self, backend_instance: MatchboxLocalDBAdapter, sqla_sqlite_warehouse: Engine + ) -> None: + self.backend: MatchboxLocalDBAdapter = backend_instance + self.scenario = partial(setup_scenario, warehouse=sqla_sqlite_warehouse) + + def test_round_trip(self) -> None: + """Insertion and retrieval of query data returns that data.""" + with self.scenario(self.backend, "index") as dag_testkit: + crn = dag_testkit.sources["crn"].path + + table = pa.table({"id": [1, 2], "value": ["a", "b"]}) + self.backend.cache_query("key1", table, depends_on=[crn]) + + fetched = self.backend.get_cached_query("key1") + assert sorted(fetched.to_pylist(), key=lambda r: r["id"]) == sorted( + table.to_pylist(), key=lambda r: r["id"] + ) + + def test_missing_returns_none(self) -> None: + """Fetching uncached query data returns None.""" + assert self.backend.get_cached_query("missing") is None + + def test_replaces(self) -> None: + """Reinserting replaces query data.""" + with self.scenario(self.backend, "index") as dag_testkit: + crn = dag_testkit.sources["crn"].path + + self.backend.cache_query("key1", pa.table({"id": [1]}), depends_on=[crn]) + self.backend.cache_query("key1", pa.table({"id": [2]}), depends_on=[crn]) + + fetched = self.backend.get_cached_query("key1") + assert fetched.to_pylist() == [{"id": 2}] + + def test_unknown_dependency_raises(self) -> None: + """Fetching query data from an unregistered dependency step raises.""" + with pytest.raises(MatchboxStepNotFoundError): + self.backend.cache_query( + "key1", pa.table({"id": [1]}), depends_on=[_NONEXISTENT_STEP] + ) + + +@pytest.mark.parametrize("backend", LOCAL_BACKENDS) +class TestDropStepData: + @pytest.fixture(autouse=True) + def setup( + self, backend_instance: MatchboxLocalDBAdapter, sqla_sqlite_warehouse: Engine + ) -> None: + self.backend: MatchboxLocalDBAdapter = backend_instance + self.scenario = partial(setup_scenario, warehouse=sqla_sqlite_warehouse) + + def test_cascades(self) -> None: + """Dropping a step's data cascades to its descendants.""" + with self.scenario(self.backend, "dedupe") as dag_testkit: + model = dag_testkit.models["naive_test_crn"] + resolver = dag_testkit.resolvers["resolver_naive_test_crn"] + + assert self.backend.get_model_data(model.path).num_rows > 0 + assert self.backend.get_resolver_data(resolver.resolver.path).num_rows > 0 + + self.backend.drop_step_data(model.path) + + assert self.backend.get_model_data(model.path).num_rows == 0 + assert self.backend.get_resolver_data(resolver.resolver.path).num_rows == 0 + + def test_clears_descendant_cache_only(self) -> None: + """Only the caches of descendents of the dropped step are cleared.""" + with self.scenario(self.backend, "dedupe") as dag_testkit: + crn = dag_testkit.sources["crn"].path + dh = dag_testkit.sources["dh"].path + + self.backend.cache_query( + "dependent", pa.table({"a": [1]}), depends_on=[crn] + ) + self.backend.cache_query( + "independent", pa.table({"a": [2]}), depends_on=[dh] + ) + + self.backend.drop_step_data(crn) + + assert self.backend.get_cached_query("dependent") is None + assert self.backend.get_cached_query("independent") is not None + + def test_raw_data_survives(self) -> None: + """drop_step_data doesn't touch RawData - it's canonical, not a cache.""" + with self.scenario(self.backend, "dedupe") as dag_testkit: + crn = dag_testkit.sources["crn"].path + self.backend.insert_raw_data( + crn, pa.table({"key": ["k1"], "name": ["Alice"]}) + ) + + self.backend.drop_step_data(crn) + + assert self.backend.get_raw_data(crn).num_rows == 1 diff --git a/test/client/local/test_data_mixin.py b/test/client/local/test_data_mixin.py new file mode 100644 index 00000000..c7fadb3d --- /dev/null +++ b/test/client/local/test_data_mixin.py @@ -0,0 +1,76 @@ +"""Tests for MatchboxLocalDuckDBDataMixin: replace-on-rerun. + +Local stores replace-on-rerun, the deliberate inversion of the server's +write-once rule: iteration is the point locally. +""" + +from functools import partial + +import pytest +from sqlalchemy import Engine + +from matchbox.client.base import MatchboxLocalDBAdapter +from matchbox.common.factories.scenarios import setup_scenario +from test.fixtures.db import LOCAL_BACKENDS + + +@pytest.mark.parametrize("backend", LOCAL_BACKENDS) +class TestReplaceOnRerun: + @pytest.fixture(autouse=True) + def setup( + self, backend_instance: MatchboxLocalDBAdapter, sqla_sqlite_warehouse: Engine + ) -> None: + self.backend: MatchboxLocalDBAdapter = backend_instance + self.scenario = partial(setup_scenario, warehouse=sqla_sqlite_warehouse) + + def test_reinsert_source_replaces_and_cascades(self) -> None: + """Re-running a source's insert clears downstream model/resolver data.""" + with self.scenario(self.backend, "dedupe") as dag_testkit: + crn = dag_testkit.sources["crn"] + model = dag_testkit.models["naive_test_crn"] + resolver = dag_testkit.resolvers["resolver_naive_test_crn"] + + assert self.backend.get_model_data(model.path).num_rows > 0 + assert self.backend.get_resolver_data(resolver.resolver.path).num_rows > 0 + + self.backend.insert_source_data(path=crn.path, data_hashes=crn.data_hashes) + + assert self.backend.get_model_data(model.path).num_rows == 0 + assert self.backend.get_resolver_data(resolver.resolver.path).num_rows == 0 + + # But the source itself still has data, correctly replaced + assert self.backend.query(source=crn.path).num_rows == crn.data.num_rows + + def test_reinsert_model_replaces_own_data_and_cascades(self) -> None: + """Re-running a model's insert replaces its own edges and cascades.""" + with self.scenario(self.backend, "dedupe") as dag_testkit: + model = dag_testkit.models["naive_test_crn"] + resolver = dag_testkit.resolvers["resolver_naive_test_crn"] + + before = self.backend.get_model_data(model.path).num_rows + assert before > 0 + + self.backend.insert_model_data( + path=model.path, results=model.scores.to_arrow() + ) + + assert self.backend.get_model_data(model.path).num_rows == before + # Cascaded: the downstream resolver's data is gone + assert self.backend.get_resolver_data(resolver.resolver.path).num_rows == 0 + + def test_reinsert_resolver_replaces_own_data(self) -> None: + """Re-running a resolver's insert replaces its own assignments.""" + with self.scenario(self.backend, "dedupe") as dag_testkit: + resolver = dag_testkit.resolvers["resolver_naive_test_crn"] + + before = self.backend.get_resolver_data(resolver.resolver.path).num_rows + assert before > 0 + + self.backend.insert_resolver_data( + path=resolver.resolver.path, data=resolver.resolver.results.to_arrow() + ) + + assert ( + self.backend.get_resolver_data(resolver.resolver.path).num_rows + == before + ) diff --git a/test/fixtures/client.py b/test/fixtures/client.py index 2f05fbe0..bda20e54 100644 --- a/test/fixtures/client.py +++ b/test/fixtures/client.py @@ -13,7 +13,7 @@ from matchbox.client._settings import ClientSettings from matchbox.client._settings import settings as client_settings from matchbox.server.api import app, dependencies -from matchbox.server.base import MatchboxBackends, MatchboxServerSettings +from matchbox.server.base import MatchboxServerBackends, MatchboxServerSettings from matchbox.server.uploads import InMemoryUploadTracker from test.scripts.authorisation import ( generate_EdDSA_key_pair, @@ -89,7 +89,7 @@ def api_client_and_mocks( # 4) Override server settings used by API test_settings = MatchboxServerSettings( - backend_type=MatchboxBackends.POSTGRES, + backend_type=MatchboxServerBackends.POSTGRES, task_runner="api", authorisation=True, public_key=SecretBytes(public_key), diff --git a/test/fixtures/db.py b/test/fixtures/db.py index f75fbebd..70f07ad1 100644 --- a/test/fixtures/db.py +++ b/test/fixtures/db.py @@ -15,7 +15,9 @@ from pydantic_settings import BaseSettings, SettingsConfigDict from sqlalchemy import Engine, MetaData, create_engine, text -from matchbox.server.base import MatchboxDatastoreSettings, MatchboxDBAdapter +from matchbox.client.adapters.duckdb import MatchboxLocalDuckDB +from matchbox.common.adapters.protocol import MatchboxClusterStoreAdapter +from matchbox.server.base import MatchboxDatastoreSettings from matchbox.server.postgresql import MatchboxPostgres, MatchboxPostgresSettings from matchbox.server.postgresql.db import MBDB from matchbox.server.uploads import InMemoryUploadTracker, RedisUploadTracker @@ -463,17 +465,29 @@ def empty_tracker() -> None: # Backends SERVER_BACKENDS = [ - pytest.param("matchbox_postgres", id="postgres"), + pytest.param("matchbox_postgres", id="postgres", marks=pytest.mark.docker), ] -LOCAL_BACKENDS: list = [] -# TODO: populate once the local duckdb adapter lands +LOCAL_BACKENDS = [ + pytest.param("matchbox_local_duckdb", id="local_duckdb"), +] CLUSTER_STORES = SERVER_BACKENDS + LOCAL_BACKENDS @pytest.fixture(scope="function") -def backend_instance(request: pytest.FixtureRequest, backend: str) -> MatchboxDBAdapter: +def matchbox_local_duckdb() -> MatchboxLocalDuckDB: + """Fresh in-memory MatchboxLocalDuckDB for each test. + + An in-memory duckdb file is cheap enough to create fresh for every test. + """ + return MatchboxLocalDuckDB() + + +@pytest.fixture(scope="function") +def backend_instance( + request: pytest.FixtureRequest, backend: str +) -> MatchboxClusterStoreAdapter: """Create a fresh backend instance for each test.""" backend_obj = request.getfixturevalue(backend) backend_obj.clear(certain=True) diff --git a/uv.lock b/uv.lock index fe613622..5854ff3d 100644 --- a/uv.lock +++ b/uv.lock @@ -705,6 +705,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/44/48/ddf8d3740e3d28582944f70d84e720b5dc28c10ec22b668a0e0bd965f2f2/duckdb-1.5.4-cp313-cp313-win_arm64.whl", hash = "sha256:73f4878a3012283024a64a1909e440aac12091ef336f671fc142f7e87449ce0c", size = 13965189, upload-time = "2026-06-17T10:48:32.251Z" }, ] +[[package]] +name = "duckdb-engine" +version = "0.17.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "duckdb" }, + { name = "packaging" }, + { name = "sqlalchemy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/89/d5/c0d8d0a4ca3ffea92266f33d92a375e2794820ad89f9be97cf0c9a9697d0/duckdb_engine-0.17.0.tar.gz", hash = "sha256:396b23869754e536aa80881a92622b8b488015cf711c5a40032d05d2cf08f3cf", size = 48054, upload-time = "2025-03-29T09:49:17.663Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/a2/e90242f53f7ae41554419b1695b4820b364df87c8350aa420b60b20cab92/duckdb_engine-0.17.0-py3-none-any.whl", hash = "sha256:3aa72085e536b43faab635f487baf77ddc5750069c16a2f8d9c6c3cb6083e979", size = 49676, upload-time = "2025-03-29T09:49:15.564Z" }, +] + [[package]] name = "email-validator" version = "2.3.0" @@ -1340,6 +1354,7 @@ dependencies = [ { name = "click" }, { name = "cryptography" }, { name = "duckdb" }, + { name = "duckdb-engine" }, { name = "email-validator" }, { name = "faker" }, { name = "frozendict" }, @@ -1412,6 +1427,7 @@ requires-dist = [ { name = "click", specifier = ">=8.1.7" }, { name = "cryptography", specifier = ">=44.0.2" }, { name = "duckdb", specifier = ">=1.1.1" }, + { name = "duckdb-engine", specifier = ">=0.13.6" }, { name = "email-validator", specifier = ">=2.3.0" }, { name = "faker", specifier = ">=36.1.1" }, { name = "fastapi", extras = ["standard-no-fastapi-cloud-cli"], marker = "extra == 'server'", specifier = ">=0.116.0" },