diff --git a/.gitignore b/.gitignore index f02dc20b..123c7f4a 100644 --- a/.gitignore +++ b/.gitignore @@ -245,3 +245,6 @@ server/osa.yaml # MCP widget bundles staged into the server package (built from packages/osa-widgets) server/osa/application/api/mcp/bundles/ + +# graphify knowledge-graph outputs (generated caches, any directory) +graphify-out/ diff --git a/server/migrations/versions/a8c4e6f19b02_table_statistics.py b/server/migrations/versions/a8c4e6f19b02_table_statistics.py new file mode 100644 index 00000000..d4a3a82b --- /dev/null +++ b/server/migrations/versions/a8c4e6f19b02_table_statistics.py @@ -0,0 +1,79 @@ +"""``table_statistics`` — lockstep counts, backfilled once from the data. + +Creates the per-(schema version, table) count state and populates it from +COUNT(*) / COUNT(DISTINCT record_srn) group-bys over the existing tables. +This backfill and the admin verifier are the only sanctioned whole-table +counting after #219; from here on the writing adapters maintain the counts +transactionally (osa/infrastructure/persistence/statistics_upsert.py). + +Feature-table names come from the ``feature_tables`` catalog and are +re-validated against the strict identifier pattern before interpolation — +never string-built from user input. + +Revision ID: a8c4e6f19b02 +Revises: f3a1c9d27e54 +Create Date: 2026-08-15 +""" + +import re + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "a8c4e6f19b02" +down_revision = "f3a1c9d27e54" +branch_labels = None +depends_on = None + +_SAFE_IDENT = re.compile(r"^[a-z][a-z0-9_]{0,62}$") + + +def upgrade() -> None: + op.create_table( + "table_statistics", + sa.Column("schema_id", sa.Text(), primary_key=True), + sa.Column("schema_version", sa.Text(), primary_key=True), + sa.Column("table_name", sa.Text(), primary_key=True), + sa.Column("row_count", sa.BigInteger(), nullable=False, server_default="0"), + sa.Column("records_covered", sa.BigInteger(), nullable=True), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + ) + + conn = op.get_bind() + # Records: one row per schema version present in the data. + conn.execute( + sa.text( + """ + INSERT INTO table_statistics + (schema_id, schema_version, table_name, row_count, records_covered, updated_at) + SELECT schema_id, schema_version, 'records', count(*), NULL, now() + FROM records + GROUP BY schema_id, schema_version + """ + ) + ) + # Features: every registered feature table, scoped per schema via the records join. + tables = conn.execute(sa.text("SELECT hook_name, pg_table FROM feature_tables")).fetchall() + for hook_name, pg_table in tables: + if not _SAFE_IDENT.match(pg_table): + continue + conn.execute( + sa.text( + f""" + INSERT INTO table_statistics + (schema_id, schema_version, table_name, row_count, + records_covered, updated_at) + SELECT r.schema_id, r.schema_version, :hook, count(ft.id), + count(DISTINCT ft.record_srn), now() + FROM features."{pg_table}" ft + JOIN records r ON r.srn = ft.record_srn + GROUP BY r.schema_id, r.schema_version + """ + ), + {"hook": hook_name}, + ) + + +def downgrade() -> None: + op.drop_table("table_statistics") diff --git a/server/migrations/versions/f3a1c9d27e54_records_read_index.py b/server/migrations/versions/f3a1c9d27e54_records_read_index.py new file mode 100644 index 00000000..5ad7c585 --- /dev/null +++ b/server/migrations/versions/f3a1c9d27e54_records_read_index.py @@ -0,0 +1,43 @@ +"""Composite read index on records; drop the subsumed schema_id index. + +``records (schema_id, schema_version, published_at, srn)`` serves the default +table read — schema equality prefix + (published_at, srn) ordering — as one +(backward) index range scan, including the row-value keyset predicate (#219 +phase 3). ``idx_records_schema_id`` is its left prefix and is dropped; +``idx_records_published_at`` stays (used by count_this_month's month filter). + +Built ``CONCURRENTLY``: live archives run this migration on a records table +serving traffic, so the build must not take a table lock. CONCURRENTLY cannot +run inside a transaction block, hence the autocommit block. + +Revision ID: f3a1c9d27e54 +Revises: b47f9c2e8a31 +Create Date: 2026-08-15 +""" + +from alembic import op + +# revision identifiers, used by Alembic. +revision = "f3a1c9d27e54" +down_revision = "b47f9c2e8a31" +branch_labels = None +depends_on = None + +_INDEX = "idx_records_schema_version_published" + + +def upgrade() -> None: + with op.get_context().autocommit_block(): + op.create_index( + _INDEX, + "records", + ["schema_id", "schema_version", "published_at", "srn"], + postgresql_concurrently=True, + ) + op.drop_index("idx_records_schema_id", table_name="records") + + +def downgrade() -> None: + op.create_index("idx_records_schema_id", "records", ["schema_id"]) + with op.get_context().autocommit_block(): + op.drop_index(_INDEX, table_name="records", postgresql_concurrently=True) diff --git a/server/osa/application/api/v1/routes/data/features_table.py b/server/osa/application/api/v1/routes/data/features_table.py index 3d7073b7..7fb2eade 100644 --- a/server/osa/application/api/v1/routes/data/features_table.py +++ b/server/osa/application/api/v1/routes/data/features_table.py @@ -16,7 +16,11 @@ from osa.application.api.v1.routes.data._streaming import build_table_response from osa.application.api.v1.routes.data.formats import DataResponseFormat from osa.application.api.v1.routes.data.tables import format_key, register_table_routes -from osa.domain.data.query.read_table import ReadFeatureTable, ReadFeatureTableHandler +from osa.domain.data.query.read_table import ( + ReadFeatureTable, + ReadFeatureTableHandler, + ReadMode, +) from osa.domain.shared.model.ids import FeatureName @@ -36,6 +40,7 @@ async def endpoint( cursor=cursor, limit=limit, sort=parse_sort(sort), + mode=ReadMode.PAGE if fmt.paginated else ReadMode.STREAM, timeout=fmt.timeout, ) ) @@ -60,6 +65,7 @@ async def endpoint( cursor=body.cursor, limit=body.limit, sort=parse_sort(body.sort), + mode=ReadMode.PAGE if fmt.paginated else ReadMode.STREAM, timeout=fmt.timeout, ) ) diff --git a/server/osa/application/api/v1/routes/data/records_table.py b/server/osa/application/api/v1/routes/data/records_table.py index 73318e74..f5215fe7 100644 --- a/server/osa/application/api/v1/routes/data/records_table.py +++ b/server/osa/application/api/v1/routes/data/records_table.py @@ -19,7 +19,11 @@ from osa.application.api.v1.routes.data._streaming import build_table_response from osa.application.api.v1.routes.data.formats import DataResponseFormat from osa.application.api.v1.routes.data.tables import format_key, register_table_routes -from osa.domain.data.query.read_table import ReadRecordsTable, ReadRecordsTableHandler +from osa.domain.data.query.read_table import ( + ReadMode, + ReadRecordsTable, + ReadRecordsTableHandler, +) def _make_get_endpoint(fmt: DataResponseFormat): @@ -36,6 +40,7 @@ async def endpoint( cursor=cursor, limit=limit, sort=parse_sort(sort), + mode=ReadMode.PAGE if fmt.paginated else ReadMode.STREAM, timeout=fmt.timeout, ) ) @@ -58,6 +63,7 @@ async def endpoint( cursor=body.cursor, limit=body.limit, sort=parse_sort(body.sort), + mode=ReadMode.PAGE if fmt.paginated else ReadMode.STREAM, timeout=fmt.timeout, ) ) diff --git a/server/osa/application/api/v1/routes/stats.py b/server/osa/application/api/v1/routes/stats.py index 7896653c..73a6aca7 100644 --- a/server/osa/application/api/v1/routes/stats.py +++ b/server/osa/application/api/v1/routes/stats.py @@ -6,7 +6,12 @@ from fastapi import APIRouter from pydantic import BaseModel -from osa.domain.record.query.get_stats import GetStats, GetStatsHandler +from osa.domain.data.command.verify_statistics import ( + StatisticsDriftReport, + VerifyTableStatistics, + VerifyTableStatisticsHandler, +) +from osa.domain.data.query.get_stats import GetStats, GetStatsHandler router = APIRouter( prefix="/stats", @@ -36,6 +41,19 @@ class StatsResponse(BaseModel): data_url: str = "/api/v1/data" +@router.post("/verify") +async def verify_statistics( + handler: FromDishka[VerifyTableStatisticsHandler], + repair: bool = False, +) -> StatisticsDriftReport: + """Recompute table_statistics truth; report drift; repair on request. + + ADMIN-gated (handler ``__auth__``). The only sanctioned whole-table + counting after deploy — defence in depth for the lockstep counts (#219). + """ + return await handler.run(VerifyTableStatistics(repair=repair)) + + @router.get("") async def get_stats( handler: FromDishka[GetStatsHandler], diff --git a/server/osa/application/workflow/process_batch.py b/server/osa/application/workflow/process_batch.py index 33be697a..c65d59ed 100644 --- a/server/osa/application/workflow/process_batch.py +++ b/server/osa/application/workflow/process_batch.py @@ -606,7 +606,9 @@ async def _publish( ingest_run_id=event.ingest_run_id, ) - # Checkpoint C: records durable before feature inserts (separate engine + FK). + # Checkpoint C: a redo boundary — records durable so a crash during the + # feature stage replays from here. (Feature DML shares the UoW session + # since #219 phase 4; the old separate-engine FK ordering is gone.) await self.uow.commit() return mapping diff --git a/server/osa/domain/data/command/__init__.py b/server/osa/domain/data/command/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/server/osa/domain/data/command/verify_statistics.py b/server/osa/domain/data/command/verify_statistics.py new file mode 100644 index 00000000..d8cdbd92 --- /dev/null +++ b/server/osa/domain/data/command/verify_statistics.py @@ -0,0 +1,41 @@ +"""Admin verifier for ``table_statistics`` (#219 phase 6). + +The lockstep counts are load-bearing for every discovery surface; this command +is the trust-but-verify safety net that lets exact maintenance replace live +counting. It recomputes truth with the backfill migration's query, reports any +drift, and overwrites only when explicitly asked. Defence in depth — never +load-bearing, never scheduled. +""" + +from osa.domain.auth.model.principal import Principal +from osa.domain.auth.model.role import Role +from osa.domain.data.model.statistics import StatisticsDrift +from osa.domain.data.port.statistics_store import StatisticsStore +from osa.domain.shared.authorization.gate import at_least +from osa.domain.shared.command import Command, CommandHandler, Result + + +class VerifyTableStatistics(Command): + repair: bool = False + + +class StatisticsDriftReport(Result): + drift: list[StatisticsDrift] + repaired: bool + + +class VerifyTableStatisticsHandler(CommandHandler[VerifyTableStatistics, StatisticsDriftReport]): + """Recompute → diff → report; mutate only on ``repair=True``.""" + + __auth__ = at_least(Role.ADMIN) + + principal: Principal + stats_store: StatisticsStore + + async def run(self, cmd: VerifyTableStatistics) -> StatisticsDriftReport: + drift = await self.stats_store.table_statistics_drift() + repaired = False + if cmd.repair and drift: + await self.stats_store.repair_table_statistics() + repaired = True + return StatisticsDriftReport(drift=drift, repaired=repaired) diff --git a/server/osa/domain/data/model/query_plan.py b/server/osa/domain/data/model/query_plan.py index 1b398389..958db145 100644 --- a/server/osa/domain/data/model/query_plan.py +++ b/server/osa/domain/data/model/query_plan.py @@ -18,7 +18,7 @@ from collections.abc import AsyncIterator, Mapping from dataclasses import dataclass from enum import StrEnum -from typing import Any +from typing import Annotated, Any, Literal from pydantic import BaseModel, Field, model_validator @@ -53,7 +53,14 @@ def __str__(self) -> str: return self.value -class PaginationParams(BaseModel): +class BoundedPage(BaseModel): + """A page read: cursor + limit, compiled to ``LIMIT limit+1`` in SQL. + + The only pagination interactive paths can construct — reading without a + bound requires naming :class:`FullStream` explicitly (#219 phase 2). + """ + + mode: Literal["page"] = "page" cursor: PaginationCursor | None = None limit: int = Field(default=50, ge=1) @@ -64,8 +71,8 @@ def clamped( cursor: PaginationCursor | None = None, limit: int, max_limit: int, - ) -> "PaginationParams": - """Build params with ``limit`` clamped into ``[1, max_limit]``. + ) -> "BoundedPage": + """Build a page with ``limit`` clamped into ``[1, max_limit]``. Clamp, don't reject: a consumer asking for "everything" with a big number gets the max page, not a 422. The ceiling is operator @@ -75,6 +82,19 @@ def clamped( return cls(cursor=cursor, limit=max(1, min(limit, max_limit))) +class FullStream(BaseModel): + """An explicitly unbounded read — CSV / gzipped-CSV dumps only. + + Carries no limit and no cursor by construction; the store serves it + through a server-side cursor so memory stays bounded. + """ + + mode: Literal["stream"] = "stream" + + +Pagination = Annotated[BoundedPage | FullStream, Field(discriminator="mode")] + + class Keyset(BaseModel): """The keyset-pagination contract for a plan — the single source of truth for which column is the effective primary sort and which breaks ties. @@ -106,7 +126,7 @@ def cursor_from_row(self, row: Mapping[str, Any]) -> str: TableKind.FEATURE: "id", } -# Default sort keys per table kind (data-model.md §PaginationParams). +# Default sort keys per table kind. _DEFAULT_SORTS: dict[TableKind, list[SortSpec]] = { TableKind.RECORDS: [ SortSpec(column="created_at", direction=SortDirection.DESC), @@ -121,7 +141,7 @@ class QueryPlan(BaseModel): table_kind: TableKind feature_name: FeatureName | None = None filter: FilterExpr | None = None - pagination: PaginationParams = Field(default_factory=PaginationParams) + pagination: Pagination = Field(default_factory=BoundedPage) sort: list[SortSpec] = Field(default_factory=list) @model_validator(mode="after") @@ -145,6 +165,8 @@ async def take_page(self, rows: AsyncIterator[Mapping[str, Any]]) -> PageSlice: the REST paginated-JSON path and the view queries build on this, so the encode side of pagination cannot fork. """ + if not isinstance(self.pagination, BoundedPage): + raise ValueError("take_page requires a BoundedPage plan; dumps never paginate") limit = self.pagination.limit page: list[Mapping[str, Any]] = [] truncated = False diff --git a/server/osa/domain/data/model/statistics.py b/server/osa/domain/data/model/statistics.py new file mode 100644 index 00000000..17b7406c --- /dev/null +++ b/server/osa/domain/data/model/statistics.py @@ -0,0 +1,81 @@ +"""Instance- and table-level statistics models for the data read surface. + +The records/feature distinction is a *type*, not a nullable column: +``RecordsCount`` has no coverage (coverage of records by records is +definitionally the row count), ``FeatureCount`` always has it. The union makes +"records row with a coverage value" unrepresentable. +""" + +from __future__ import annotations + +from datetime import datetime +from typing import Annotated, Literal + +from pydantic import BaseModel, Field + + +class InstanceStats(BaseModel): + """Materialized instance-wide aggregates (storage, feature rows).""" + + storage_bytes: int + feature_rows: int + computed_at: datetime + + +class RecordsCount(BaseModel): + """Counts for a schema version's records table.""" + + kind: Literal["records"] = "records" + row_count: int = 0 + + +class FeatureCount(BaseModel): + """Counts for one feature table scoped to a schema version. + + ``records_covered`` = how many of the schema's records have ≥1 row here. + """ + + kind: Literal["feature"] = "feature" + row_count: int = 0 + records_covered: int = 0 + + +TableCount = Annotated[RecordsCount | FeatureCount, Field(discriminator="kind")] + + +class SchemaTableCounts(BaseModel): + """Lockstep counts for one schema version — the manifest's count source. + + Absent state is zero by construction: a fresh schema yields default + ``RecordsCount()`` / ``FeatureCount()`` instances, never an error. + """ + + records: RecordsCount = Field(default_factory=RecordsCount) + features: dict[str, FeatureCount] = Field(default_factory=dict) + + def feature(self, name: str) -> FeatureCount: + return self.features.get(name, FeatureCount()) + + +class TableCountEntry(BaseModel): + """One table's counts with its full identity — stored or recomputed truth.""" + + schema_id: str + schema_version: str + table_name: str + counts: TableCount + + +class StatisticsDrift(BaseModel): + """A stored count that disagrees with the recomputed truth. + + ``None`` on either side means the row is absent there: ``stored=None`` is a + table the stats missed; ``actual=None`` is an orphan stats row whose data + is gone. + """ + + schema_id: str + schema_version: str + table_name: str + stored: TableCount | None + actual: TableCount | None diff --git a/server/osa/domain/data/port/data_read_store.py b/server/osa/domain/data/port/data_read_store.py index 70524307..7dbea8f8 100644 --- a/server/osa/domain/data/port/data_read_store.py +++ b/server/osa/domain/data/port/data_read_store.py @@ -20,11 +20,11 @@ if TYPE_CHECKING: from osa.domain.data.model.catalog import NodeCatalog - from osa.domain.data.model.manifest import SchemaManifest + from osa.domain.data.model.manifest import ColumnSpec, SchemaManifest from osa.domain.data.model.query_plan import QueryPlan from osa.domain.data.model.record_summary import RecordSummary from osa.domain.data.model.skill import AuthorDocs, SampleValue - from osa.domain.shared.model.ids import RecordId + from osa.domain.shared.model.ids import FeatureName, RecordId from osa.domain.shared.model.srn import SchemaId @@ -53,6 +53,24 @@ async def get_schema_manifest(self, schema_id: "SchemaId") -> "SchemaManifest | """Full manifest for a schema. ``None`` if unknown.""" ... + async def get_record_columns(self, schema_id: "SchemaId") -> "list[ColumnSpec] | None": + """The records table's column schema (implicit + declared fields). + + A catalog lookup only — never touches row data (#219: table resolution + must stay O(1) as tables grow). ``None`` if the schema is unknown. + """ + ... + + async def get_feature_columns( + self, schema_id: "SchemaId", feature_name: "FeatureName" + ) -> "list[ColumnSpec] | None": + """A feature table's column schema (implicit + declared columns). + + Same catalog-only contract as :meth:`get_record_columns`. ``None`` if + the feature is not registered on this schema. + """ + ... + async def get_latest_schema_id(self, schema_short_id: str) -> "SchemaId | None": """Resolve a bare schema id to its latest published version. ``None`` if unknown.""" ... diff --git a/server/osa/domain/data/port/statistics_store.py b/server/osa/domain/data/port/statistics_store.py new file mode 100644 index 00000000..f52fa1a2 --- /dev/null +++ b/server/osa/domain/data/port/statistics_store.py @@ -0,0 +1,57 @@ +"""Port for instance statistics and ``table_statistics`` verification (#219). + +Lives in the data domain: statistics are read-surface content (dashboard, +manifest, SKILL), and the verifier is the read surface's defence in depth. +""" + +from __future__ import annotations + +from typing import Protocol + +from osa.domain.data.model.statistics import InstanceStats, StatisticsDrift + + +class StatisticsStore(Protocol): + """Instance snapshot + lockstep-count reads + the verifier's truth query. + + The snapshot holds what only polling the storage engine can observe + (storage bytes). Row counts come from the lockstep-maintained + ``table_statistics`` — never recounted on request paths. The verifier + methods hold the ONLY sanctioned post-deploy whole-table counting. + """ + + async def count_this_month(self) -> int: + """Records published since the start of the current month. + + Live but bounded: an index-served month window over + ``idx_records_published_at``, never a full-table count. + """ + ... + + async def records_total(self) -> int: + """Total records across schemas — SUM over ``table_statistics``.""" + ... + + async def read_snapshot(self) -> InstanceStats | None: + """The last materialized snapshot, or None if never refreshed.""" + ... + + async def compute_snapshot(self) -> InstanceStats: + """Compute the snapshot: sampled storage bytes + summed lockstep counts.""" + ... + + async def refresh(self) -> None: + """Recompute and upsert the singleton snapshot row.""" + ... + + async def table_statistics_drift(self) -> list[StatisticsDrift]: + """Recompute true counts and diff them against ``table_statistics``. + + The truth query is the backfill migration's, kept runnable — this and + repair are the only sanctioned whole-table counting after deploy. + """ + ... + + async def repair_table_statistics(self) -> None: + """Overwrite ``table_statistics`` with recomputed truth (admin repair).""" + ... diff --git a/server/osa/domain/record/query/get_stats.py b/server/osa/domain/data/query/get_stats.py similarity index 63% rename from server/osa/domain/record/query/get_stats.py rename to server/osa/domain/data/query/get_stats.py index d6ca7445..0e9529ea 100644 --- a/server/osa/domain/record/query/get_stats.py +++ b/server/osa/domain/data/query/get_stats.py @@ -1,9 +1,15 @@ -"""GetStats query handler — public node statistics.""" +"""GetStats query handler — public node statistics (data read surface). + +Relocated from the record domain with #219 phase 6: statistics are read-surface +content, and the records total now comes from the lockstep-maintained +``table_statistics`` (SUM over stored counts) instead of a full-table COUNT on +the request path. ``records_this_month`` stays live — an index-served month +window, the one sanctioned counting statement here. +""" from datetime import datetime -from osa.domain.record.port.statistics_store import StatisticsStore -from osa.domain.record.service.record import RecordService +from osa.domain.data.port.statistics_store import StatisticsStore from osa.domain.shared.authorization.gate import public from osa.domain.shared.query import Query, QueryHandler, Result @@ -21,20 +27,13 @@ class StatsResult(Result): class GetStatsHandler(QueryHandler[GetStats, StatsResult]): - """Node statistics: live counts + the materialized storage/feature snapshot. - - ``records`` and ``records_this_month`` are read live (cheap, indexed) so they - stay fresh; ``storage_bytes`` and ``features_per_record`` come from the - periodically-refreshed snapshot, falling back to a live computation on cold - start (before the first refresh). - """ + """Node statistics: lockstep counts + the materialized storage snapshot.""" __auth__ = public() - record_service: RecordService stats_store: StatisticsStore async def run(self, cmd: GetStats) -> StatsResult: - records = await self.record_service.count() + records = await self.stats_store.records_total() records_this_month = await self.stats_store.count_this_month() snapshot = await self.stats_store.read_snapshot() diff --git a/server/osa/domain/data/query/read_table.py b/server/osa/domain/data/query/read_table.py index 9cb40827..f1b15bc3 100644 --- a/server/osa/domain/data/query/read_table.py +++ b/server/osa/domain/data/query/read_table.py @@ -14,6 +14,7 @@ from collections.abc import AsyncIterator, Mapping from dataclasses import dataclass from datetime import timedelta +from enum import StrEnum from typing import Any from pydantic import Field @@ -22,8 +23,9 @@ from osa.domain.data.model.filter import FilterExpr from osa.domain.data.model.manifest import ColumnSpec from osa.domain.data.model.query_plan import ( + BoundedPage, + FullStream, PaginationCursor, - PaginationParams, QueryPlan, SortSpec, TableKind, @@ -35,12 +37,24 @@ from osa.domain.shared.query import Query, QueryHandler +class ReadMode(StrEnum): + """How much of the table a read may return — chosen by the response format. + + ``PAGE`` is the default everywhere; ``STREAM`` (the whole table) must be + named explicitly and only the CSV/gzip dump formats do (#219 phase 2). + """ + + PAGE = "page" + STREAM = "stream" + + class ReadRecordsTable(Query): schema: str # URL segment: ```` or ``@`` filter: FilterExpr | None = None cursor: str | None = None limit: int = 50 sort: list[SortSpec] = Field(default_factory=list) + mode: ReadMode = ReadMode.PAGE timeout: timedelta | None = None # execution budget chosen by the response format @@ -58,8 +72,10 @@ class TableRead: rows: AsyncIterator[Mapping[str, Any]] -def _pagination(cmd: ReadRecordsTable, config: Config) -> PaginationParams: - return PaginationParams.clamped( +def _pagination(cmd: ReadRecordsTable, config: Config) -> BoundedPage | FullStream: + if cmd.mode == ReadMode.STREAM: + return FullStream() + return BoundedPage.clamped( cursor=PaginationCursor(value=cmd.cursor) if cmd.cursor else None, limit=cmd.limit, max_limit=config.data.max_page_limit, diff --git a/server/osa/domain/data/service/data_catalog.py b/server/osa/domain/data/service/data_catalog.py index 376f3f12..fc82fe1f 100644 --- a/server/osa/domain/data/service/data_catalog.py +++ b/server/osa/domain/data/service/data_catalog.py @@ -82,24 +82,32 @@ async def resolve_table( Unknown schema or table raises ``NotFoundError`` (404 before bytes). """ schema_id = await self.resolve_schema(schema) - manifest = await self.get_schema_manifest(schema_id) - # TableResource.name is a plain str ("records" or a feature-table name). - name = ( - "records" - if table_kind == TableKind.RECORDS - else (feature_name.root if feature_name is not None else None) - ) - resource = next( - (tr for tr in manifest.table_resources if tr.name == name and tr.kind == table_kind), - None, + # Columns come straight from the schema/feature catalogs (#219): table + # resolution must stay O(1) as tables grow, so the manifest — which + # carries per-table row counts — is never built on this path. Matching + # stays by name AND kind by construction: the records lookup never + # consults features, and the feature lookup can never yield records. + record_columns = await self.read_store.get_record_columns(schema_id) + if record_columns is None: + raise NotFoundError( + f"No schema '{schema_id.render()}'. See /api/v1/data for the catalog.", + code="schema_not_found", + ) + if table_kind == TableKind.RECORDS: + return ResolvedTable(schema_id=schema_id, columns=record_columns) + name = feature_name.root if feature_name is not None else None + columns = ( + await self.read_store.get_feature_columns(schema_id, feature_name) + if feature_name is not None + else None ) - if resource is None: + if columns is None: raise NotFoundError( f"No table '{name}' on schema '{schema}'. " f"See /api/v1/data/{schema_id.render()} for its table resources.", code="table_not_found", ) - return ResolvedTable(schema_id=schema_id, columns=resource.columns) + return ResolvedTable(schema_id=schema_id, columns=columns) async def get_record_by_id(self, id: RecordId, version: int | None) -> RecordSummary: record = await self.read_store.get_record_by_id(id, version) diff --git a/server/osa/domain/data/service/data_view.py b/server/osa/domain/data/service/data_view.py index bc18da0a..88ff591a 100644 --- a/server/osa/domain/data/service/data_view.py +++ b/server/osa/domain/data/service/data_view.py @@ -19,7 +19,7 @@ from osa.domain.data.model.manifest import ColumnSpec from osa.domain.data.model.query_plan import ( PaginationCursor, - PaginationParams, + BoundedPage, QueryPlan, SortSpec, TableKind, @@ -75,16 +75,17 @@ async def page( schema, table_kind, feature_name=feature_name ) self._check_required_columns(resolved.columns, require_columns) + page = BoundedPage.clamped( + cursor=PaginationCursor(value=cursor) if cursor else None, + limit=limit, + max_limit=self.config.data.max_page_limit, + ) plan = QueryPlan( schema_id=resolved.schema_id, table_kind=table_kind, feature_name=feature_name, filter=filter, - pagination=PaginationParams.clamped( - cursor=PaginationCursor(value=cursor) if cursor else None, - limit=limit, - max_limit=self.config.data.max_page_limit, - ), + pagination=page, sort=list(sort), ) if table_kind == TableKind.RECORDS: @@ -98,7 +99,7 @@ async def page( table=table, filter=filter, sort=list(sort), - limit=plan.pagination.limit, + limit=page.limit, ), columns=resolved.columns, rows=[self._render_row(row, resolved.columns) for row in slice_.rows], diff --git a/server/osa/domain/record/model/statistics.py b/server/osa/domain/record/model/statistics.py deleted file mode 100644 index 5d111073..00000000 --- a/server/osa/domain/record/model/statistics.py +++ /dev/null @@ -1,20 +0,0 @@ -"""Instance-wide statistics — the materialized snapshot of O(rows) aggregates.""" - -from __future__ import annotations - -from datetime import datetime - -from pydantic import BaseModel - - -class InstanceStats(BaseModel): - """Precomputed instance-wide aggregates. - - Only the expensive-to-compute figures are materialized here (total storage - footprint and total feature-table rows); cheap counts (records, this-month) - are read live. Refreshed periodically by the WorkerPool. - """ - - storage_bytes: int - feature_rows: int - computed_at: datetime diff --git a/server/osa/domain/record/port/repository.py b/server/osa/domain/record/port/repository.py index 43bdb216..5e13dbab 100644 --- a/server/osa/domain/record/port/repository.py +++ b/server/osa/domain/record/port/repository.py @@ -32,6 +32,3 @@ async def srns_for_ingest_batch( that batch's index and are correctly excluded. """ ... - - @abstractmethod - async def count(self) -> int: ... diff --git a/server/osa/domain/record/port/statistics_store.py b/server/osa/domain/record/port/statistics_store.py deleted file mode 100644 index 424a6f9d..00000000 --- a/server/osa/domain/record/port/statistics_store.py +++ /dev/null @@ -1,32 +0,0 @@ -"""Port for reading and refreshing the instance-statistics snapshot.""" - -from __future__ import annotations - -from typing import Protocol - -from osa.domain.record.model.statistics import InstanceStats - - -class StatisticsStore(Protocol): - """Reads the materialized instance-statistics snapshot and refreshes it. - - The snapshot holds the expensive aggregates (storage, feature rows). Cheap - counts (this-month) are exposed here as live queries so the read path stays - fresh without a refresh cycle. - """ - - async def count_this_month(self) -> int: - """Records published since the start of the current month (live).""" - ... - - async def read_snapshot(self) -> InstanceStats | None: - """The last materialized snapshot, or None if never refreshed.""" - ... - - async def compute_snapshot(self) -> InstanceStats: - """Compute the aggregates live (cold-start fallback; O(rows)).""" - ... - - async def refresh(self) -> None: - """Recompute and upsert the singleton snapshot row.""" - ... diff --git a/server/osa/domain/record/service/record.py b/server/osa/domain/record/service/record.py index 8a7d8a37..635f5709 100644 --- a/server/osa/domain/record/service/record.py +++ b/server/osa/domain/record/service/record.py @@ -55,10 +55,6 @@ async def get(self, srn: RecordSRN) -> Record: raise NotFoundError(f"Record not found: {srn}") return record - async def count(self) -> int: - """Total published records on this node.""" - return await self.record_repo.count() - async def srns_for_ingest_batch( self, ingest_run_id: str, batch_index: int ) -> dict[str, RecordSRN]: diff --git a/server/osa/infrastructure/data/postgres_catalog_read_store.py b/server/osa/infrastructure/data/postgres_catalog_read_store.py index 72deffdf..a739ee28 100644 --- a/server/osa/infrastructure/data/postgres_catalog_read_store.py +++ b/server/osa/infrastructure/data/postgres_catalog_read_store.py @@ -9,8 +9,9 @@ from __future__ import annotations import logging +from dataclasses import dataclass -from sqlalchemy import func, select +from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from osa.domain.data.model.catalog import ( @@ -28,6 +29,11 @@ ) from osa.domain.data.model.query_plan import TableKind from osa.domain.data.model.record_summary import RecordSummary +from osa.domain.data.model.statistics import ( + FeatureCount, + RecordsCount, + SchemaTableCounts, +) from osa.domain.data.model.skill import AuthorDocs, SampleValue from osa.domain.semantics.model.value import ( FieldDefinition, @@ -35,7 +41,7 @@ NumberConstraints, TermConstraints, ) -from osa.domain.shared.model.ids import RecordId +from osa.domain.shared.model.ids import FeatureName, RecordId from osa.domain.shared.model.srn import Domain, RecordSRN, SchemaId from osa.infrastructure.data.schema_feature_reader import SchemaFeatureReader from osa.infrastructure.persistence.feature_table import ( @@ -46,6 +52,7 @@ conventions_table, records_table, schemas_table, + table_statistics_table, ) logger = logging.getLogger(__name__) @@ -64,6 +71,14 @@ _ALL_FORMATS = ["", "csv", "csv.gz"] +@dataclass(frozen=True) +class _SchemaSpecs: + """A schema's manifest projections: rich field specs + bare column specs.""" + + fields: list[FieldSpec] + columns: list[ColumnSpec] + + class PostgresCatalogReadStore: def __init__(self, session: AsyncSession, node_domain: Domain) -> None: self.session = session @@ -152,9 +167,35 @@ async def get_schema_manifest(self, schema_id: SchemaId) -> SchemaManifest | Non if row is None: return None + specs = self._field_and_column_specs(row["fields"]) + counts = await self._table_counts(schema_id) + records_resource = TableResource( + name="records", + kind=TableKind.RECORDS, + # Implicit columns (id, srn, schema_id, version, created_at) precede + # the schema's declared metadata fields — this is the CSV header order. + columns=[*IMPLICIT_RECORD_COLUMN_SPECS, *specs.columns], + row_count=counts.records.row_count, + formats=list(_ALL_FORMATS), + ) + feature_resources = await self._feature_resources(schema_id, counts) + return SchemaManifest( + id=schema_id.id.root, + version=schema_id.version.root, + srn=schema_id.to_srn(self.node_domain).render(), + title=row["title"], + fields=specs.fields, + table_resources=[records_resource, *feature_resources], + ) + + @staticmethod + def _field_and_column_specs( + fields_blob: list[dict], + ) -> _SchemaSpecs: + """Map a schema's serialized fields to manifest field/column specs.""" field_specs: list[FieldSpec] = [] column_specs: list[ColumnSpec] = [] - for f in row["fields"]: + for f in fields_blob: # The blob IS a serialized FieldDefinition — validate it back into # the domain model and read typed attributes, never raw dict keys. fd = FieldDefinition.model_validate(f) @@ -179,34 +220,66 @@ async def get_schema_manifest(self, schema_id: SchemaId) -> SchemaManifest | Non ) ) column_specs.append(ColumnSpec(name=fd.name, type=fd.type)) + return _SchemaSpecs(fields=field_specs, columns=column_specs) - record_count = await self._records_count(schema_id) - records_resource = TableResource( - name="records", - kind=TableKind.RECORDS, - # Implicit columns (id, srn, schema_id, version, created_at) precede - # the schema's declared metadata fields — this is the CSV header order. - columns=[*IMPLICIT_RECORD_COLUMN_SPECS, *column_specs], - row_count=record_count, - formats=list(_ALL_FORMATS), + # ------------------------------------------------------------------ # + # Columns-only table resolution (#219 phase 1) + # ------------------------------------------------------------------ # + + async def get_record_columns(self, schema_id: SchemaId) -> list[ColumnSpec] | None: + """Records column schema from the ``schemas`` catalog — no row data touched.""" + stmt = select(schemas_table.c.fields).where( + schemas_table.c.id == schema_id.id.root, + schemas_table.c.version == schema_id.version.root, ) - feature_resources = await self._feature_resources(schema_id) - return SchemaManifest( - id=schema_id.id.root, - version=schema_id.version.root, - srn=schema_id.to_srn(self.node_domain).render(), - title=row["title"], - fields=field_specs, - table_resources=[records_resource, *feature_resources], + result = await self.session.execute(stmt) + row = result.mappings().first() + if row is None: + return None + specs = self._field_and_column_specs(row["fields"]) + return [*IMPLICIT_RECORD_COLUMN_SPECS, *specs.columns] + + async def get_feature_columns( + self, schema_id: SchemaId, feature_name: FeatureName + ) -> list[ColumnSpec] | None: + """Feature column schema from the ``feature_tables`` catalog — no row data.""" + for hook_name, fschema in await self._features.feature_tables(schema_id): + if hook_name == feature_name.root: + return [*IMPLICIT_FEATURE_COLUMN_SPECS, *self._feature_column_specs(fschema)] + return None + + async def _table_counts(self, schema_id: SchemaId) -> SchemaTableCounts: + """Lockstep counts per table for one schema version (#219 phase 6). + + One indexed select over ``table_statistics``; an absent row is zero + (the model defaults). Manifest renders must never recount tables. + """ + stmt = select( + table_statistics_table.c.table_name, + table_statistics_table.c.row_count, + table_statistics_table.c.records_covered, + ).where( + table_statistics_table.c.schema_id == schema_id.id.root, + table_statistics_table.c.schema_version == schema_id.version.root, ) + result = await self.session.execute(stmt) + counts = SchemaTableCounts() + for table_name, row_count, covered in result.all(): + if table_name == "records": + counts.records = RecordsCount(row_count=row_count) + else: + counts.features[table_name] = FeatureCount( + row_count=row_count, records_covered=covered or 0 + ) + return counts - async def _feature_resources(self, schema_id: SchemaId) -> list[TableResource]: + async def _feature_resources( + self, schema_id: SchemaId, counts: SchemaTableCounts + ) -> list[TableResource]: """Build a TableResource for each feature table registered on the schema.""" resources: list[TableResource] = [] for hook_name, fschema in await self._features.feature_tables(schema_id): - ft = build_feature_table(hook_name, fschema) - count = await self._features.count_rows(ft, schema_id) - covered = await self._features.count_covered_records(ft, schema_id) + count = counts.feature(hook_name) resources.append( TableResource( name=hook_name, @@ -214,8 +287,8 @@ async def _feature_resources(self, schema_id: SchemaId) -> list[TableResource]: # Implicit columns (id, record_srn, created_at) precede the # hook's declared data columns — this is the CSV header order. columns=[*IMPLICIT_FEATURE_COLUMN_SPECS, *self._feature_column_specs(fschema)], - row_count=count, - records_covered=covered, + row_count=count.row_count, + records_covered=count.records_covered, formats=list(_ALL_FORMATS), ) ) @@ -305,18 +378,6 @@ async def get_latest_schema_id(self, schema_short_id: str) -> SchemaId | None: latest = max(versions, key=lambda v: tuple(int(p) for p in v.split("-")[0].split("."))) return SchemaId.parse(f"{schema_short_id}@{latest}") - async def _records_count(self, schema_id: SchemaId) -> int: - t = records_table - stmt = ( - select(func.count()) - .select_from(t) - .where( - t.c.schema_id == schema_id.id.root, - t.c.schema_version == schema_id.version.root, - ) - ) - return int((await self.session.execute(stmt)).scalar_one()) - @staticmethod def _feature_column_specs(fschema: FeatureSchema) -> list[ColumnSpec]: """Map a feature table's declared columns to manifest ColumnSpecs.""" diff --git a/server/osa/infrastructure/data/postgres_statistics_store.py b/server/osa/infrastructure/data/postgres_statistics_store.py index 2036a555..c21aa2a0 100644 --- a/server/osa/infrastructure/data/postgres_statistics_store.py +++ b/server/osa/infrastructure/data/postgres_statistics_store.py @@ -1,10 +1,16 @@ -"""Postgres adapter for the instance-statistics snapshot. +"""Postgres adapter for instance statistics + table_statistics verification. -Storage size is summed via ``pg_total_relation_size`` over ``records`` plus every -dynamic ``features.*`` and ``metadata.*`` table (enumerated from their catalogs — -never string-built from user input; ``to_regclass`` yields NULL for a missing -table so a dropped table can't error the sum). Feature-row totals are a genuine -O(rows) scan, which is exactly why the result is materialized. +Storage size is the one fact only observable by polling the storage engine — +``pg_total_relation_size`` over ``records`` plus every dynamic ``features.*`` +and ``metadata.*`` table (enumerated from their catalogs — never string-built +from user input; ``to_regclass`` yields NULL for a missing table so a dropped +table can't error the sum). Everything countable comes from the +lockstep-maintained ``table_statistics`` (#219): the snapshot SUMs stored +counts instead of sweeping tables with COUNT(*). + +The verifier methods (:meth:`table_statistics_drift` / :meth:`repair_table_statistics`) +hold the truth query — the same group-bys the backfill migration ran — and are +the only sanctioned whole-table counting after deploy. """ from __future__ import annotations @@ -16,15 +22,21 @@ from sqlalchemy import func, select, text from sqlalchemy.ext.asyncio import AsyncSession -from osa.domain.record.model.statistics import InstanceStats +from osa.domain.data.model.statistics import ( + FeatureCount, + InstanceStats, + RecordsCount, + StatisticsDrift, + TableCountEntry, +) from osa.infrastructure.persistence.api_naming import ( feature_pg_schema, metadata_pg_schema, ) from osa.infrastructure.persistence.tables import ( - feature_tables_table, instance_statistics_table, records_table, + table_statistics_table, ) # Feature/metadata pg_table names are system-generated and validated on creation @@ -37,6 +49,7 @@ def __init__(self, session: AsyncSession) -> None: self.session = session async def count_this_month(self) -> int: + # Live but bounded: index-served month window (idx_records_published_at). stmt = ( select(func.count()) .select_from(records_table) @@ -44,6 +57,12 @@ async def count_this_month(self) -> int: ) return int((await self.session.execute(stmt)).scalar_one()) + async def records_total(self) -> int: + stmt = select(func.coalesce(func.sum(table_statistics_table.c.row_count), 0)).where( + table_statistics_table.c.table_name == "records" + ) + return int((await self.session.execute(stmt)).scalar_one()) + async def read_snapshot(self) -> InstanceStats | None: row = (await self.session.execute(select(instance_statistics_table))).mappings().first() if row is None: @@ -95,14 +114,139 @@ async def _storage_bytes(self) -> int: return int(result.scalar_one() or 0) async def _feature_rows(self) -> int: - names = ( - (await self.session.execute(select(feature_tables_table.c.pg_table))).scalars().all() + """Total feature rows = SUM over the lockstep counts — no table sweep.""" + stmt = select(func.coalesce(func.sum(table_statistics_table.c.row_count), 0)).where( + table_statistics_table.c.table_name != "records" ) - schema = feature_pg_schema() - total = 0 - for name in names: - if not _SAFE_IDENT.match(name): + return int((await self.session.execute(stmt)).scalar_one()) + + # ------------------------------------------------------------------ # + # Verifier: recompute truth, diff, repair (#219 phase 6) + # ------------------------------------------------------------------ # + + async def _lock_statistics(self) -> None: + """Serialize the verifier against every counted write (PR #220 review). + + Taken BEFORE the truth read, held to commit. Lockstep is what makes one + lock sufficient: every write to a counted table bumps + ``table_statistics`` in its own transaction, so an in-flight writer + blocks here while its data rows are still uncommitted (correctly absent + from our truth) and re-applies its additive delta on the repaired base + after we commit. Without this, a write landing between truth read and + delete+reinsert is clobbered — and being additive, the base stays wrong + forever, not just until the next repair. EXCLUSIVE blocks writers only; + manifest reads proceed. + """ + await self.session.execute(text("LOCK TABLE table_statistics IN EXCLUSIVE MODE")) + + async def table_statistics_drift(self) -> list[StatisticsDrift]: + await self._lock_statistics() + truth = {_key(e): e for e in await self._recompute_truth()} + stored = {_key(e): e for e in await self._read_stored()} + drift: list[StatisticsDrift] = [] + for key in sorted(truth.keys() | stored.keys()): + t, s = truth.get(key), stored.get(key) + if (s.counts if s else None) != (t.counts if t else None): + drift.append( + StatisticsDrift( + schema_id=key[0], + schema_version=key[1], + table_name=key[2], + stored=s.counts if s else None, + actual=t.counts if t else None, + ) + ) + return drift + + async def repair_table_statistics(self) -> None: + """Replace stored counts wholesale with recomputed truth, in one tx.""" + await self._lock_statistics() + truth = await self._recompute_truth() + await self.session.execute(sa.delete(table_statistics_table)) + if truth: + await self.session.execute( + sa.insert(table_statistics_table), + [ + { + "schema_id": e.schema_id, + "schema_version": e.schema_version, + "table_name": e.table_name, + "row_count": e.counts.row_count, + "records_covered": ( + e.counts.records_covered if isinstance(e.counts, FeatureCount) else None + ), + "updated_at": datetime.now(UTC), + } + for e in truth + ], + ) + + async def _read_stored(self) -> list[TableCountEntry]: + result = await self.session.execute(select(table_statistics_table)) + return [ + TableCountEntry( + schema_id=row["schema_id"], + schema_version=row["schema_version"], + table_name=row["table_name"], + counts=( + RecordsCount(row_count=row["row_count"]) + if row["table_name"] == "records" + else FeatureCount( + row_count=row["row_count"], + records_covered=row["records_covered"] or 0, + ) + ), + ) + for row in result.mappings() + ] + + async def _recompute_truth(self) -> list[TableCountEntry]: + """The backfill migration's truth query, kept runnable (#219).""" + entries: list[TableCountEntry] = [] + records = await self.session.execute( + text( + """ + SELECT schema_id, schema_version, count(*) AS n + FROM records GROUP BY schema_id, schema_version + """ + ) + ) + for schema_id, schema_version, n in records.fetchall(): + entries.append( + TableCountEntry( + schema_id=schema_id, + schema_version=schema_version, + table_name="records", + counts=RecordsCount(row_count=n), + ) + ) + tables = await self.session.execute(text("SELECT hook_name, pg_table FROM feature_tables")) + fschema = feature_pg_schema() + for hook_name, pg_table in tables.fetchall(): + if not _SAFE_IDENT.match(pg_table): continue - stmt = text(f'SELECT count(*) FROM "{schema}"."{name}"') - total += int((await self.session.execute(stmt)).scalar_one()) - return total + result = await self.session.execute( + text( + f""" + SELECT r.schema_id, r.schema_version, count(ft.id) AS n, + count(DISTINCT ft.record_srn) AS covered + FROM "{fschema}"."{pg_table}" ft + JOIN records r ON r.srn = ft.record_srn + GROUP BY r.schema_id, r.schema_version + """ + ) + ) + for schema_id, schema_version, n, covered in result.fetchall(): + entries.append( + TableCountEntry( + schema_id=schema_id, + schema_version=schema_version, + table_name=hook_name, + counts=FeatureCount(row_count=n, records_covered=covered), + ) + ) + return entries + + +def _key(e: TableCountEntry) -> tuple[str, str, str]: + return (e.schema_id, e.schema_version, e.table_name) diff --git a/server/osa/infrastructure/data/postgres_table_read_store.py b/server/osa/infrastructure/data/postgres_table_read_store.py index 8fb2dcf0..428600cd 100644 --- a/server/osa/infrastructure/data/postgres_table_read_store.py +++ b/server/osa/infrastructure/data/postgres_table_read_store.py @@ -33,6 +33,7 @@ Predicate, ) from osa.domain.data.model.query_plan import ( + BoundedPage, QueryPlan, SortDirection, TableKind, @@ -147,15 +148,24 @@ async def _stream_records(self, plan: QueryPlan) -> AsyncIterator[Mapping[str, A ) col_names = [c.name for c in metadata_schema.columns] - # ``stream()`` opens a server-side cursor. The try/finally closes it on - # client disconnect (the generator is thrown a CancelledError), returning - # the connection to the pool (research §2). - result = await self.session.stream(stmt) + if isinstance(plan.pagination, BoundedPage): + # LIMIT in the statement lets PG top-N / stop the index scan after + # limit+1 rows instead of sorting the whole joined result; at page + # sizes a server-side cursor is pure overhead (#219 phase 2). + stmt = stmt.limit(plan.pagination.limit + 1) + result = await self.session.execute(stmt) + for row in result.mappings(): + yield self._records_row_to_mapping(row, col_names) + return + # FullStream: ``stream()`` opens a server-side cursor. The try/finally + # closes it on client disconnect (the generator is thrown a + # CancelledError), returning the connection to the pool (research §2). + stream_result = await self.session.stream(stmt) try: - async for row in result.mappings(): + async for row in stream_result.mappings(): yield self._records_row_to_mapping(row, col_names) finally: - await result.close() + await stream_result.close() def _records_row_to_mapping(self, row: RowMapping, col_names: list[str]) -> dict[str, Any]: srn = RecordSRN.parse(row["srn"]) @@ -190,12 +200,26 @@ def _records_sort(self, plan: QueryPlan, metadata_table: Any) -> tuple[list[Any] ) page = KeysetPage( [ - SortKey(sort_expr, descending=is_desc, nulls_last=True), - SortKey(tiebreak_expr, descending=is_desc), + SortKey( + sort_expr, + descending=is_desc, + nulls_last=True, + nullable=self._column_nullable(sort_expr), + ), + SortKey(tiebreak_expr, descending=is_desc, nullable=False), ] ) return page.order_by(), self._cursor_after(plan, page, sort_expr, tiebreak_expr) + @staticmethod + def _column_nullable(expr: sa.ColumnElement[Any]) -> bool: + """A column's DDL nullability, conservatively True for non-Column + expressions. NOT NULL is what licenses plain ASC/DESC emission and the + row-value cursor predicate — both index-servable (#219 phase 3).""" + if isinstance(expr, sa.Column): + return bool(expr.nullable) + return True + def _cursor_after( self, plan: QueryPlan, @@ -209,10 +233,11 @@ def _cursor_after( coerced to their column's Python type; a non-conforming value raises ``ValueError`` → 400 (the cursor is client-supplied input). """ - if plan.pagination.cursor is None: + cursor = plan.pagination.cursor if isinstance(plan.pagination, BoundedPage) else None + if cursor is None: return None try: - decoded = decode_cursor(str(plan.pagination.cursor)) + decoded = decode_cursor(str(cursor)) return page.after( ( self._coerce_cursor_value(decoded["s"], sort_expr), @@ -282,12 +307,18 @@ async def _stream_features(self, plan: QueryPlan) -> AsyncIterator[Mapping[str, .order_by(*order_keys) ) - result = await self.session.stream(stmt) + if isinstance(plan.pagination, BoundedPage): + stmt = stmt.limit(plan.pagination.limit + 1) + result = await self.session.execute(stmt) + for row in result.mappings(): + yield dict(row) + return + stream_result = await self.session.stream(stmt) try: - async for row in result.mappings(): + async for row in stream_result.mappings(): yield dict(row) finally: - await result.close() + await stream_result.close() async def _resolve_feature_table( self, schema_id: SchemaId, feature_name: str @@ -325,8 +356,13 @@ def _features_sort(self, plan: QueryPlan, ft: sa.Table) -> tuple[list[Any], Any ) page = KeysetPage( [ - SortKey(sort_expr, descending=is_desc, nulls_last=True), - SortKey(tiebreak_expr, descending=is_desc), + SortKey( + sort_expr, + descending=is_desc, + nulls_last=True, + nullable=self._column_nullable(sort_expr), + ), + SortKey(tiebreak_expr, descending=is_desc, nullable=False), ] ) return page.order_by(), self._cursor_after(plan, page, sort_expr, tiebreak_expr) diff --git a/server/osa/infrastructure/data/schema_feature_reader.py b/server/osa/infrastructure/data/schema_feature_reader.py index 39bd3665..0c940ae1 100644 --- a/server/osa/infrastructure/data/schema_feature_reader.py +++ b/server/osa/infrastructure/data/schema_feature_reader.py @@ -2,17 +2,17 @@ A ``features.`` table is global (UNIQUE(hook_name)) and shared by every convention that registers the hook name, across schemas. This reader answers -"which feature tables does this schema expose" (through its conventions) and -counts a feature table's rows scoped to one schema's records. Composed by both -``/data/`` read adapters (table streaming + catalog/manifest). +"which feature tables does this schema expose" (through its conventions). +Row/coverage counts come from the lockstep ``table_statistics`` (#219) — this +reader never counts. Composed by both ``/data/`` read adapters (table +streaming + catalog/manifest). """ from __future__ import annotations from typing import Any -import sqlalchemy as sa -from sqlalchemy import and_, func, select +from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from osa.domain.shared.model.srn import SchemaId @@ -42,29 +42,6 @@ async def feature_tables(self, schema_id: SchemaId) -> list[tuple[str, FeatureSc for row in result.mappings() ] - async def count_rows(self, ft: sa.Table, schema_id: SchemaId) -> int: - """Row count of a feature table scoped to the schema's records.""" - stmt = ( - select(func.count()) - .select_from(ft.join(records_table, records_table.c.srn == ft.c.record_srn)) - .where(and_(*self.records_scope(schema_id))) - ) - return int((await self.session.execute(stmt)).scalar_one()) - - async def count_covered_records(self, ft: sa.Table, schema_id: SchemaId) -> int: - """Distinct records with ≥1 row in this feature table (join coverage). - - Feature tables are 1-to-many with records, so ``count_rows`` alone can't - tell a 1-row table that covers 1 record from one that covers many; this - is ``COUNT(DISTINCT record_srn)`` over the same schema-scoped join. - """ - stmt = ( - select(func.count(func.distinct(ft.c.record_srn))) - .select_from(ft.join(records_table, records_table.c.srn == ft.c.record_srn)) - .where(and_(*self.records_scope(schema_id))) - ) - return int((await self.session.execute(stmt)).scalar_one()) - @staticmethod def records_scope(schema_id: SchemaId) -> list[Any]: """Records-join conditions scoping a shared feature table to one schema.""" diff --git a/server/osa/infrastructure/event/worker.py b/server/osa/infrastructure/event/worker.py index ca359d74..b3b1537a 100644 --- a/server/osa/infrastructure/event/worker.py +++ b/server/osa/infrastructure/event/worker.py @@ -722,7 +722,7 @@ async def _run_device_auth_cleanup(self) -> None: async def _run_statistics_refresh(self) -> None: """Periodically refresh the materialized instance-statistics snapshot.""" - from osa.domain.record.port.statistics_store import StatisticsStore + from osa.domain.data.port.statistics_store import StatisticsStore while not self._shutdown: try: diff --git a/server/osa/infrastructure/persistence/di.py b/server/osa/infrastructure/persistence/di.py index dc4f5503..11f6fa37 100644 --- a/server/osa/infrastructure/persistence/di.py +++ b/server/osa/infrastructure/persistence/di.py @@ -15,9 +15,10 @@ from osa.domain.metadata.service.metadata import MetadataService from osa.domain.record.port.feature_reader import FeatureReader from osa.domain.record.port.repository import RecordRepository -from osa.domain.record.port.statistics_store import StatisticsStore +from osa.domain.data.port.statistics_store import StatisticsStore from osa.domain.record.query.get_record import GetRecordHandler -from osa.domain.record.query.get_stats import GetStatsHandler +from osa.domain.data.command.verify_statistics import VerifyTableStatisticsHandler +from osa.domain.data.query.get_stats import GetStatsHandler from osa.domain.record.service import RecordService from osa.infrastructure.persistence.adapter.feature_reader import PostgresFeatureReader from osa.domain.feature.port.storage import FeatureStoragePort @@ -223,3 +224,4 @@ def get_statistics_store(self, session: AsyncSession) -> PostgresStatisticsStore # Record query handlers get_record_handler = provide(GetRecordHandler, scope=Scope.UOW) get_stats_handler = provide(GetStatsHandler, scope=Scope.UOW) + verify_table_statistics_handler = provide(VerifyTableStatisticsHandler, scope=Scope.UOW) diff --git a/server/osa/infrastructure/persistence/feature_store.py b/server/osa/infrastructure/persistence/feature_store.py index 564a17b0..510025a6 100644 --- a/server/osa/infrastructure/persistence/feature_store.py +++ b/server/osa/infrastructure/persistence/feature_store.py @@ -7,17 +7,23 @@ import sqlalchemy as sa from sqlalchemy import select, text +from sqlalchemy.engine import CursorResult from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession from osa.domain.feature.port.feature_store import FeatureStore -from osa.domain.shared.error import ConflictError, ValidationError +from osa.domain.shared.error import ConflictError, NotFoundError, ValidationError from osa.domain.shared.model.hook import ColumnDef from osa.infrastructure.persistence.api_naming import feature_pg_schema, feature_pg_table from osa.infrastructure.persistence.feature_table import ( FeatureSchema, build_feature_table, ) -from osa.infrastructure.persistence.tables import feature_tables_table +from osa.domain.shared.model.srn import SchemaId +from osa.infrastructure.persistence.statistics_upsert import ( + FeatureDelta, + bump_table_statistics, +) +from osa.infrastructure.persistence.tables import feature_tables_table, records_table _PG_IDENTIFIER = re.compile(r"^[a-z][a-z0-9_]{0,62}$") @@ -87,11 +93,19 @@ async def insert_features( Redoing an insert after a partial failure converges instead of duplicating rows (#160): existing rows for ``record_srn`` in this feature table are deleted before the insert, in the same transaction. + + DML runs on the injected session (#219 phase 4) — the caller's unit of + work owns commit/rollback, so records, metadata, and feature rows + written in one stage land or vanish together. The table object comes + from the ``feature_tables`` catalog (as the read path builds it); + runtime reflection was the only reason this ever needed a raw engine + connection. """ if not rows: return 0 _validate_pg_identifier(feature) + table = await self._catalog_table(feature) now = datetime.now(UTC) enriched_rows = [ @@ -104,23 +118,60 @@ async def insert_features( for row in rows ] - # Bulk insert in chunks of 1000 + # Replace-by-record: drop any prior rows for this record so a redo + # after a partial failure converges instead of duplicating (#160). + delete_result = await self._session.execute( + table.delete().where(table.c.record_srn == record_srn) + ) + # DML always yields a CursorResult; the isinstance narrows the union + # session.execute is typed with. max() guards the DBAPI's -1 sentinel. + deleted = max(delete_result.rowcount, 0) if isinstance(delete_result, CursorResult) else 0 + chunk_size = 1000 total = 0 - pg_schema = feature_pg_schema() - pg_table = feature_pg_table(feature) - async with self._engine.begin() as conn: - # Reflect the actual table to get correct column types for casts - metadata = sa.MetaData(schema=pg_schema) - await conn.run_sync(metadata.reflect, only=[pg_table]) - table = metadata.tables[f"{pg_schema}.{pg_table}"] - - # Replace-by-record: drop any prior rows for this record so a redo - # after a partial failure converges instead of duplicating (#160). - await conn.execute(table.delete().where(table.c.record_srn == record_srn)) - - for i in range(0, len(enriched_rows), chunk_size): - chunk = enriched_rows[i : i + chunk_size] - await conn.execute(table.insert(), chunk) - total += len(chunk) + for i in range(0, len(enriched_rows), chunk_size): + chunk = enriched_rows[i : i + chunk_size] + await self._session.execute(table.insert(), chunk) + total += len(chunk) + + # Lockstep statistics (#219 phase 5): replace-by-record yields exact + # in-transaction deltas — rows = inserted − deleted; a record enters + # coverage on its first feature write only. The schema identity comes + # from the record row itself (feature tables are shared across + # schemas), so attribution cannot drift from the data. + schema = await self._record_schema(record_srn) + await bump_table_statistics( + self._session, + schema=schema, + delta=FeatureDelta( + feature=feature, + rows=total - deleted, + covered=1 if deleted == 0 else 0, + ), + ) + await self._session.flush() return total + + async def _record_schema(self, record_srn: str) -> SchemaId: + """The owning record's schema identity (PK lookup, in-transaction).""" + result = await self._session.execute( + select(records_table.c.schema_id, records_table.c.schema_version).where( + records_table.c.srn == record_srn + ) + ) + row = result.first() + if row is None: + raise NotFoundError(f"No record '{record_srn}' to attach feature rows to.") + return SchemaId.parse(f"{row[0]}@{row[1]}") + + async def _catalog_table(self, feature: str) -> sa.Table: + """Build the feature's table object from the ``feature_tables`` catalog.""" + result = await self._session.execute( + select(feature_tables_table.c.feature_schema).where( + feature_tables_table.c.hook_name == feature + ) + ) + row = result.first() + if row is None: + raise NotFoundError(f"No feature table registered for hook '{feature}'.") + return build_feature_table(feature, FeatureSchema.model_validate(row[0])) diff --git a/server/osa/infrastructure/persistence/keyset.py b/server/osa/infrastructure/persistence/keyset.py index dc42392f..1744103a 100644 --- a/server/osa/infrastructure/persistence/keyset.py +++ b/server/osa/infrastructure/persistence/keyset.py @@ -3,7 +3,20 @@ Derives both ORDER BY and WHERE predicate from a single sort specification so that NULL handling is consistent between the two. -Key insight for NULLS LAST ordering: +NULLS emission is nullability-aware (#219 phase 3): a NOT NULL sort column +emits plain ``ASC``/``DESC`` — semantically identical when no NULLs exist, and +textually matchable to a default btree in either scan direction (the planner +matches orderings textually, so an explicit ``NULLS LAST`` on DESC forces a +Sort node no index can absorb). Only genuinely nullable columns carry explicit +``NULLS LAST``/``NULLS FIRST``. + +The cursor predicate follows the same split: when every key is NOT NULL and +same-direction, the row-value form ``(k0, k1) < (:v0, :v1)`` is emitted — PG +collapses it to a single index range scan. The OR-form (which defeats that +collapse) is kept only where row-values are not equivalent: nullable or +mixed-direction sorts. + +Key insight for NULLS LAST ordering on nullable columns: - Non-null cursor value: "strictly after" must include ``OR expr IS NULL`` because NULLs sort after all non-null values. - Null cursor value: only the tiebreaker applies @@ -18,7 +31,7 @@ from dataclasses import dataclass from typing import Any, Sequence -from sqlalchemy import ColumnElement, UnaryExpression, and_, false, or_ +from sqlalchemy import ColumnElement, UnaryExpression, and_, false, or_, tuple_ @dataclass(frozen=True) @@ -28,9 +41,14 @@ class SortKey: expression: ColumnElement[Any] descending: bool = False nulls_last: bool = True + nullable: bool = True def order_clause(self) -> UnaryExpression[Any]: clause = self.expression.desc() if self.descending else self.expression.asc() + if not self.nullable: + # No NULLs can exist; plain ASC/DESC matches a default btree in + # both scan directions where an explicit NULLS clause would not. + return clause return clause.nullslast() if self.nulls_last else clause.nullsfirst() @@ -40,8 +58,8 @@ class KeysetPage: Usage:: page = KeysetPage([ - SortKey(sort_expr, descending=is_desc, nulls_last=True), - SortKey(t.c.id, descending=is_desc), + SortKey(sort_expr, descending=is_desc, nullable=...), + SortKey(t.c.id, descending=is_desc, nullable=False), ]) stmt = stmt.order_by(*page.order_by()) if cursor: @@ -61,6 +79,12 @@ def after(self, cursor_values: tuple[Any, ...]) -> ColumnElement[Any]: f"Cursor length {len(cursor_values)} does not match key length {len(self._keys)}" ) + if self._row_value_eligible(cursor_values): + row = tuple_(*[k.expression for k in self._keys]) + # All keys share one direction: strictly-after is a single + # row-wise comparison, which PG serves as one index range scan. + return row < cursor_values if self._keys[0].descending else row > cursor_values + # Build from right to left: for keys (k0, k1), the predicate is # strictly_after(k0, v0) OR (eq(k0, v0) AND strictly_after(k1, v1)) result: ColumnElement[Any] = false() @@ -79,6 +103,17 @@ def after(self, cursor_values: tuple[Any, ...]) -> ColumnElement[Any]: return result + def _row_value_eligible(self, cursor_values: tuple[Any, ...]) -> bool: + """Row-value comparison is exactly equivalent to the OR-form only when + NULLs are impossible (every key NOT NULL, every cursor value present) + and all keys scan the same direction.""" + directions = {k.descending for k in self._keys} + return ( + len(directions) == 1 + and all(not k.nullable for k in self._keys) + and all(v is not None for v in cursor_values) + ) + def _null_eq(expr: ColumnElement[Any], value: Any) -> ColumnElement[Any]: """``IS NULL`` when value is None, else ``= value``.""" @@ -106,6 +141,10 @@ def _strictly_after(key: SortKey, value: Any) -> ColumnElement[Any] | None: # Cursor is at a non-null value gt = expr < value if key.descending else expr > value + if not key.nullable: + # No NULL region exists; the plain comparison is complete. + return gt + if key.nulls_last: # NULLs come after all non-nulls → include them return or_(gt, expr.is_(None)) diff --git a/server/osa/infrastructure/persistence/repository/record.py b/server/osa/infrastructure/persistence/repository/record.py index b6cb3925..370d39e5 100644 --- a/server/osa/infrastructure/persistence/repository/record.py +++ b/server/osa/infrastructure/persistence/repository/record.py @@ -1,13 +1,19 @@ """PostgreSQL implementation of RecordRepository.""" -from sqlalchemy import Integer, func, select, text +from sqlalchemy import Integer, select, text from sqlalchemy.dialects.postgresql import insert from sqlalchemy.ext.asyncio import AsyncSession +from collections import Counter + from osa.domain.record.model.aggregate import Record from osa.domain.record.port.repository import RecordRepository from osa.domain.shared.model.srn import RecordSRN from osa.infrastructure.persistence.mappers.record import record_to_dict, row_to_record +from osa.infrastructure.persistence.statistics_upsert import ( + RecordsDelta, + bump_table_statistics, +) from osa.infrastructure.persistence.tables import records_table @@ -18,16 +24,25 @@ def __init__(self, session: AsyncSession) -> None: self.session = session async def save(self, record: Record) -> None: - """Persist a record. Records are immutable, so this is insert-only.""" + """Persist a record. Records are immutable, so this is insert-only. + + The records count in ``table_statistics`` is bumped in the same + transaction (#219 phase 5) — counts always equal committed data. + """ record_dict = record_to_dict(record) stmt = insert(records_table).values(**record_dict) await self.session.execute(stmt) + await bump_table_statistics( + self.session, schema=record.schema_id, delta=RecordsDelta(rows=1) + ) await self.session.flush() async def save_many(self, records: list[Record]) -> list[Record]: """Multi-row INSERT with ON CONFLICT DO NOTHING. - Returns the records that were actually inserted (duplicates are skipped). + Returns the records that were actually inserted (duplicates are + skipped). The per-schema statistics delta is exactly the rows this + statement actually inserted — redoing a batch nets zero (#219 phase 5). """ if not records: return [] @@ -44,9 +59,16 @@ async def save_many(self, records: list[Record]) -> list[Record]: .returning(records_table.c.srn) ) result = await self.session.execute(stmt) - await self.session.flush() inserted_srns = {row[0] for row in result.fetchall()} - return [r for r in records if str(r.srn) in inserted_srns] + inserted = [r for r in records if str(r.srn) in inserted_srns] + rows_per_schema = Counter(r.schema_id.render() for r in inserted) + schema_by_key = {r.schema_id.render(): r.schema_id for r in inserted} + for key, rows in rows_per_schema.items(): + await bump_table_statistics( + self.session, schema=schema_by_key[key], delta=RecordsDelta(rows=rows) + ) + await self.session.flush() + return inserted async def get(self, srn: RecordSRN) -> Record | None: """Get a record by SRN.""" @@ -76,9 +98,3 @@ async def srns_for_ingest_batch( ) result = await self.session.execute(stmt) return {upstream_source: RecordSRN.parse(srn) for srn, upstream_source in result.fetchall()} - - async def count(self) -> int: - """Count total records in the database.""" - stmt = select(func.count()).select_from(records_table) - result = await self.session.execute(stmt) - return result.scalar() or 0 diff --git a/server/osa/infrastructure/persistence/statistics_upsert.py b/server/osa/infrastructure/persistence/statistics_upsert.py new file mode 100644 index 00000000..9c5ab4f8 --- /dev/null +++ b/server/osa/infrastructure/persistence/statistics_upsert.py @@ -0,0 +1,98 @@ +"""Lockstep ``table_statistics`` upsert (#219 phase 5). + +Counts are write-model derived state: the delta is computable from the write +itself, so the writing adapter upserts it inside its own transaction and the +displayed counts always equal committed data. Additive deltas compose under +concurrency — the ON CONFLICT row lock serializes increments, so no update is +lost. Call this ONLY from the adapter that performed the counted DML, on the +same session, before its transaction commits. + +The delta is an algebraic type, mirroring the count models: ``RecordsDelta`` +carries no coverage (unrepresentable, not merely unused), ``FeatureDelta`` +always does. These are write-side persistence constructs, so they live here +with the helper rather than in the domain model. + +The truth query (recompute-from-source) lives with the backfill migration and +the admin verifier — reconciliation is mandatory there, never load-bearing +here. +""" + +from __future__ import annotations + +from datetime import UTC, datetime +from typing import Literal + +from pydantic import BaseModel, Field +from sqlalchemy import func +from sqlalchemy.dialects.postgresql import insert +from sqlalchemy.ext.asyncio import AsyncSession + +from osa.domain.shared.model.srn import SchemaId +from osa.infrastructure.persistence.tables import table_statistics_table + + +class RecordsDelta(BaseModel): + """Rows added to a schema version's records table.""" + + kind: Literal["records"] = "records" + rows: int + + @property + def table_name(self) -> str: + return "records" + + +class FeatureDelta(BaseModel): + """Rows and coverage added to one feature table for a schema version. + + ``covered`` is 0 or 1: writes are per-record (replace-by-record), and a + record enters coverage on its first feature write only. The bound is + enforced — a batch-level writer (#218) must widen it deliberately, not + drift into it. + """ + + kind: Literal["feature"] = "feature" + feature: str + rows: int + covered: int = Field(ge=0, le=1) + + @property + def table_name(self) -> str: + return self.feature + + +CountDelta = RecordsDelta | FeatureDelta + + +async def bump_table_statistics( + session: AsyncSession, + *, + schema: SchemaId, + delta: CountDelta, +) -> None: + """Apply *delta* to one table's counts, inside the caller's transaction.""" + if delta.rows == 0 and (isinstance(delta, RecordsDelta) or delta.covered == 0): + return + t = table_statistics_table + now = datetime.now(UTC) + covered = delta.covered if isinstance(delta, FeatureDelta) else None + stmt = insert(t).values( + schema_id=schema.id.root, + schema_version=schema.version.root, + table_name=delta.table_name, + row_count=delta.rows, + records_covered=covered, + updated_at=now, + ) + set_: dict = { + "row_count": t.c.row_count + delta.rows, + "updated_at": now, + } + if isinstance(delta, FeatureDelta): + set_["records_covered"] = func.coalesce(t.c.records_covered, 0) + delta.covered + await session.execute( + stmt.on_conflict_do_update( + index_elements=["schema_id", "schema_version", "table_name"], + set_=set_, + ) + ) diff --git a/server/osa/infrastructure/persistence/tables.py b/server/osa/infrastructure/persistence/tables.py index 7e7a54a2..5a70b468 100644 --- a/server/osa/infrastructure/persistence/tables.py +++ b/server/osa/infrastructure/persistence/tables.py @@ -78,7 +78,16 @@ ) Index("idx_records_convention_id", records_table.c.convention_id) -Index("idx_records_schema_id", records_table.c.schema_id) +# Serves the default table read — schema equality prefix + (published_at, srn) +# ordering — as one (backward) index range scan, including the row-value keyset +# predicate (#219). Subsumes the old idx_records_schema_id (left prefix). +Index( + "idx_records_schema_version_published", + records_table.c.schema_id, + records_table.c.schema_version, + records_table.c.published_at, + records_table.c.srn, +) # Expression must be the raw ``->>`` text accessor (NOT .as_string(), which adds a # redundant CAST) so it matches the bulk-publish ON CONFLICT ((source->>'type'), # (source->>'id')) — Postgres matches ON CONFLICT to a unique index by exact @@ -333,6 +342,29 @@ ) +# ============================================================================ +# TABLE STATISTICS (lockstep row/coverage counts, #219) +# ============================================================================ +# One row per (schema version, table): the records table or a feature table. +# Maintained by additive upsert INSIDE the writing adapter's transaction +# (osa/infrastructure/persistence/statistics_upsert.py), so counts always equal +# committed data. Absent row = zero. records_covered is NULL for the records +# row — coverage ("records with ≥1 feature row") is a feature-table concept; +# for records it is definitionally row_count, and storing a duplicate invites +# drift. Only three writers exist: the lockstep upsert, the rev-B backfill +# migration, and the admin verifier's repair path. +table_statistics_table = Table( + "table_statistics", + metadata, + Column("schema_id", Text, primary_key=True), + Column("schema_version", Text, primary_key=True), + Column("table_name", Text, primary_key=True), + Column("row_count", BigInteger, nullable=False, server_default="0"), + Column("records_covered", BigInteger, nullable=True), + Column("updated_at", DateTime(timezone=True), nullable=False), +) + + # ============================================================================ # INSTANCE STATISTICS (materialized snapshot of O(rows) aggregates) # ============================================================================ diff --git a/server/tests/contract/test_mcp_surface.py b/server/tests/contract/test_mcp_surface.py index 620bf03c..e4255d19 100644 --- a/server/tests/contract/test_mcp_surface.py +++ b/server/tests/contract/test_mcp_surface.py @@ -172,6 +172,26 @@ async def get_node_catalog(self) -> NodeCatalog: async def get_schema_manifest(self, schema_id: SchemaId) -> SchemaManifest | None: return _manifest() if schema_id == SCHEMA_ID else None + async def get_record_columns(self, schema_id: SchemaId) -> list | None: + # Mirrors the adapter contract (#219 phase 1): columns without counts. + if schema_id != SCHEMA_ID: + return None + return next( + tr.columns for tr in _manifest().table_resources if tr.kind == TableKind.RECORDS + ) + + async def get_feature_columns(self, schema_id: SchemaId, feature_name) -> list | None: + if schema_id != SCHEMA_ID: + return None + return next( + ( + tr.columns + for tr in _manifest().table_resources + if tr.kind == TableKind.FEATURE and tr.name == feature_name.root + ), + None, + ) + async def get_latest_schema_id(self, schema_short_id: str) -> SchemaId | None: return SCHEMA_ID if schema_short_id == "sample-data" else None diff --git a/server/tests/integration/conftest.py b/server/tests/integration/conftest.py index 4cd6aa21..b1597441 100644 --- a/server/tests/integration/conftest.py +++ b/server/tests/integration/conftest.py @@ -8,7 +8,7 @@ import pytest import pytest_asyncio -from sqlalchemy import text +from sqlalchemy import event, text from sqlalchemy.ext.asyncio import ( AsyncEngine, async_sessionmaker, @@ -69,6 +69,23 @@ async def seed_record( "published_at": published_at or datetime.now(UTC), }, ) + # Mirror the production writer's lockstep statistics bump (#219): read + # surfaces render counts from table_statistics, so seeded records must + # count exactly like published ones. + await conn.execute( + text( + """ + INSERT INTO table_statistics + (schema_id, schema_version, table_name, row_count, + records_covered, updated_at) + VALUES (:schema_id, :schema_version, 'records', 1, NULL, now()) + ON CONFLICT (schema_id, schema_version, table_name) + DO UPDATE SET row_count = table_statistics.row_count + 1, + updated_at = now() + """ + ), + {"schema_id": schema_id, "schema_version": schema_version}, + ) async def seed_hook_run( @@ -128,6 +145,24 @@ async def pg_engine(): await engine.dispose() +@pytest_asyncio.fixture +async def captured_sql(pg_engine: AsyncEngine): + """Record every SQL statement the test's engine emits (#219). + + Yields a mutable list of statement strings; ``.clear()`` it after the + arrange phase so assertions see only the act phase. Backs the request-path + tripwires: zero ``count(`` statements, ``LIMIT`` present on bounded reads. + """ + statements: list[str] = [] + + def _capture(conn, cursor, statement, parameters, context, executemany) -> None: + statements.append(statement) + + event.listen(pg_engine.sync_engine, "before_cursor_execute", _capture) + yield statements + event.remove(pg_engine.sync_engine, "before_cursor_execute", _capture) + + @pytest_asyncio.fixture async def pg_session(pg_engine: AsyncEngine): """Per-test session with TRUNCATE cleanup.""" @@ -146,7 +181,7 @@ async def pg_session(pg_engine: AsyncEngine): "TRUNCATE TABLE depositions, conventions, schemas, ontologies, " "ontology_terms, events, deliveries, records, validation_runs, " "feature_tables, metadata_tables, hooks, hook_releases, hook_runs, " - "users, identities, refresh_tokens, " + "table_statistics, users, identities, refresh_tokens, " "role_assignments CASCADE" ) ) diff --git a/server/tests/integration/persistence/test_feature_store.py b/server/tests/integration/persistence/test_feature_store.py index a5cbeddd..5b258a61 100644 --- a/server/tests/integration/persistence/test_feature_store.py +++ b/server/tests/integration/persistence/test_feature_store.py @@ -122,6 +122,9 @@ async def test_insert_features(self, pg_engine: AsyncEngine, pg_session: AsyncSe ] count = await store.insert_features("insert_hook", record_srn, rows, run_id) assert count == 3 + # DML rides the caller's unit of work (#219 phase 4): commit before + # verifying through a separate engine connection. + await pg_session.commit() # Verify data is in the table async with pg_engine.begin() as conn: @@ -175,6 +178,7 @@ async def test_insert_twice_replaces_rows_for_the_record( run_id, ) assert second == 2 + await pg_session.commit() async with pg_engine.begin() as conn: result = await conn.execute( @@ -210,6 +214,7 @@ async def test_replace_only_touches_the_target_record( await store.insert_features("replace_scope_hook", srn_b, [{"score": 0.2}], run_id) # Redo record A — B must be untouched. await store.insert_features("replace_scope_hook", srn_a, [{"score": 0.5}], run_id) + await pg_session.commit() async with pg_engine.begin() as conn: result = await conn.execute( @@ -279,3 +284,4 @@ async def test_jsonb_column_for_array_and_object( ] count = await store.insert_features("jsonb_hook", record_srn, rows, run_id) assert count == 1 + await pg_session.commit() diff --git a/server/tests/integration/persistence/test_stage_atomicity_postgres.py b/server/tests/integration/persistence/test_stage_atomicity_postgres.py new file mode 100644 index 00000000..df1ea423 --- /dev/null +++ b/server/tests/integration/persistence/test_stage_atomicity_postgres.py @@ -0,0 +1,99 @@ +"""Feature DML joins the unit of work (#219 phase 4). + +``PostgresFeatureStore.insert_features`` was the only DML in the system running +on a private engine connection — it committed independently of the caller's +session, so a stage could half-land: records rolled back, feature rows durable. +These tests demand the property that could not hold before: feature rows +written in a stage commit and roll back WITH the session. + +``create_table`` (DDL) stays engine-scoped — the sanctioned MetadataStore +split (DDL on engine, DML on session). + +Skips automatically unless OSA_DATABASE__URL points at PostgreSQL. +""" + +from datetime import UTC, datetime + +import pytest +import sqlalchemy as sa +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession + +from osa.domain.shared.model.hook import ColumnDef +from osa.infrastructure.persistence.feature_store import PostgresFeatureStore + +from tests.integration.conftest import seed_hook_run, seed_record + +HOOK = "atomic_features" +SRN = "urn:osa:localhost:rec:atomic1@1" + + +def _columns() -> list[ColumnDef]: + return [ColumnDef(name="score", json_type="number", required=True)] + + +async def _count_rows(engine: AsyncEngine) -> int: + async with engine.connect() as conn: + result = await conn.execute(sa.text(f'SELECT count(*) FROM features."{HOOK}"')) + return int(result.scalar_one()) + + +async def _setup(engine: AsyncEngine, session: AsyncSession) -> str: + run_id = await seed_hook_run(engine, feature_name=HOOK, columns=_columns()) + await PostgresFeatureStore(engine, session).create_table(HOOK, _columns()) + await seed_record( + engine, + srn=SRN, + schema_id="compound", + schema_version="1.0.0", + metadata={}, + published_at=datetime.now(UTC), + ) + return run_id + + +@pytest.mark.asyncio +class TestFeatureDmlJoinsTheUnitOfWork: + async def test_uncommitted_feature_rows_roll_back_with_the_session( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + run_id = await _setup(pg_engine, pg_session) + store = PostgresFeatureStore(pg_engine, pg_session) + + inserted = await store.insert_features(HOOK, SRN, [{"score": 1.0}], run_id) + assert inserted == 1 + await pg_session.rollback() + + assert await _count_rows(pg_engine) == 0, ( + "feature rows survived a session rollback — insert_features is " + "committing outside the unit of work" + ) + + async def test_committed_feature_rows_are_durable( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + run_id = await _setup(pg_engine, pg_session) + store = PostgresFeatureStore(pg_engine, pg_session) + + inserted = await store.insert_features(HOOK, SRN, [{"score": 1.0}, {"score": 2.0}], run_id) + assert inserted == 2 + await pg_session.commit() + assert await _count_rows(pg_engine) == 2 + + async def test_replace_by_record_stays_in_transaction( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + """Redo converges (replace semantics) and the replacement itself is + transactional: a rollback restores the previous rows.""" + run_id = await _setup(pg_engine, pg_session) + store = PostgresFeatureStore(pg_engine, pg_session) + + await store.insert_features(HOOK, SRN, [{"score": 1.0}], run_id) + await pg_session.commit() + + await store.insert_features(HOOK, SRN, [{"score": 9.0}, {"score": 8.0}], run_id) + await pg_session.rollback() + assert await _count_rows(pg_engine) == 1, "rollback must restore the replaced rows" + + await store.insert_features(HOOK, SRN, [{"score": 9.0}, {"score": 8.0}], run_id) + await pg_session.commit() + assert await _count_rows(pg_engine) == 2 diff --git a/server/tests/integration/persistence/test_table_statistics_postgres.py b/server/tests/integration/persistence/test_table_statistics_postgres.py new file mode 100644 index 00000000..138f30b4 --- /dev/null +++ b/server/tests/integration/persistence/test_table_statistics_postgres.py @@ -0,0 +1,160 @@ +"""``table_statistics`` maintained in lockstep with the writes (#219 phase 5). + +Counts are write-model derived state: the writing adapter upserts the delta +inside its own transaction, so displayed counts always equal committed data — +no sweep, no projection lag, no live COUNT(*). The invariants under test: + +- rows and stats commit or roll back together (I1); +- redo converges — replaying a batch nets zero delta (I2); +- deltas are ON CONFLICT-aware: only rows actually inserted count; +- feature coverage gains +1 only on a record's first feature write; +- a fresh table simply has no stats row (absent = zero). + +Skips automatically unless OSA_DATABASE__URL points at PostgreSQL. +""" + +from datetime import UTC, datetime + +import pytest +import sqlalchemy as sa +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession + +from osa.domain.record.model.aggregate import Record +from osa.domain.shared.model.hook import ColumnDef +from osa.domain.shared.model.source import IngestSource +from osa.domain.shared.model.srn import ConventionSlug, RecordSRN, SchemaId +from osa.infrastructure.persistence.feature_store import PostgresFeatureStore +from osa.infrastructure.persistence.repository.record import PostgresRecordRepository +from osa.infrastructure.persistence.tables import table_statistics_table + +from tests.integration.conftest import seed_hook_run + +SCHEMA = SchemaId.parse("compound@1.0.0") +HOOK = "stats_features" + + +def _record(i: int) -> Record: + return Record( + srn=RecordSRN.parse(f"urn:osa:localhost:rec:stat{i}@1"), + source=IngestSource( + id=f"stat-{i}", ingest_run_id="run-1", upstream_source=f"up-{i}", batch_index=0 + ), + convention_id=ConventionSlug.parse("stats-conv"), + schema_id=SCHEMA, + metadata={}, + published_at=datetime(2026, 1, 1, 12, i, tzinfo=UTC), + ) + + +async def _stats(engine: AsyncEngine, table: str) -> tuple[int, int | None] | None: + async with engine.connect() as conn: + result = await conn.execute( + sa.select( + table_statistics_table.c.row_count, + table_statistics_table.c.records_covered, + ).where( + table_statistics_table.c.schema_id == SCHEMA.id.root, + table_statistics_table.c.schema_version == SCHEMA.version.root, + table_statistics_table.c.table_name == table, + ) + ) + row = result.first() + return (row[0], row[1]) if row is not None else None + + +@pytest.mark.asyncio +class TestRecordCountsLockstep: + async def test_save_many_bumps_by_rows_actually_inserted( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + repo = PostgresRecordRepository(pg_session) + inserted = await repo.save_many([_record(1), _record(2), _record(3)]) + assert len(inserted) == 3 + await pg_session.commit() + assert await _stats(pg_engine, "records") == (3, None) + + # Redo the same batch: ON CONFLICT skips all three — delta must be 0, + # not 3 (the delta is rows actually inserted, not rows attempted). + again = await repo.save_many([_record(1), _record(2), _record(3)]) + assert again == [] + await pg_session.commit() + assert await _stats(pg_engine, "records") == (3, None) + + async def test_rows_and_stats_roll_back_together( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + repo = PostgresRecordRepository(pg_session) + await repo.save_many([_record(1)]) + await pg_session.commit() + + await repo.save_many([_record(2), _record(3)]) + await pg_session.rollback() + assert await _stats(pg_engine, "records") == (1, None) + + async def test_single_save_counts_one(self, pg_engine: AsyncEngine, pg_session: AsyncSession): + repo = PostgresRecordRepository(pg_session) + await repo.save(_record(7)) + await pg_session.commit() + assert await _stats(pg_engine, "records") == (1, None) + + async def test_fresh_table_has_no_stats_row( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + assert await _stats(pg_engine, "records") is None + + +@pytest.mark.asyncio +class TestFeatureCountsLockstep: + async def _setup(self, engine: AsyncEngine, session: AsyncSession) -> str: + columns = [ColumnDef(name="score", json_type="number", required=True)] + run_id = await seed_hook_run(engine, feature_name=HOOK, columns=columns) + await PostgresFeatureStore(engine, session).create_table(HOOK, columns) + repo = PostgresRecordRepository(session) + await repo.save_many([_record(1), _record(2)]) + await session.commit() + return run_id + + async def test_deltas_and_coverage(self, pg_engine: AsyncEngine, pg_session: AsyncSession): + run_id = await self._setup(pg_engine, pg_session) + store = PostgresFeatureStore(pg_engine, pg_session) + srn1, srn2 = str(_record(1).srn), str(_record(2).srn) + + await store.insert_features(HOOK, srn1, [{"score": 1.0}, {"score": 2.0}], run_id) + await pg_session.commit() + assert await _stats(pg_engine, HOOK) == (2, 1) + + await store.insert_features( + HOOK, srn2, [{"score": 1.0}, {"score": 2.0}, {"score": 3.0}], run_id + ) + await pg_session.commit() + assert await _stats(pg_engine, HOOK) == (5, 2) + + async def test_redo_converges_to_zero_delta( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + run_id = await self._setup(pg_engine, pg_session) + store = PostgresFeatureStore(pg_engine, pg_session) + srn1 = str(_record(1).srn) + + await store.insert_features(HOOK, srn1, [{"score": 1.0}, {"score": 2.0}], run_id) + await pg_session.commit() + + # Redo, same row count: replace-by-record nets 0 rows, 0 coverage. + await store.insert_features(HOOK, srn1, [{"score": 9.0}, {"score": 8.0}], run_id) + await pg_session.commit() + assert await _stats(pg_engine, HOOK) == (2, 1) + + # Redo with MORE rows: delta = inserted - deleted = +2, coverage still 1. + await store.insert_features(HOOK, srn1, [{"score": float(i)} for i in range(4)], run_id) + await pg_session.commit() + assert await _stats(pg_engine, HOOK) == (4, 1) + + async def test_feature_rows_and_stats_roll_back_together( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + run_id = await self._setup(pg_engine, pg_session) + store = PostgresFeatureStore(pg_engine, pg_session) + + await store.insert_features(HOOK, str(_record(1).srn), [{"score": 1.0}], run_id) + await pg_session.rollback() + assert await _stats(pg_engine, HOOK) is None diff --git a/server/tests/integration/test_bounded_reads_postgres.py b/server/tests/integration/test_bounded_reads_postgres.py new file mode 100644 index 00000000..bbbebeda --- /dev/null +++ b/server/tests/integration/test_bounded_reads_postgres.py @@ -0,0 +1,172 @@ +"""LIMIT pushdown for bounded table reads (#219 phase 2). + +A ``BoundedPage`` read must compile its limit into SQL (``LIMIT limit+1``) so +Postgres can top-N instead of sorting the entire joined result; the previous +plan carried a limit the store ignored, and ``take_page`` discarded the surplus +in Python over a server-side cursor. ``FullStream`` (the dump path) must remain +genuinely unbounded. + +Skips automatically unless OSA_DATABASE__URL points at PostgreSQL. +""" + +from datetime import UTC, datetime + +import pytest +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession + +from osa.domain.data.model.query_plan import ( + BoundedPage, + FullStream, + PaginationCursor, + QueryPlan, + TableKind, +) +from osa.domain.semantics.model.schema import Schema +from osa.domain.semantics.model.value import Cardinality, FieldDefinition, FieldType +from osa.domain.shared.model.srn import RecordSRN, SchemaId +from osa.infrastructure.data.postgres_table_read_store import PostgresTableReadStore +from osa.infrastructure.persistence.metadata_store import PostgresMetadataStore +from osa.infrastructure.persistence.repository.schema import ( + PostgresSemanticsSchemaRepository, +) + +from tests.integration.conftest import seed_record + +SCHEMA = SchemaId.parse("compound@1.0.0") + + +def _fields() -> list[FieldDefinition]: + return [ + FieldDefinition( + name="species", + type=FieldType.TEXT, + required=True, + cardinality=Cardinality.EXACTLY_ONE, + ), + ] + + +async def _setup_schema(engine: AsyncEngine, session: AsyncSession) -> PostgresMetadataStore: + store = PostgresMetadataStore(engine, session) + await store.ensure_table(SCHEMA, _fields()) + await PostgresSemanticsSchemaRepository(session).save( + Schema(id=SCHEMA, title="compound", fields=_fields(), created_at=datetime.now(UTC)) + ) + return store + + +async def _seed(engine: AsyncEngine, store: PostgresMetadataStore, n: int) -> None: + for i in range(n): + srn = RecordSRN.parse(f"urn:osa:localhost:rec:rec{i:03d}@1") + await seed_record( + engine, + srn=str(srn), + schema_id=SCHEMA.id.root, + schema_version=SCHEMA.version.root, + metadata={"species": f"sp{i}"}, + published_at=datetime(2026, 1, 1, i % 24, i % 60, tzinfo=UTC), + ) + await store.insert(SCHEMA, srn, {"species": f"sp{i}"}) + + +def _plan(pagination) -> QueryPlan: + return QueryPlan(schema_id=SCHEMA, table_kind=TableKind.RECORDS, pagination=pagination) + + +async def _drain(store: PostgresTableReadStore, plan: QueryPlan) -> list[dict]: + return [dict(row) async for row in store.stream_rows(plan)] + + +@pytest.mark.asyncio +class TestLimitPushdown: + async def test_bounded_read_compiles_limit_into_sql( + self, pg_engine: AsyncEngine, pg_session: AsyncSession, captured_sql: list[str] + ): + store = await _setup_schema(pg_engine, pg_session) + await _seed(pg_engine, store, 8) + await pg_session.commit() + + captured_sql.clear() + rows = await _drain(PostgresTableReadStore(pg_session), _plan(BoundedPage(limit=3))) + + selects = [s for s in captured_sql if s.lstrip().upper().startswith("SELECT")] + page_selects = [s for s in selects if "records" in s] + assert page_selects, f"no page SELECT captured: {captured_sql}" + assert any("LIMIT" in s.upper() for s in page_selects), ( + f"bounded read did not push LIMIT into SQL: {page_selects}" + ) + # limit+1: exactly one look-ahead row beyond the page, never the table. + assert len(rows) == 4 + + async def test_full_stream_is_unbounded_and_has_no_limit( + self, pg_engine: AsyncEngine, pg_session: AsyncSession, captured_sql: list[str] + ): + store = await _setup_schema(pg_engine, pg_session) + await _seed(pg_engine, store, 8) + await pg_session.commit() + + captured_sql.clear() + rows = await _drain(PostgresTableReadStore(pg_session), _plan(FullStream())) + + assert len(rows) == 8 + page_selects = [ + s for s in captured_sql if s.lstrip().upper().startswith("SELECT") and "records" in s + ] + assert all("LIMIT" not in s.upper() for s in page_selects), ( + f"FullStream must not carry a LIMIT: {page_selects}" + ) + + +@pytest.mark.asyncio +class TestPageBoundaries: + """take_page semantics pinned across the pushdown (was Python-side truncation).""" + + async def test_exact_fill_has_cursor_iff_more_rows( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + store = await _setup_schema(pg_engine, pg_session) + await _seed(pg_engine, store, 5) + await pg_session.commit() + rs = PostgresTableReadStore(pg_session) + + plan = _plan(BoundedPage(limit=3)) + page = await plan.take_page(rs.stream_rows(plan)) + assert len(page.rows) == 3 + assert page.truncated and page.next_cursor is not None + + plan2 = _plan(BoundedPage(limit=2, cursor=PaginationCursor(value=page.next_cursor))) + page2 = await plan2.take_page(rs.stream_rows(plan2)) + assert len(page2.rows) == 2 + assert not page2.truncated and page2.next_cursor is None + + async def test_empty_page_has_no_cursor(self, pg_engine: AsyncEngine, pg_session: AsyncSession): + await _setup_schema(pg_engine, pg_session) + await pg_session.commit() + rs = PostgresTableReadStore(pg_session) + plan = _plan(BoundedPage(limit=3)) + page = await plan.take_page(rs.stream_rows(plan)) + assert page.rows == [] and page.next_cursor is None and not page.truncated + + async def test_pagination_covers_all_rows_without_gaps_or_dupes( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + store = await _setup_schema(pg_engine, pg_session) + await _seed(pg_engine, store, 7) + await pg_session.commit() + rs = PostgresTableReadStore(pg_session) + + seen: list[str] = [] + cursor: str | None = None + for _ in range(10): + plan = _plan( + BoundedPage( + limit=3, + cursor=PaginationCursor(value=cursor) if cursor else None, + ) + ) + page = await plan.take_page(rs.stream_rows(plan)) + seen.extend(r["srn"] for r in page.rows) + if not page.truncated: + break + cursor = page.next_cursor + assert len(seen) == 7 and len(set(seen)) == 7 diff --git a/server/tests/integration/test_cursor_compat_postgres.py b/server/tests/integration/test_cursor_compat_postgres.py new file mode 100644 index 00000000..6935810c --- /dev/null +++ b/server/tests/integration/test_cursor_compat_postgres.py @@ -0,0 +1,119 @@ +"""Cursor wire-compatibility across the #219 predicate change (phase 3). + +Live consumers hold ``next_cursor`` tokens across deploys. The cursor payload +(``{"s": sort_value, "id": tiebreak}``, urlsafe base64) is wire contract; only +predicate *compilation* may change. These tests mint cursors exactly as a +pre-#219 server handed them out — including the datetime-as-ISO-string +rendering of ``encode_cursor`` — and prove pagination resumes correctly +(no gaps, no duplicates) through the current predicate path. + +Skips automatically unless OSA_DATABASE__URL points at PostgreSQL. +""" + +from datetime import UTC, datetime + +import pytest +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession + +from osa.domain.data.model.query_plan import ( + BoundedPage, + PaginationCursor, + QueryPlan, + TableKind, + encode_cursor, +) +from osa.domain.semantics.model.schema import Schema +from osa.domain.semantics.model.value import Cardinality, FieldDefinition, FieldType +from osa.domain.shared.model.srn import RecordSRN, SchemaId +from osa.infrastructure.data.postgres_table_read_store import PostgresTableReadStore +from osa.infrastructure.persistence.metadata_store import PostgresMetadataStore +from osa.infrastructure.persistence.repository.schema import ( + PostgresSemanticsSchemaRepository, +) + +from tests.integration.conftest import seed_record + +SCHEMA = SchemaId.parse("compound@1.0.0") + +# Deterministic publish times: rec0 oldest … rec9 newest. Default records sort +# is published_at DESC, srn DESC, so the first page is rec9, rec8, … +_PUBLISHED = [datetime(2026, 1, 1, 12, i, tzinfo=UTC) for i in range(10)] + + +def _fields() -> list[FieldDefinition]: + return [ + FieldDefinition( + name="species", + type=FieldType.TEXT, + required=True, + cardinality=Cardinality.EXACTLY_ONE, + ), + ] + + +async def _seed(engine: AsyncEngine, session: AsyncSession) -> None: + store = PostgresMetadataStore(engine, session) + await store.ensure_table(SCHEMA, _fields()) + await PostgresSemanticsSchemaRepository(session).save( + Schema(id=SCHEMA, title="compound", fields=_fields(), created_at=datetime.now(UTC)) + ) + for i, published in enumerate(_PUBLISHED): + srn = RecordSRN.parse(f"urn:osa:localhost:rec:rec{i}@1") + await seed_record( + engine, + srn=str(srn), + schema_id=SCHEMA.id.root, + schema_version=SCHEMA.version.root, + metadata={"species": f"sp{i}"}, + published_at=published, + ) + await store.insert(SCHEMA, srn, {"species": f"sp{i}"}) + await session.commit() + + +@pytest.mark.asyncio +class TestPreChangeCursorsStillPaginate: + async def test_cursor_minted_by_a_pre_219_server_resumes_correctly( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + await _seed(pg_engine, pg_session) + rs = PostgresTableReadStore(pg_session) + + # A pre-#219 server encoded the last row of page one — rec7 at + # position 3 of the DESC ordering — with default=str datetime + # rendering. Mint that token byte-for-byte, bypassing today's + # page-taking machinery entirely. + legacy_cursor = encode_cursor(_PUBLISHED[7], "urn:osa:localhost:rec:rec7@1") + + plan = QueryPlan( + schema_id=SCHEMA, + table_kind=TableKind.RECORDS, + pagination=BoundedPage(limit=4, cursor=PaginationCursor(value=legacy_cursor)), + ) + page = await plan.take_page(rs.stream_rows(plan)) + ids = [r["id"] for r in page.rows] + # Strictly after rec7 in DESC order: rec6..rec3, no gap, no repeat. + assert ids == ["rec6", "rec5", "rec4", "rec3"] + assert page.truncated and page.next_cursor is not None + + async def test_full_walk_from_legacy_cursor_covers_the_tail_exactly_once( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + await _seed(pg_engine, pg_session) + rs = PostgresTableReadStore(pg_session) + cursor = encode_cursor(_PUBLISHED[7], "urn:osa:localhost:rec:rec7@1") + + seen: list[str] = [] + for _ in range(10): + plan = QueryPlan( + schema_id=SCHEMA, + table_kind=TableKind.RECORDS, + pagination=BoundedPage(limit=3, cursor=PaginationCursor(value=cursor)), + ) + page = await plan.take_page(rs.stream_rows(plan)) + seen.extend(r["id"] for r in page.rows) + if not page.truncated: + break + assert page.next_cursor is not None + cursor = page.next_cursor + assert seen == ["rec6", "rec5", "rec4", "rec3", "rec2", "rec1", "rec0"] diff --git a/server/tests/integration/test_data_features_postgres.py b/server/tests/integration/test_data_features_postgres.py index 61970a43..3e98e5dc 100644 --- a/server/tests/integration/test_data_features_postgres.py +++ b/server/tests/integration/test_data_features_postgres.py @@ -16,7 +16,7 @@ from osa.domain.data.model.filter import FilterOperator, FeatureFieldRef, Predicate from osa.domain.data.model.query_plan import ( PaginationCursor, - PaginationParams, + BoundedPage, QueryPlan, SortDirection, SortSpec, @@ -138,7 +138,7 @@ def _feature_plan(filter_expr=None, limit=50) -> QueryPlan: table_kind=TableKind.FEATURE, feature_name=HOOK, filter=filter_expr, - pagination=PaginationParams(limit=limit), + pagination=BoundedPage(limit=limit), ) @@ -259,7 +259,7 @@ async def test_created_at_sort_cursor_round_trips( table_kind=TableKind.FEATURE, feature_name=HOOK, sort=sort, - pagination=PaginationParams(cursor=PaginationCursor(value=cursor)), + pagination=BoundedPage(cursor=PaginationCursor(value=cursor)), ), ) assert [r["id"] for r in page2] == [all_rows[2]["id"]] diff --git a/server/tests/integration/test_data_read_store_postgres.py b/server/tests/integration/test_data_read_store_postgres.py index 24bc7060..3b270855 100644 --- a/server/tests/integration/test_data_read_store_postgres.py +++ b/server/tests/integration/test_data_read_store_postgres.py @@ -16,7 +16,7 @@ from osa.domain.data.model.filter import FilterOperator, MetadataFieldRef, Predicate from osa.domain.data.model.query_plan import ( PaginationCursor, - PaginationParams, + BoundedPage, QueryPlan, SortDirection, SortSpec, @@ -92,7 +92,7 @@ def _records_plan(filter_expr=None, limit=50, cursor=None) -> QueryPlan: schema_id=SCHEMA, table_kind=TableKind.RECORDS, filter=filter_expr, - pagination=PaginationParams(limit=limit, cursor=cursor), + pagination=BoundedPage(limit=limit, cursor=cursor), ) @@ -259,7 +259,7 @@ async def test_cursor_advances_without_overlap( plan2 = QueryPlan( schema_id=SCHEMA, table_kind=TableKind.RECORDS, - pagination=PaginationParams(limit=50, cursor=PaginationCursor(value=cursor)), + pagination=BoundedPage(limit=50, cursor=PaginationCursor(value=cursor)), sort=[SortSpec(column="created_at", direction=SortDirection.DESC)], ) page2 = await _drain(rs, plan2) @@ -297,7 +297,7 @@ async def test_id_sort_cursor_round_trips( schema_id=SCHEMA, table_kind=TableKind.RECORDS, sort=sort, - pagination=PaginationParams(cursor=PaginationCursor(value=cursor)), + pagination=BoundedPage(cursor=PaginationCursor(value=cursor)), ), ) assert [r["id"] for r in page2] == ["rec2"] @@ -359,7 +359,7 @@ async def test_date_sort_cursor_round_trips( schema_id=ASSAY_SCHEMA, table_kind=TableKind.RECORDS, sort=sort, - pagination=PaginationParams(cursor=PaginationCursor(value=cursor)), + pagination=BoundedPage(cursor=PaginationCursor(value=cursor)), ), ) assert [r["id"] for r in page2] == ["arec2"] diff --git a/server/tests/integration/test_data_routes_e2e_postgres.py b/server/tests/integration/test_data_routes_e2e_postgres.py index bf314241..1c4ace48 100644 --- a/server/tests/integration/test_data_routes_e2e_postgres.py +++ b/server/tests/integration/test_data_routes_e2e_postgres.py @@ -120,6 +120,9 @@ async def _seed_feature( [{"score": float(i), "label": f"l{i}"} for i in range(n)], run_id, ) + # Feature DML rides the caller's unit of work (#219 phase 4) — commit so + # the app under test (its own sessions) can see the seeded rows. + await session.commit() @pytest.fixture diff --git a/server/tests/integration/test_data_streaming_guarantees_postgres.py b/server/tests/integration/test_data_streaming_guarantees_postgres.py index 52fdd3aa..576dccc2 100644 --- a/server/tests/integration/test_data_streaming_guarantees_postgres.py +++ b/server/tests/integration/test_data_streaming_guarantees_postgres.py @@ -28,7 +28,7 @@ from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession -from osa.domain.data.model.query_plan import PaginationParams, QueryPlan, TableKind +from osa.domain.data.model.query_plan import FullStream, QueryPlan, TableKind from osa.domain.semantics.model.schema import Schema from osa.domain.semantics.model.value import Cardinality, FieldDefinition, FieldType from osa.domain.shared.model.srn import SchemaId @@ -101,11 +101,13 @@ async def _bulk_seed(engine: AsyncEngine, n: int) -> None: ) -def _records_plan(limit: int = 1000) -> QueryPlan: +def _records_plan() -> QueryPlan: + # These are the DUMP-path guarantees: the server-side-cursor behaviours + # only FullStream reads retain after #219's LIMIT pushdown. return QueryPlan( schema_id=SCHEMA, table_kind=TableKind.RECORDS, - pagination=PaginationParams(limit=limit), + pagination=FullStream(), ) diff --git a/server/tests/integration/test_explain_default_sort_postgres.py b/server/tests/integration/test_explain_default_sort_postgres.py new file mode 100644 index 00000000..ee449117 --- /dev/null +++ b/server/tests/integration/test_explain_default_sort_postgres.py @@ -0,0 +1,258 @@ +"""EXPLAIN assertions for the default-sort bounded reads (#219 phase 3). + +The planner matches sort orderings textually: ``DESC NULLS LAST`` matches no +default btree in either scan direction, so the always-emitted NULLS LAST +guaranteed a Sort node — a pipeline breaker that materializes the entire +joined result before the first row. With PG-default NULLS emission on NOT NULL +sort columns, the composite index ``records (schema_id, schema_version, +published_at, srn)``, and row-value cursor predicates, the default-sort page +must plan as a pure index scan: no Sort node, rows examined ≈ limit+1. + +Skips automatically unless OSA_DATABASE__URL points at PostgreSQL. +""" + +from datetime import UTC, datetime +from typing import Any + +import pytest +import sqlalchemy as sa +from sqlalchemy import event +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession + +from osa.domain.data.model.query_plan import BoundedPage, QueryPlan, TableKind +from osa.domain.semantics.model.schema import Schema +from osa.domain.semantics.model.value import Cardinality, FieldDefinition, FieldType +from osa.domain.shared.model.hook import ColumnDef +from osa.domain.shared.model.ids import FeatureName +from osa.domain.shared.model.srn import SchemaId +from osa.infrastructure.data.postgres_table_read_store import PostgresTableReadStore +from osa.infrastructure.persistence.feature_store import PostgresFeatureStore +from osa.infrastructure.persistence.metadata_store import PostgresMetadataStore +from osa.infrastructure.persistence.repository.schema import ( + PostgresSemanticsSchemaRepository, +) +from osa.infrastructure.persistence.tables import conventions_table + +from tests.factories import make_convention_docs_dict +from tests.integration.conftest import seed_hook_run + +SCHEMA = SchemaId.parse("compound@1.0.0") + + +def _fields() -> list[FieldDefinition]: + return [ + FieldDefinition( + name="species", + type=FieldType.TEXT, + required=True, + cardinality=Cardinality.EXACTLY_ONE, + ), + ] + + +async def _setup(engine: AsyncEngine, session: AsyncSession, n: int) -> None: + store = PostgresMetadataStore(engine, session) + await store.ensure_table(SCHEMA, _fields()) + await PostgresSemanticsSchemaRepository(session).save( + Schema(id=SCHEMA, title="compound", fields=_fields(), created_at=datetime.now(UTC)) + ) + await session.commit() + # Bulk set-based seeding: planner assertions need row volume, and per-row + # inserts at this count dominate the test's wall clock. + async with engine.begin() as conn: + await conn.execute( + sa.text( + """ + INSERT INTO records (srn, convention_id, schema_id, schema_version, + source, metadata, published_at) + SELECT 'urn:osa:localhost:rec:rec' || lpad(g::text, 4, '0') || '@1', + 'urn:osa:localhost:conv:test@1.0.0', :sid, :sver, + jsonb_build_object('type', 'seed', 'id', g::text), + jsonb_build_object('species', 'sp' || g), + TIMESTAMPTZ '2026-01-01' + g * INTERVAL '1 minute' + FROM generate_series(0, :n - 1) g + """ + ), + {"sid": SCHEMA.id.root, "sver": SCHEMA.version.root, "n": n}, + ) + await conn.execute( + sa.text( + """ + INSERT INTO metadata.compound_v1 (record_srn, species) + SELECT 'urn:osa:localhost:rec:rec' || lpad(g::text, 4, '0') || '@1', + 'sp' || g + FROM generate_series(0, :n - 1) g + """ + ), + {"n": n}, + ) + + +async def _captured_page_query( + engine: AsyncEngine, session: AsyncSession, plan: QueryPlan +) -> tuple[str, Any]: + """Run the store's read and capture the page SELECT + its bind parameters.""" + captured: list[tuple[str, Any]] = [] + + def _capture(conn, cursor, statement, parameters, context, executemany) -> None: + if statement.lstrip().upper().startswith("SELECT") and "FROM records" in statement: + captured.append((statement, parameters)) + + event.listen(engine.sync_engine, "before_cursor_execute", _capture) + try: + async for _row in PostgresTableReadStore(session).stream_rows(plan): + pass + finally: + event.remove(engine.sync_engine, "before_cursor_execute", _capture) + assert captured, "no page SELECT captured" + return captured[-1] + + +def _node_types(plan_node: dict, acc: list[str]) -> list[str]: + acc.append(plan_node["Node Type"]) + for child in plan_node.get("Plans", []): + _node_types(child, acc) + return acc + + +async def _explain(engine: AsyncEngine, stmt: str, params: Any) -> dict: + async with engine.connect() as conn: + result = await conn.exec_driver_sql("EXPLAIN (ANALYZE, FORMAT JSON) " + stmt, params) + payload = result.scalar_one() + return payload[0]["Plan"] + + +@pytest.mark.asyncio +class TestDefaultSortPlansAsIndexScan: + async def test_records_default_sort_has_no_sort_node( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + await _setup(pg_engine, pg_session, 300) + plan = QueryPlan( + schema_id=SCHEMA, + table_kind=TableKind.RECORDS, + pagination=BoundedPage(limit=10), + ) + stmt, params = await _captured_page_query(pg_engine, pg_session, plan) + top = await _explain(pg_engine, stmt, params) + nodes = _node_types(top, []) + + assert "Sort" not in nodes and "Incremental Sort" not in nodes, ( + f"default-sort bounded read still requires a Sort: {nodes}" + ) + assert any("Index Scan" in n for n in nodes), ( + f"expected an index scan over records, got: {nodes}" + ) + # O(page), not O(table): the whole plan examined ≈ limit+1 rows. + assert top["Actual Rows"] <= 11 + + async def test_cursor_page_also_plans_without_sort( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + """The keyset predicate must collapse to one index range (row-value + form) — the OR-form defeats the scan even with the right index.""" + await _setup(pg_engine, pg_session, 300) + rs = PostgresTableReadStore(pg_session) + first = QueryPlan( + schema_id=SCHEMA, + table_kind=TableKind.RECORDS, + pagination=BoundedPage(limit=10), + ) + page = await first.take_page(rs.stream_rows(first)) + assert page.next_cursor is not None + + from osa.domain.data.model.query_plan import PaginationCursor + + follow = QueryPlan( + schema_id=SCHEMA, + table_kind=TableKind.RECORDS, + pagination=BoundedPage(limit=10, cursor=PaginationCursor(value=page.next_cursor)), + ) + stmt, params = await _captured_page_query(pg_engine, pg_session, follow) + top = await _explain(pg_engine, stmt, params) + nodes = _node_types(top, []) + assert "Sort" not in nodes and "Incremental Sort" not in nodes, ( + f"cursor-follow page still requires a Sort: {nodes}" + ) + assert top["Actual Rows"] <= 11 + + +HOOK = "chem_features" + + +@pytest.mark.asyncio +class TestFeatureDefaultSortNeedsNoNewIndex: + """#219 explicitly does NOT add feature-table indexes — this verifies (not + assumes) that the default feature read (id ASC = the PK) plans without a + Sort using only the indexes the tables already have.""" + + async def test_feature_default_sort_has_no_sort_node( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + await _setup(pg_engine, pg_session, 2000) + await pg_session.execute( + conventions_table.insert().values( + id=f"{SCHEMA.id.root}-{HOOK}", + title="conv", + description="convention", + schema_id=SCHEMA.id.root, + schema_version=SCHEMA.version.root, + file_requirements={}, + hooks=[HOOK], + source=None, + docs=make_convention_docs_dict(), + created_at=datetime.now(UTC), + ) + ) + await pg_session.commit() + columns = [ColumnDef(name="score", json_type="number", required=True)] + run_id = await seed_hook_run(pg_engine, feature_name=HOOK, columns=columns) + await PostgresFeatureStore(pg_engine, pg_session).create_table(HOOK, columns) + async with pg_engine.begin() as conn: + await conn.execute( + sa.text( + f""" + INSERT INTO features."{HOOK}" (record_srn, run_id, score) + SELECT 'urn:osa:localhost:rec:rec' || lpad((g / 2)::text, 4, '0') || '@1', + :run_id, g::float + FROM generate_series(0, 3999) g + """ + ), + {"run_id": run_id}, + ) + # Planner choices at toy scale are costing noise; give it real stats. + await conn.execute(sa.text("ANALYZE records")) + await conn.execute(sa.text(f'ANALYZE features."{HOOK}"')) + + plan = QueryPlan( + schema_id=SCHEMA, + table_kind=TableKind.FEATURE, + feature_name=FeatureName(HOOK), + pagination=BoundedPage(limit=10), + ) + stmt, params = await _captured_feature_query(pg_engine, pg_session, plan) + top = await _explain(pg_engine, stmt, params) + nodes = _node_types(top, []) + assert "Sort" not in nodes and "Incremental Sort" not in nodes, ( + f"feature default sort requires a Sort — a new index may be needed " + f"(record findings on #219 before adding one): {nodes}" + ) + + +async def _captured_feature_query( + engine: AsyncEngine, session: AsyncSession, plan: QueryPlan +) -> tuple[str, Any]: + captured: list[tuple[str, Any]] = [] + + def _capture(conn, cursor, statement, parameters, context, executemany) -> None: + if statement.lstrip().upper().startswith("SELECT") and "features." in statement: + captured.append((statement, parameters)) + + event.listen(engine.sync_engine, "before_cursor_execute", _capture) + try: + async for _row in PostgresTableReadStore(session).stream_rows(plan): + pass + finally: + event.remove(engine.sync_engine, "before_cursor_execute", _capture) + assert captured, "no feature SELECT captured" + return captured[-1] diff --git a/server/tests/integration/test_nullable_sort_postgres.py b/server/tests/integration/test_nullable_sort_postgres.py new file mode 100644 index 00000000..20ffc886 --- /dev/null +++ b/server/tests/integration/test_nullable_sort_postgres.py @@ -0,0 +1,139 @@ +"""Nullable-column sort semantics pinned across #219 phase 3. + +NOT NULL sort columns switch to PG-default NULLS emission (index-servable); +genuinely nullable columns — dynamic metadata fields — MUST keep explicit +``NULLS LAST`` in both directions and the OR-form keyset predicate, because +row-value comparison is not equivalent in the presence of NULLs. Absent +values sort last either way, and pagination crosses the NULL boundary without +gaps or duplicates. + +Skips automatically unless OSA_DATABASE__URL points at PostgreSQL. +""" + +from datetime import UTC, datetime + +import pytest +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession + +from osa.domain.data.model.query_plan import ( + BoundedPage, + PaginationCursor, + QueryPlan, + SortDirection, + SortSpec, + TableKind, +) +from osa.domain.semantics.model.schema import Schema +from osa.domain.semantics.model.value import Cardinality, FieldDefinition, FieldType +from osa.domain.shared.model.srn import RecordSRN, SchemaId +from osa.infrastructure.data.postgres_table_read_store import PostgresTableReadStore +from osa.infrastructure.persistence.metadata_store import PostgresMetadataStore +from osa.infrastructure.persistence.repository.schema import ( + PostgresSemanticsSchemaRepository, +) + +from tests.integration.conftest import seed_record + +SCHEMA = SchemaId.parse("compound@1.0.0") + +# rid → mw (None = absent). Three non-null values and two NULLs. +_ROWS: list[tuple[str, float | None]] = [ + ("rec0", 3.0), + ("rec1", None), + ("rec2", 1.0), + ("rec3", None), + ("rec4", 2.0), +] + + +def _fields() -> list[FieldDefinition]: + return [ + FieldDefinition( + name="species", + type=FieldType.TEXT, + required=True, + cardinality=Cardinality.EXACTLY_ONE, + ), + FieldDefinition( + name="mw", + type=FieldType.NUMBER, + required=False, + cardinality=Cardinality.EXACTLY_ONE, + ), + ] + + +async def _seed(engine: AsyncEngine, session: AsyncSession) -> None: + store = PostgresMetadataStore(engine, session) + await store.ensure_table(SCHEMA, _fields()) + await PostgresSemanticsSchemaRepository(session).save( + Schema(id=SCHEMA, title="compound", fields=_fields(), created_at=datetime.now(UTC)) + ) + for i, (rid, mw) in enumerate(_ROWS): + srn = RecordSRN.parse(f"urn:osa:localhost:rec:{rid}@1") + meta: dict = {"species": f"sp{i}"} + if mw is not None: + meta["mw"] = mw + await seed_record( + engine, + srn=str(srn), + schema_id=SCHEMA.id.root, + schema_version=SCHEMA.version.root, + metadata=meta, + published_at=datetime(2026, 1, 1, 12, i, tzinfo=UTC), + ) + await store.insert(SCHEMA, srn, meta) + await session.commit() + + +def _plan(direction: SortDirection, limit: int = 10, cursor: str | None = None) -> QueryPlan: + return QueryPlan( + schema_id=SCHEMA, + table_kind=TableKind.RECORDS, + pagination=BoundedPage( + limit=limit, cursor=PaginationCursor(value=cursor) if cursor else None + ), + sort=[SortSpec(column="mw", direction=direction)], + ) + + +async def _walk(rs: PostgresTableReadStore, direction: SortDirection, limit: int) -> list[str]: + seen: list[str] = [] + cursor: str | None = None + for _ in range(10): + plan = _plan(direction, limit=limit, cursor=cursor) + page = await plan.take_page(rs.stream_rows(plan)) + seen.extend(r["id"] for r in page.rows) + if not page.truncated: + break + cursor = page.next_cursor + return seen + + +@pytest.mark.asyncio +class TestNullableSortSemantics: + async def test_absent_values_sort_last_ascending( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + await _seed(pg_engine, pg_session) + rows = await _walk(PostgresTableReadStore(pg_session), SortDirection.ASC, limit=10) + # 1.0, 2.0, 3.0 then the two NULLs (tiebreak srn asc within regions). + assert rows[:3] == ["rec2", "rec4", "rec0"] + assert set(rows[3:]) == {"rec1", "rec3"} + + async def test_absent_values_sort_last_descending( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + await _seed(pg_engine, pg_session) + rows = await _walk(PostgresTableReadStore(pg_session), SortDirection.DESC, limit=10) + assert rows[:3] == ["rec0", "rec4", "rec2"] + assert set(rows[3:]) == {"rec1", "rec3"} + + async def test_pagination_crosses_the_null_boundary_without_gaps( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + await _seed(pg_engine, pg_session) + rs = PostgresTableReadStore(pg_session) + for direction in (SortDirection.ASC, SortDirection.DESC): + rows = await _walk(rs, direction, limit=2) + assert len(rows) == 5 and len(set(rows)) == 5, (direction, rows) diff --git a/server/tests/integration/test_resolve_table_perf_postgres.py b/server/tests/integration/test_resolve_table_perf_postgres.py new file mode 100644 index 00000000..94270deb --- /dev/null +++ b/server/tests/integration/test_resolve_table_perf_postgres.py @@ -0,0 +1,218 @@ +"""``resolve_table`` must not build manifests (#219 phase 1). + +Every table read resolves its column schema through +``DataCatalogService.resolve_table``. Before #219 that walked the full +``SchemaManifest`` — three aggregate queries per feature table — and kept only +a column list. These tests hold the tripwire (zero ``count(`` statements) and +pin the resolution semantics the manifest path provided incidentally, so the +rebuild underneath cannot change observable behaviour. + +Skips automatically unless OSA_DATABASE__URL points at PostgreSQL. +""" + +from datetime import UTC, datetime + +import pytest +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession + +from osa.domain.data.model.query_plan import TableKind +from osa.domain.data.service.data_catalog import DataCatalogService +from osa.domain.semantics.model.schema import Schema +from osa.domain.semantics.model.value import Cardinality, FieldDefinition, FieldType +from osa.domain.shared.error import NotFoundError +from osa.domain.shared.model.hook import ColumnDef +from osa.domain.shared.model.ids import FeatureName +from osa.domain.shared.model.srn import Domain, RecordSRN, SchemaId +from osa.infrastructure.data.postgres_catalog_read_store import PostgresCatalogReadStore +from osa.infrastructure.persistence.feature_store import PostgresFeatureStore +from osa.infrastructure.persistence.metadata_store import PostgresMetadataStore +from osa.infrastructure.persistence.repository.schema import ( + PostgresSemanticsSchemaRepository, +) +from osa.infrastructure.persistence.tables import conventions_table + +from tests.factories import make_convention_docs_dict +from tests.integration.conftest import seed_hook_run, seed_record + +SCHEMA = SchemaId.parse("compound@1.0.0") +SCHEMA_V2 = SchemaId.parse("compound@1.2.0") +HOOK = "chem_features" + + +def _fields() -> list[FieldDefinition]: + return [ + FieldDefinition( + name="species", + type=FieldType.TEXT, + required=True, + cardinality=Cardinality.EXACTLY_ONE, + ), + ] + + +def _feature_columns() -> list[ColumnDef]: + return [ + ColumnDef(name="score", json_type="number", required=True), + ColumnDef(name="label", json_type="string", required=False), + ] + + +def _service(session: AsyncSession) -> DataCatalogService: + return DataCatalogService(read_store=PostgresCatalogReadStore(session, Domain("localhost"))) + + +async def _setup_schema( + engine: AsyncEngine, session: AsyncSession, schema: SchemaId = SCHEMA +) -> PostgresMetadataStore: + store = PostgresMetadataStore(engine, session) + await store.ensure_table(schema, _fields()) + await PostgresSemanticsSchemaRepository(session).save( + Schema(id=schema, title=schema.id.root, fields=_fields(), created_at=datetime.now(UTC)) + ) + return store + + +async def _register_hook( + engine: AsyncEngine, session: AsyncSession, schema: SchemaId = SCHEMA +) -> str: + await session.execute( + conventions_table.insert().values( + id=f"{schema.id.root}-{HOOK}", + title="conv", + description="convention", + schema_id=schema.id.root, + schema_version=schema.version.root, + file_requirements={}, + hooks=[HOOK], + source=None, + docs=make_convention_docs_dict(), + created_at=datetime.now(UTC), + ) + ) + await session.commit() + run_id = await seed_hook_run(engine, feature_name=HOOK, columns=_feature_columns()) + await PostgresFeatureStore(engine, session).create_table(HOOK, _feature_columns()) + return run_id + + +async def _seed_rows(engine: AsyncEngine, store: PostgresMetadataStore, n: int = 3) -> None: + for i in range(n): + srn = RecordSRN.parse(f"urn:osa:localhost:rec:rec{i}@1") + await seed_record( + engine, + srn=str(srn), + schema_id=SCHEMA.id.root, + schema_version=SCHEMA.version.root, + metadata={"species": f"sp{i}"}, + published_at=datetime(2026, 1, 1 + i, tzinfo=UTC), + ) + await store.insert(SCHEMA, srn, {"species": f"sp{i}"}) + + +@pytest.mark.asyncio +class TestResolveTableEmitsNoAggregates: + """The #219 tripwire: table resolution is O(1) catalog lookups, no counting.""" + + async def test_records_resolution_emits_zero_count_statements( + self, pg_engine: AsyncEngine, pg_session: AsyncSession, captured_sql: list[str] + ): + store = await _setup_schema(pg_engine, pg_session) + await _register_hook(pg_engine, pg_session) + await _seed_rows(pg_engine, store) + await pg_session.commit() + + captured_sql.clear() + resolved = await _service(pg_session).resolve_table("compound@1.0.0", TableKind.RECORDS) + + counting = [s for s in captured_sql if "count(" in s.lower()] + assert counting == [], f"resolve_table issued aggregate queries: {counting}" + assert resolved.schema_id == SCHEMA + assert [c.name for c in resolved.columns[:2]] == ["id", "srn"] + + async def test_feature_resolution_emits_zero_count_statements( + self, pg_engine: AsyncEngine, pg_session: AsyncSession, captured_sql: list[str] + ): + store = await _setup_schema(pg_engine, pg_session) + await _register_hook(pg_engine, pg_session) + await _seed_rows(pg_engine, store) + await pg_session.commit() + + captured_sql.clear() + resolved = await _service(pg_session).resolve_table( + "compound@1.0.0", TableKind.FEATURE, FeatureName(HOOK) + ) + + counting = [s for s in captured_sql if "count(" in s.lower()] + assert counting == [], f"resolve_table issued aggregate queries: {counting}" + assert resolved.schema_id == SCHEMA + assert "score" in [c.name for c in resolved.columns] + + +@pytest.mark.asyncio +class TestResolveTableContract: + """Semantics the manifest path provided incidentally, pinned through the rebuild.""" + + async def test_bare_id_resolves_latest_schema_version( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + await _setup_schema(pg_engine, pg_session, SCHEMA) + await _setup_schema(pg_engine, pg_session, SCHEMA_V2) + await pg_session.commit() + + resolved = await _service(pg_session).resolve_table("compound", TableKind.RECORDS) + assert resolved.schema_id == SCHEMA_V2 + + async def test_reserved_names_are_404(self, pg_session: AsyncSession): + svc = _service(pg_session) + for reserved in ("records", "datasets"): + with pytest.raises(NotFoundError): + await svc.resolve_table(reserved, TableKind.RECORDS) + + async def test_unknown_schema_is_404(self, pg_session: AsyncSession): + with pytest.raises(NotFoundError): + await _service(pg_session).resolve_table("nope@1.0.0", TableKind.RECORDS) + + async def test_unknown_feature_on_known_schema_is_404( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + await _setup_schema(pg_engine, pg_session) + await pg_session.commit() + with pytest.raises(NotFoundError): + await _service(pg_session).resolve_table( + "compound@1.0.0", TableKind.FEATURE, FeatureName("no_such_hook") + ) + + async def test_feature_matched_by_name_and_kind( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + """A hook can never shadow the records slot, nor vice versa.""" + store = await _setup_schema(pg_engine, pg_session) + await _register_hook(pg_engine, pg_session) + await _seed_rows(pg_engine, store, n=1) + await pg_session.commit() + svc = _service(pg_session) + + records = await svc.resolve_table("compound@1.0.0", TableKind.RECORDS) + feature = await svc.resolve_table("compound@1.0.0", TableKind.FEATURE, FeatureName(HOOK)) + record_cols = [c.name for c in records.columns] + feature_cols = [c.name for c in feature.columns] + assert "species" in record_cols and "score" not in record_cols + assert "score" in feature_cols and "species" not in feature_cols + + async def test_resolved_columns_match_manifest_columns( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + """The rebuilt lookup returns exactly the columns the manifest declares.""" + store = await _setup_schema(pg_engine, pg_session) + await _register_hook(pg_engine, pg_session) + await _seed_rows(pg_engine, store, n=1) + await pg_session.commit() + svc = _service(pg_session) + + manifest = await svc.get_schema_manifest(SCHEMA) + by_name = {(tr.name, tr.kind): tr.columns for tr in manifest.table_resources} + + records = await svc.resolve_table("compound@1.0.0", TableKind.RECORDS) + assert records.columns == by_name[("records", TableKind.RECORDS)] + feature = await svc.resolve_table("compound@1.0.0", TableKind.FEATURE, FeatureName(HOOK)) + assert feature.columns == by_name[(HOOK, TableKind.FEATURE)] diff --git a/server/tests/integration/test_surfaces_use_statistics_postgres.py b/server/tests/integration/test_surfaces_use_statistics_postgres.py new file mode 100644 index 00000000..dd560e26 --- /dev/null +++ b/server/tests/integration/test_surfaces_use_statistics_postgres.py @@ -0,0 +1,283 @@ +"""Read surfaces consume ``table_statistics`` (#219 phase 6). + +The manifest (and through it the catalog, SKILL.md, and the MCP views), the +dashboard stats query, and the instance snapshot all previously recounted +tables per render. They now read the lockstep-maintained counts: zero +``count(`` statements on any request path — the only sanctioned counting is +the one-time backfill migration and the admin verifier. + +Skips automatically unless OSA_DATABASE__URL points at PostgreSQL. +""" + +import asyncio +from datetime import UTC, datetime + +import pytest +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker + +from osa.domain.auth.model.principal import Principal, ProviderIdentity +from osa.domain.auth.model.role import Role +from osa.domain.auth.model.value import UserId +from osa.domain.data.command.verify_statistics import ( + VerifyTableStatistics, + VerifyTableStatisticsHandler, +) +from osa.domain.data.query.get_stats import GetStats, GetStatsHandler +from osa.domain.data.service.data_catalog import DataCatalogService +from osa.domain.record.model.aggregate import Record +from osa.domain.semantics.model.schema import Schema +from osa.domain.semantics.model.value import Cardinality, FieldDefinition, FieldType +from osa.domain.shared.model.hook import ColumnDef +from osa.domain.shared.model.source import IngestSource +from osa.domain.shared.model.srn import ConventionSlug, Domain, RecordSRN, SchemaId +from osa.infrastructure.data.postgres_catalog_read_store import PostgresCatalogReadStore +from osa.infrastructure.data.postgres_statistics_store import PostgresStatisticsStore +from osa.infrastructure.persistence.feature_store import PostgresFeatureStore +from osa.infrastructure.persistence.metadata_store import PostgresMetadataStore +from osa.infrastructure.persistence.repository.record import PostgresRecordRepository +from osa.infrastructure.persistence.repository.schema import ( + PostgresSemanticsSchemaRepository, +) +from osa.infrastructure.persistence.tables import conventions_table, table_statistics_table + +from tests.factories import make_convention_docs_dict +from tests.integration.conftest import seed_hook_run + +SCHEMA = SchemaId.parse("compound@1.0.0") +HOOK = "surface_features" + + +def _fields() -> list[FieldDefinition]: + return [ + FieldDefinition( + name="species", + type=FieldType.TEXT, + required=True, + cardinality=Cardinality.EXACTLY_ONE, + ), + ] + + +def _record(i: int) -> Record: + return Record( + srn=RecordSRN.parse(f"urn:osa:localhost:rec:surf{i}@1"), + source=IngestSource( + id=f"surf-{i}", ingest_run_id="run-1", upstream_source=f"up-{i}", batch_index=0 + ), + convention_id=ConventionSlug.parse("surf-conv"), + schema_id=SCHEMA, + metadata={"species": f"sp{i}"}, + published_at=datetime(2026, 1, 1, 12, i, tzinfo=UTC), + ) + + +async def _seed_all(engine: AsyncEngine, session: AsyncSession) -> None: + """Seed through the LOCKSTEP writers so table_statistics is maintained.""" + store = PostgresMetadataStore(engine, session) + await store.ensure_table(SCHEMA, _fields()) + await PostgresSemanticsSchemaRepository(session).save( + Schema(id=SCHEMA, title="compound", fields=_fields(), created_at=datetime.now(UTC)) + ) + await session.execute( + conventions_table.insert().values( + id=f"{SCHEMA.id.root}-{HOOK}", + title="conv", + description="convention", + schema_id=SCHEMA.id.root, + schema_version=SCHEMA.version.root, + file_requirements={}, + hooks=[HOOK], + source=None, + docs=make_convention_docs_dict(), + created_at=datetime.now(UTC), + ) + ) + await session.commit() + columns = [ColumnDef(name="score", json_type="number", required=True)] + run_id = await seed_hook_run(engine, feature_name=HOOK, columns=columns) + fstore = PostgresFeatureStore(engine, session) + await fstore.create_table(HOOK, columns) + + repo = PostgresRecordRepository(session) + await repo.save_many([_record(1), _record(2), _record(3)]) + for r in [_record(1), _record(2), _record(3)]: + await store.insert(SCHEMA, r.srn, r.metadata) + # Feature rows for two of the three records: 2 + 3 = 5 rows, coverage 2. + await fstore.insert_features( + HOOK, str(_record(1).srn), [{"score": 1.0}, {"score": 2.0}], run_id + ) + await fstore.insert_features( + HOOK, str(_record(2).srn), [{"score": 3.0}, {"score": 4.0}, {"score": 5.0}], run_id + ) + await session.commit() + + +def _catalog_service(session: AsyncSession) -> DataCatalogService: + return DataCatalogService(read_store=PostgresCatalogReadStore(session, Domain("localhost"))) + + +def _counting(captured: list[str]) -> list[str]: + return [s for s in captured if "count(" in s.lower()] + + +@pytest.mark.asyncio +class TestManifestConsumesStatistics: + async def test_manifest_counts_match_truth_with_zero_count_statements( + self, pg_engine: AsyncEngine, pg_session: AsyncSession, captured_sql: list[str] + ): + await _seed_all(pg_engine, pg_session) + + captured_sql.clear() + manifest = await _catalog_service(pg_session).get_schema_manifest(SCHEMA) + + assert _counting(captured_sql) == [], ( + f"manifest render issued aggregate queries: {_counting(captured_sql)}" + ) + by_name = {tr.name: tr for tr in manifest.table_resources} + assert by_name["records"].row_count == 3 + assert by_name[HOOK].row_count == 5 + assert by_name[HOOK].records_covered == 2 + + async def test_fresh_schema_manifest_renders_zeros( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + """Absent stats row means zero — never an error, never a recount.""" + store = PostgresMetadataStore(pg_engine, pg_session) + await store.ensure_table(SCHEMA, _fields()) + await PostgresSemanticsSchemaRepository(pg_session).save( + Schema(id=SCHEMA, title="compound", fields=_fields(), created_at=datetime.now(UTC)) + ) + await pg_session.commit() + + manifest = await _catalog_service(pg_session).get_schema_manifest(SCHEMA) + by_name = {tr.name: tr for tr in manifest.table_resources} + assert by_name["records"].row_count == 0 + + +@pytest.mark.asyncio +class TestStatsQueryConsumesStatistics: + async def test_records_total_comes_from_statistics_not_a_full_count( + self, pg_engine: AsyncEngine, pg_session: AsyncSession, captured_sql: list[str] + ): + await _seed_all(pg_engine, pg_session) + handler = GetStatsHandler(stats_store=PostgresStatisticsStore(pg_session)) + + captured_sql.clear() + result = await handler.run(GetStats()) + + assert result.records == 3 + # count_this_month is the one sanctioned counting statement left on + # this path: an index-served month window, never a full-table count. + for stmt in _counting(captured_sql): + assert "published_at" in stmt, f"full-table count on the stats path: {stmt}" + + async def test_instance_snapshot_sums_statistics_instead_of_sweeping( + self, pg_engine: AsyncEngine, pg_session: AsyncSession, captured_sql: list[str] + ): + await _seed_all(pg_engine, pg_session) + store = PostgresStatisticsStore(pg_session) + + captured_sql.clear() + snapshot = await store.compute_snapshot() + + assert snapshot.feature_rows == 5 + assert _counting(captured_sql) == [], ( + f"snapshot swept tables with count(): {_counting(captured_sql)}" + ) + + +def _admin() -> Principal: + return Principal( + user_id=UserId.generate(), + provider_identity=ProviderIdentity(provider="test", external_id="admin-1"), + roles=frozenset({Role.ADMIN}), + ) + + +@pytest.mark.asyncio +class TestVerifier: + async def test_clean_state_reports_zero_drift_and_writes_nothing( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + await _seed_all(pg_engine, pg_session) + handler = VerifyTableStatisticsHandler( + principal=_admin(), + stats_store=PostgresStatisticsStore(pg_session), + ) + report = await handler.run(VerifyTableStatistics(repair=False)) + assert report.drift == [] + assert report.repaired is False + + async def test_corrupted_count_is_reported_and_repaired_only_on_request( + self, pg_engine: AsyncEngine, pg_session: AsyncSession + ): + await _seed_all(pg_engine, pg_session) + # Corrupt the records count behind the system's back. + await pg_session.execute( + table_statistics_table.update() + .where(table_statistics_table.c.table_name == "records") + .values(row_count=999) + ) + await pg_session.commit() + + handler = VerifyTableStatisticsHandler( + principal=_admin(), + stats_store=PostgresStatisticsStore(pg_session), + ) + report = await handler.run(VerifyTableStatistics(repair=False)) + assert len(report.drift) == 1 + entry = report.drift[0] + assert entry.table_name == "records" + assert entry.stored is not None and entry.stored.row_count == 999 + assert entry.actual is not None and entry.actual.row_count == 3 + assert report.repaired is False + + # Still drifted (report-only made no writes) — now repair. + report2 = await handler.run(VerifyTableStatistics(repair=True)) + assert len(report2.drift) == 1 and report2.repaired is True + await pg_session.commit() + + report3 = await handler.run(VerifyTableStatistics(repair=False)) + assert report3.drift == [] + + async def test_repair_does_not_clobber_concurrent_lockstep_writes( + self, pg_engine: AsyncEngine, pg_session: AsyncSession, monkeypatch + ): + """Greptile P1 on #220: a write committing between repair's truth read + and its delete+reinsert must not be lost — and because increments are + additive, a lost one would skew the base FOREVER, not just until the + next repair. The fix: repair locks ``table_statistics`` before reading + truth; lockstep means every counted write blocks at its stats bump and + re-applies its delta on the repaired base after repair commits. + """ + await _seed_all(pg_engine, pg_session) + store = PostgresStatisticsStore(pg_session) + + concurrent_write_started = asyncio.Event() + original_truth = store._recompute_truth + + async def paused_truth(): + truth = await original_truth() + # Window between truth read and delete+reinsert: let a concurrent + # lockstep writer run. Under the fix it blocks on the table lock; + # under the bug it commits here and repair clobbers its increment. + concurrent_write_started.set() + await asyncio.sleep(0.3) + return truth + + monkeypatch.setattr(store, "_recompute_truth", paused_truth) + + async def concurrent_publish() -> None: + await concurrent_write_started.wait() + factory = async_sessionmaker(pg_engine, expire_on_commit=False) + async with factory() as session: + await PostgresRecordRepository(session).save_many([_record(9)]) + await session.commit() + + writer = asyncio.create_task(concurrent_publish()) + await store.repair_table_statistics() + await pg_session.commit() + await asyncio.wait_for(writer, timeout=10) + + # 3 seeded + 1 concurrent — the concurrent increment must survive. + assert await store.records_total() == 4 diff --git a/server/tests/unit/application/api/v1/routes/data/test_streaming.py b/server/tests/unit/application/api/v1/routes/data/test_streaming.py index 3e7ccd79..cefd794e 100644 --- a/server/tests/unit/application/api/v1/routes/data/test_streaming.py +++ b/server/tests/unit/application/api/v1/routes/data/test_streaming.py @@ -14,6 +14,7 @@ from osa.application.api.v1.routes.data.formats import FORMATS from osa.domain.data.model.manifest import ColumnSpec from osa.domain.data.model.query_plan import ( + BoundedPage, QueryPlan, SortDirection, SortSpec, @@ -47,7 +48,9 @@ async def _body(resp) -> bytes: def _plan(limit: int = 50) -> QueryPlan: - return QueryPlan(schema_id=SCHEMA, table_kind=TableKind.RECORDS, pagination={"limit": limit}) + return QueryPlan( + schema_id=SCHEMA, table_kind=TableKind.RECORDS, pagination=BoundedPage(limit=limit) + ) @pytest.mark.asyncio @@ -124,7 +127,7 @@ async def test_paginated_records_id_sort_encodes_srn() -> None: plan = QueryPlan( schema_id=SCHEMA, table_kind=TableKind.RECORDS, - pagination={"limit": 2}, + pagination=BoundedPage(limit=2), sort=[SortSpec(column="id", direction=SortDirection.ASC)], ) resp = await build_table_response(_aiter(rows), JSON_FMT, COLUMNS, plan) @@ -147,7 +150,7 @@ async def test_paginated_features_id_sort_encodes_row_id() -> None: schema_id=SCHEMA, table_kind=TableKind.FEATURE, feature_name="chem_features", - pagination={"limit": 2}, + pagination=BoundedPage(limit=2), # default FEATURE sort is id asc ) resp = await build_table_response(_aiter(rows), JSON_FMT, COLUMNS, plan) @@ -172,7 +175,7 @@ async def test_paginated_feature_hook_column_named_srn_does_not_hijack_tiebreak( schema_id=SCHEMA, table_kind=TableKind.FEATURE, feature_name="chem_features", - pagination={"limit": 2}, + pagination=BoundedPage(limit=2), # default FEATURE sort is id asc ) resp = await build_table_response(_aiter(rows), JSON_FMT, COLUMNS, plan) diff --git a/server/tests/unit/domain/data/test_data_catalog_service.py b/server/tests/unit/domain/data/test_data_catalog_service.py index 94cf6551..a351edc6 100644 --- a/server/tests/unit/domain/data/test_data_catalog_service.py +++ b/server/tests/unit/domain/data/test_data_catalog_service.py @@ -77,6 +77,28 @@ async def get_node_catalog(self) -> NodeCatalog: async def get_schema_manifest(self, schema_id: SchemaId) -> SchemaManifest | None: return self.manifest + async def get_record_columns(self, schema_id: SchemaId) -> list[ColumnSpec] | None: + # Mirrors the adapter contract: columns from the catalog, no manifest. + if self.manifest is None: + return None + return next( + tr.columns for tr in self.manifest.table_resources if tr.kind == TableKind.RECORDS + ) + + async def get_feature_columns( + self, schema_id: SchemaId, feature_name: FeatureName + ) -> list[ColumnSpec] | None: + if self.manifest is None: + return None + return next( + ( + tr.columns + for tr in self.manifest.table_resources + if tr.kind == TableKind.FEATURE and tr.name == feature_name.root + ), + None, + ) + async def get_latest_schema_id(self, schema_short_id: str) -> SchemaId | None: return SchemaId.parse(f"{schema_short_id}@1.0.0") diff --git a/server/tests/unit/domain/record/test_get_stats_handler.py b/server/tests/unit/domain/data/test_get_stats_handler.py similarity index 67% rename from server/tests/unit/domain/record/test_get_stats_handler.py rename to server/tests/unit/domain/data/test_get_stats_handler.py index 2cc48633..95b39780 100644 --- a/server/tests/unit/domain/record/test_get_stats_handler.py +++ b/server/tests/unit/domain/data/test_get_stats_handler.py @@ -1,7 +1,7 @@ -"""GetStats query handler — the stats route must not touch RecordRepository. +"""GetStats query handler — counts come from the statistics store (#219 ph6). -Layering regression guard: Router → QueryHandler → Service → Repository. The -/stats route previously injected RecordRepository directly. +The records total is the SUM over lockstep table_statistics, never a live +full-table COUNT; the snapshot covers storage/feature aggregates. """ from datetime import UTC, datetime @@ -9,17 +9,15 @@ import pytest -from osa.domain.record.model.statistics import InstanceStats -from osa.domain.record.query.get_stats import GetStats, GetStatsHandler +from osa.domain.data.model.statistics import InstanceStats +from osa.domain.data.query.get_stats import GetStats, GetStatsHandler class TestGetStatsHandler: @pytest.mark.asyncio async def test_returns_record_count_from_service(self): - service = AsyncMock() - service.count.return_value = 42 - stats_store = AsyncMock() + stats_store.records_total.return_value = 42 stats_store.count_this_month.return_value = 5 stats_store.read_snapshot.return_value = InstanceStats( storage_bytes=1024, @@ -27,7 +25,7 @@ async def test_returns_record_count_from_service(self): computed_at=datetime(2026, 7, 27, tzinfo=UTC), ) - handler = GetStatsHandler(record_service=service, stats_store=stats_store) + handler = GetStatsHandler(stats_store=stats_store) result = await handler.run(GetStats()) assert result.records == 42 @@ -35,14 +33,12 @@ async def test_returns_record_count_from_service(self): assert result.storage_bytes == 1024 assert result.features_per_record == 2.0 # 84 feature rows / 42 records assert result.computed_at == datetime(2026, 7, 27, tzinfo=UTC) - service.count.assert_awaited_once() + stats_store.records_total.assert_awaited_once() @pytest.mark.asyncio async def test_falls_back_to_live_compute_before_first_refresh(self): - service = AsyncMock() - service.count.return_value = 10 - stats_store = AsyncMock() + stats_store.records_total.return_value = 10 stats_store.count_this_month.return_value = 0 stats_store.read_snapshot.return_value = None # no snapshot yet stats_store.compute_snapshot.return_value = InstanceStats( @@ -51,7 +47,7 @@ async def test_falls_back_to_live_compute_before_first_refresh(self): computed_at=datetime(2026, 7, 27, tzinfo=UTC), ) - handler = GetStatsHandler(record_service=service, stats_store=stats_store) + handler = GetStatsHandler(stats_store=stats_store) result = await handler.run(GetStats()) assert result.storage_bytes == 2048 diff --git a/server/tests/unit/domain/data/test_query_plan.py b/server/tests/unit/domain/data/test_query_plan.py index 963abf82..5aa8a083 100644 --- a/server/tests/unit/domain/data/test_query_plan.py +++ b/server/tests/unit/domain/data/test_query_plan.py @@ -7,7 +7,7 @@ from osa.domain.data.model.query_plan import ( PaginationCursor, - PaginationParams, + BoundedPage, QueryPlan, SortDirection, SortSpec, @@ -49,23 +49,23 @@ def test_limit_above_max_clamps_to_max() -> None: # Clamp, don't reject: a consumer asking for "everything" with a big # number gets the max page, not a 422. The clamp *policy* lives HERE, on # the canonical model; the *bound* comes from config (DataConfig). - assert PaginationParams.clamped(limit=5000, max_limit=1000).limit == 1000 - assert PaginationParams.clamped(limit=5000, max_limit=200).limit == 200 + assert BoundedPage.clamped(limit=5000, max_limit=1000).limit == 1000 + assert BoundedPage.clamped(limit=5000, max_limit=200).limit == 200 def test_limit_below_one_clamps_to_one() -> None: - assert PaginationParams.clamped(limit=0, max_limit=1000).limit == 1 - assert PaginationParams.clamped(limit=-5, max_limit=1000).limit == 1 + assert BoundedPage.clamped(limit=0, max_limit=1000).limit == 1 + assert BoundedPage.clamped(limit=-5, max_limit=1000).limit == 1 def test_clamped_preserves_cursor() -> None: - p = PaginationParams.clamped(cursor=PaginationCursor(value="CUR"), limit=10, max_limit=1000) + p = BoundedPage.clamped(cursor=PaginationCursor(value="CUR"), limit=10, max_limit=1000) assert str(p.cursor) == "CUR" assert p.limit == 10 def test_limit_default_is_50() -> None: - assert PaginationParams().limit == 50 + assert BoundedPage().limit == 50 def test_cursor_roundtrip() -> None: diff --git a/server/tests/unit/domain/data/test_query_plan_pagination.py b/server/tests/unit/domain/data/test_query_plan_pagination.py new file mode 100644 index 00000000..a11310a0 --- /dev/null +++ b/server/tests/unit/domain/data/test_query_plan_pagination.py @@ -0,0 +1,63 @@ +"""``BoundedPage | FullStream`` — unbounded reads are opt-in by name (#219 phase 2). + +The old ``PaginationParams`` carried a limit on every plan that the store +ignored and the CSV path disregarded — every reader had to know whether the +limit was real. The discriminated union encodes it structurally: interactive +paths can only construct ``BoundedPage`` (LIMIT pushed into SQL); ``FullStream`` +exists solely for the dump serializers. +""" + +import pytest + +from osa.domain.data.model.query_plan import ( + BoundedPage, + FullStream, + PaginationCursor, + QueryPlan, + TableKind, +) +from osa.domain.shared.model.srn import SchemaId + +SCHEMA = SchemaId.parse("compound@1.0.0") + + +def _plan(pagination) -> QueryPlan: + return QueryPlan(schema_id=SCHEMA, table_kind=TableKind.RECORDS, pagination=pagination) + + +class TestPaginationUnion: + def test_default_pagination_is_a_bounded_page(self) -> None: + plan = QueryPlan(schema_id=SCHEMA, table_kind=TableKind.RECORDS) + assert isinstance(plan.pagination, BoundedPage) + assert plan.pagination.limit == 50 + + def test_full_stream_carries_no_limit_or_cursor(self) -> None: + # The type has no such fields at all — "no limit" is unrepresentable + # as a forgotten default. + fields = set(FullStream.model_fields) + assert "limit" not in fields and "cursor" not in fields + + def test_plan_accepts_full_stream(self) -> None: + plan = _plan(FullStream()) + assert isinstance(plan.pagination, FullStream) + + def test_clamped_lives_on_bounded_page(self) -> None: + page = BoundedPage.clamped(limit=5000, max_limit=1000) + assert page.limit == 1000 + page = BoundedPage.clamped(limit=0, max_limit=1000) + assert page.limit == 1 + + def test_clamped_preserves_cursor(self) -> None: + cur = PaginationCursor(value="abc") + page = BoundedPage.clamped(cursor=cur, limit=10, max_limit=1000) + assert page.cursor == cur + + +@pytest.mark.asyncio +class TestTakePageRequiresBoundedPage: + async def test_take_page_on_full_stream_plan_is_a_programming_error(self) -> None: + async def rows(): + yield {"srn": "urn:osa:localhost:rec:r1@1"} + + with pytest.raises(ValueError, match="BoundedPage"): + await _plan(FullStream()).take_page(rows()) diff --git a/server/tests/unit/domain/data/test_take_page.py b/server/tests/unit/domain/data/test_take_page.py index ed1472a4..778dbf86 100644 --- a/server/tests/unit/domain/data/test_take_page.py +++ b/server/tests/unit/domain/data/test_take_page.py @@ -9,7 +9,7 @@ from typing import Any from osa.domain.data.model.query_plan import ( - PaginationParams, + BoundedPage, QueryPlan, TableKind, decode_cursor, @@ -30,7 +30,7 @@ def _plan(limit: int) -> QueryPlan: schema_id=SchemaId.parse("alloy-sample@1.0.0"), table_kind=TableKind.FEATURE, feature_name="tensile_test", - pagination=PaginationParams(limit=limit), + pagination=BoundedPage(limit=limit), ) diff --git a/server/tests/unit/infrastructure/data/test_cursor_validation.py b/server/tests/unit/infrastructure/data/test_cursor_validation.py index 5f2f2cd3..01bb0b25 100644 --- a/server/tests/unit/infrastructure/data/test_cursor_validation.py +++ b/server/tests/unit/infrastructure/data/test_cursor_validation.py @@ -18,7 +18,7 @@ from osa.domain.data.model.query_plan import ( PaginationCursor, - PaginationParams, + BoundedPage, QueryPlan, SortDirection, SortSpec, @@ -48,7 +48,7 @@ def _records_plan(cursor: str) -> QueryPlan: return QueryPlan( schema_id=SCHEMA, table_kind=TableKind.RECORDS, - pagination=PaginationParams(cursor=PaginationCursor(value=cursor)), + pagination=BoundedPage(cursor=PaginationCursor(value=cursor)), ) @@ -57,7 +57,7 @@ def _feature_plan(cursor: str) -> QueryPlan: schema_id=SCHEMA, table_kind=TableKind.FEATURE, feature_name="chem_features", - pagination=PaginationParams(cursor=PaginationCursor(value=cursor)), + pagination=BoundedPage(cursor=PaginationCursor(value=cursor)), sort=[SortSpec(column="id", direction=SortDirection.ASC)], ) diff --git a/server/tests/unit/infrastructure/test_postgres_feature_store.py b/server/tests/unit/infrastructure/test_postgres_feature_store.py index b72bc8fb..90e50a17 100644 --- a/server/tests/unit/infrastructure/test_postgres_feature_store.py +++ b/server/tests/unit/infrastructure/test_postgres_feature_store.py @@ -6,10 +6,10 @@ import pytest import sqlalchemy as sa -from osa.domain.shared.error import ConflictError, ValidationError +from osa.domain.shared.error import ConflictError, NotFoundError, ValidationError from osa.domain.shared.model.hook import ColumnDef from osa.infrastructure.persistence.feature_store import PostgresFeatureStore -from osa.infrastructure.persistence.feature_table import FEATURES_SCHEMA +from osa.infrastructure.persistence.feature_table import FEATURES_SCHEMA, FeatureSchema _RUN_ID = "0190a1b2-c3d4-7e5f-8a9b-0c1d2e3f4a5b" @@ -142,10 +142,39 @@ async def test_create_table_raises_conflict_on_duplicate(self): class TestInsertFeatures: + """insert_features runs DML on the injected session (#219 phase 4), with + the table built from the feature_tables catalog — call order is + catalog SELECT, replace-DELETE, then insert chunks.""" + + @staticmethod + def _mock_session(feature_columns: list[str] | None = None) -> AsyncMock: + session = AsyncMock() + fschema = FeatureSchema( + columns=[ + ColumnDef(name=c, json_type="number", required=False) + for c in (feature_columns or ["score"]) + ] + ) + catalog_result = MagicMock() + catalog_result.first.return_value = (fschema.model_dump(),) + results = [catalog_result] + # Subsequent calls: the replace-DELETE (rowcount consumed for the stats + # delta), insert chunks, the record-schema PK lookup, and the stats + # upsert — one generic result covers them all. + generic = MagicMock() + generic.rowcount = 0 + generic.first.return_value = ("compound", "1.0.0") + + async def _execute(*args, **kwargs): + return results.pop(0) if results else generic + + session.execute = AsyncMock(side_effect=_execute) + return session + @pytest.mark.asyncio async def test_inserts_rows(self): - engine, conn = _mock_engine_with_reflect("pocket_detect", ["score", "pocket_id"]) - store = PostgresFeatureStore(engine=engine, session=AsyncMock()) + session = self._mock_session(["score", "pocket_id"]) + store = PostgresFeatureStore(engine=AsyncMock(), session=session) rows = [ {"score": 0.95, "pocket_id": "P1"}, {"score": 0.82, "pocket_id": "P2"}, @@ -154,41 +183,41 @@ async def test_inserts_rows(self): count = await store.insert_features("pocket_detect", "urn:rec:1", rows, _RUN_ID) assert count == 2 - # Replace semantics (#160): one DELETE (scoped to the record) + one insert chunk. - assert conn.execute.call_count == 2 + # Catalog SELECT + replace-DELETE (#160) + one insert chunk + # + record-schema lookup + lockstep stats upsert (#219 ph5). + assert session.execute.call_count == 5 @pytest.mark.asyncio async def test_deletes_existing_rows_for_record_before_insert(self): """Replace-by-record: a DELETE scoped to record_srn precedes the insert (#160).""" - engine, conn = _mock_engine_with_reflect("pocket_detect", ["score"]) - store = PostgresFeatureStore(engine=engine, session=AsyncMock()) + session = self._mock_session() + store = PostgresFeatureStore(engine=AsyncMock(), session=session) await store.insert_features("pocket_detect", "urn:rec:1", [{"score": 0.95}], _RUN_ID) - # First execute is the DELETE; second is the insert. - delete_stmt = conn.execute.call_args_list[0][0][0] - compiled = delete_stmt.compile() - assert "DELETE FROM" in str(compiled) - assert "record_srn" in str(compiled) + delete_stmt = session.execute.call_args_list[1][0][0] + compiled = str(delete_stmt.compile()) + assert "DELETE FROM" in compiled + assert "record_srn" in compiled @pytest.mark.asyncio async def test_empty_rows_returns_zero(self): - engine = AsyncMock() - store = PostgresFeatureStore(engine=engine, session=AsyncMock()) + session = AsyncMock() + store = PostgresFeatureStore(engine=AsyncMock(), session=session) count = await store.insert_features("pocket_detect", "urn:rec:1", [], _RUN_ID) assert count == 0 + session.execute.assert_not_called() @pytest.mark.asyncio async def test_enriches_rows_with_record_srn(self): - engine, conn = _mock_engine_with_reflect("pocket_detect", ["score"]) - store = PostgresFeatureStore(engine=engine, session=AsyncMock()) + session = self._mock_session() + store = PostgresFeatureStore(engine=AsyncMock(), session=session) await store.insert_features("pocket_detect", "urn:rec:1", [{"score": 0.95}], _RUN_ID) - call_args = conn.execute.call_args - params = call_args[0][1] # second positional arg is the params list + params = session.execute.call_args_list[2][0][1] assert len(params) == 1 assert params[0]["record_srn"] == "urn:rec:1" assert params[0]["run_id"] == _RUN_ID @@ -197,32 +226,31 @@ async def test_enriches_rows_with_record_srn(self): @pytest.mark.asyncio async def test_chunks_large_inserts(self): - engine, conn = _mock_engine_with_reflect("hook", ["score"]) - store = PostgresFeatureStore(engine=engine, session=AsyncMock()) + session = self._mock_session() + store = PostgresFeatureStore(engine=AsyncMock(), session=session) rows = [{"score": float(i)} for i in range(2500)] count = await store.insert_features("hook", "urn:rec:1", rows, _RUN_ID) assert count == 2500 - # 1 replace-DELETE + 3 insert chunks (1000 + 1000 + 500). - assert conn.execute.call_count == 4 + # Catalog SELECT + replace-DELETE + 3 insert chunks (1000 + 1000 + 500) + # + record-schema lookup + stats upsert. + assert session.execute.call_count == 7 @pytest.mark.asyncio async def test_single_chunk_for_small_batch(self): - engine, conn = _mock_engine_with_reflect("hook", ["score"]) - store = PostgresFeatureStore(engine=engine, session=AsyncMock()) + session = self._mock_session() + store = PostgresFeatureStore(engine=AsyncMock(), session=session) rows = [{"score": float(i)} for i in range(999)] count = await store.insert_features("hook", "urn:rec:1", rows, _RUN_ID) assert count == 999 - # 1 replace-DELETE + 1 insert chunk. - assert conn.execute.call_count == 2 + assert session.execute.call_count == 5 @pytest.mark.asyncio async def test_insert_rejects_invalid_hook_name(self): - engine = AsyncMock() - store = PostgresFeatureStore(engine=engine, session=AsyncMock()) + store = PostgresFeatureStore(engine=AsyncMock(), session=AsyncMock()) with pytest.raises(ValidationError, match="Invalid identifier"): await store.insert_features("'; DROP TABLE --", "urn:rec:1", [{"score": 1}], _RUN_ID) @@ -236,12 +264,14 @@ async def test_create_rejects_invalid_hook_name(self): await store.create_table("'; DROP TABLE --", _make_columns()) @pytest.mark.asyncio - async def test_reflects_table_before_insert(self): - """insert_features reflects the real table schema instead of guessing types.""" - engine, conn = _mock_engine_with_reflect("hook", ["score"]) - store = PostgresFeatureStore(engine=engine, session=AsyncMock()) - - await store.insert_features("hook", "urn:rec:1", [{"score": 0.95}], _RUN_ID) - - # run_sync should have been called for reflection - conn.run_sync.assert_called_once() + async def test_unregistered_feature_raises_not_found(self): + """The table is built from the feature_tables catalog — an unregistered + hook is a NotFoundError, never a guessed table shape.""" + session = AsyncMock() + missing = MagicMock() + missing.first.return_value = None + session.execute = AsyncMock(return_value=missing) + store = PostgresFeatureStore(engine=AsyncMock(), session=session) + + with pytest.raises(NotFoundError): + await store.insert_features("ghost_hook", "urn:rec:1", [{"score": 1.0}], _RUN_ID)