feat: harden datasource materialization payloads

src/govoplan_datasources/backend/service.py

@@ -4,7 +4,7 @@ import hashlib
 import json
 import re
 from collections.abc import Mapping, Sequence
-from dataclasses import replace
+from dataclasses import dataclass, replace
 from typing import Any, cast
 
 from sqlalchemy import func, or_, select
@@ -34,10 +34,18 @@ from govoplan_core.core.datasources import (
 from govoplan_core.db.base import utcnow
 from govoplan_datasources.backend.db.models import (
     DatasourceMaterializationRecord,
+    DatasourcePayloadRecord,
     DatasourcePublicationRecord,
     DatasourceRecord,
     DatasourceStageRecord,
 )
+from govoplan_datasources.backend.payloads import (
+    DatasourcePayloadBackend,
+    PayloadBackendRegistry,
+    create_database_rows_payload,
+    payload_for_materialization,
+    validate_payload_size,
+)
 from govoplan_datasources.backend.tabular import (
     MAX_READ_ROWS,
     MAX_STAGE_ROWS,
@@ -55,9 +63,27 @@ STAGE_WRITE_SCOPE = "datasources:stage:write"
 ADMIN_SCOPE = "datasources:source:admin"
 
 
+@dataclass(frozen=True, slots=True)
+class _PreparedPublication:
+    producer_module: str
+    producer_run_ref: str
+    idempotency_key: str
+    rows: tuple[dict[str, Any], ...]
+    schema: tuple[DatasourceField, ...]
+    fingerprint: str
+    byte_count: int
+    request_hash: str
+
+
 class SqlDatasourceProvider:
-    def __init__(self, *, registry: object | None = None) -> None:
+    def __init__(
+        self,
+        *,
+        registry: object | None = None,
+        payload_backends: Sequence[DatasourcePayloadBackend] = (),
+    ) -> None:
         self._registry = registry
+        self._payload_backends = PayloadBackendRegistry(payload_backends)
 
     def publish_rows(
         self,
@@ -67,162 +93,49 @@ class SqlDatasourceProvider:
         request: DatasourcePublicationRequest,
     ) -> DatasourcePublicationResult:
         db, api_principal = _publication_context(session, principal)
-        producer_module = request.producer_module.strip()
-        producer_run_ref = request.producer_run_ref.strip()
-        idempotency_key = request.idempotency_key.strip()
-        if not producer_module or len(producer_module) > 100:
-            raise DatasourceValidationError(
-                "A producer module of at most 100 characters is required."
-            )
-        if not producer_run_ref or len(producer_run_ref) > 500:
-            raise DatasourceValidationError(
-                "A producer run reference of at most 500 characters is required."
-            )
-        if not idempotency_key or len(idempotency_key) > 255:
-            raise DatasourceValidationError(
-                "An idempotency key of at most 255 characters is required."
-            )
-        normalized = normalize_rows(request.rows)
-        schema = infer_schema(normalized)
-        fingerprint = fingerprint_rows(normalized, schema)
-        request_hash = _publication_request_hash(
-            request,
-            normalized=normalized,
-            fingerprint=fingerprint,
-        )
-        existing = db.scalar(
-            select(DatasourcePublicationRecord).where(
-                DatasourcePublicationRecord.tenant_id
-                == api_principal.tenant_id,
-                DatasourcePublicationRecord.producer_module == producer_module,
-                DatasourcePublicationRecord.idempotency_key == idempotency_key,
-            )
+        prepared = _prepare_publication(request)
+        existing = _existing_publication_result(
+            db,
+            tenant_id=api_principal.tenant_id,
+            producer_module=prepared.producer_module,
+            idempotency_key=prepared.idempotency_key,
+            request_hash=prepared.request_hash,
         )
         if existing is not None:
-            if existing.request_hash != request_hash:
-                raise DatasourceValidationError(
-                    "The publication idempotency key was already used with "
-                    "different output."
-                )
-            datasource = db.get(DatasourceRecord, existing.datasource_id)
-            materialization = db.get(
-                DatasourceMaterializationRecord,
-                existing.materialization_id,
-            )
-            if (
-                datasource is None
-                or materialization is None
-                or datasource.tenant_id != api_principal.tenant_id
-                or materialization.tenant_id != api_principal.tenant_id
-            ):
-                raise DatasourceUnavailableError(
-                    "The prior publication result is no longer available."
-                )
-            return DatasourcePublicationResult(
-                ref=_publication_ref(existing.id),
-                status=existing.status,
-                datasource=_datasource_dto(datasource),
-                materialization=_materialization_dto(materialization),
-                replayed=True,
-            )
+            return existing
 
-        datasource = None
-        if request.target_datasource_ref:
-            datasource = _required_datasource(
-                db,
-                tenant_id=api_principal.tenant_id,
-                datasource_ref=request.target_datasource_ref,
-            )
-            if datasource.mode == "live" or datasource.shape != "tabular":
-                raise DatasourceValidationError(
-                    "Produced rows require a static or cached tabular datasource."
-                )
-        else:
-            name = str(request.name or "").strip()
-            source_name = str(request.source_name or "").strip()
-            if not name:
-                raise DatasourceValidationError(
-                    "A datasource name is required for a new publication target."
-                )
-            if not _valid_source_name(source_name):
-                raise DatasourceValidationError(
-                    "Datasource keys must start with a letter or underscore and "
-                    "contain only letters, numbers, and underscores."
-                )
-            _ensure_source_name_available(
-                db,
-                tenant_id=api_principal.tenant_id,
-                source_name=source_name,
-            )
-            datasource = DatasourceRecord(
-                tenant_id=api_principal.tenant_id,
-                source_name=source_name,
-                name=name,
-                description=_clean_optional(request.description),
-                kind="custom",
-                mode="static",
-                shape="tabular",
-                status="active",
-                provider=producer_module,
-                provider_ref=producer_run_ref,
-                schema_version=1,
-                schema_=[field_payload(field) for field in schema],
-                fingerprint=fingerprint,
-                row_count=len(normalized),
-                byte_count=encoded_size(normalized),
-                provenance_={
-                    **dict(request.provenance),
-                    "producer_module": producer_module,
-                    "producer_run_ref": producer_run_ref,
-                },
-                metadata_=dict(request.metadata),
-                created_by=_actor_id(api_principal),
-                updated_by=_actor_id(api_principal),
-            )
-            db.add(datasource)
-            db.flush()
-
-        publication_provenance = {
-            **dict(request.provenance),
-            "producer_module": producer_module,
-            "producer_run_ref": producer_run_ref,
-            "idempotency_key": idempotency_key,
-            "published_at": utcnow().isoformat(),
-        }
+        actor_id = _actor_id(api_principal)
+        datasource = _publication_target(
+            db,
+            tenant_id=api_principal.tenant_id,
+            actor_id=actor_id,
+            request=request,
+            prepared=prepared,
+        )
         materialization = _append_materialization(
             db,
             datasource=datasource,
-            rows=normalized,
-            schema=[field_payload(field) for field in schema],
-            fingerprint=fingerprint,
-            byte_count=encoded_size(normalized),
-            actor_id=_actor_id(api_principal),
+            rows=prepared.rows,
+            schema=[field_payload(field) for field in prepared.schema],
+            fingerprint=prepared.fingerprint,
+            byte_count=prepared.byte_count,
+            actor_id=actor_id,
             frozen=request.freeze,
             frozen_label=request.frozen_label,
             source_timestamp=request.source_timestamp,
-            provenance=publication_provenance,
+            provenance=_publication_provenance(request, prepared),
             metadata=dict(request.metadata),
             set_current=request.set_current,
         )
-        publication = DatasourcePublicationRecord(
+        publication = _create_publication_record(
+            db,
             tenant_id=api_principal.tenant_id,
-            producer_module=producer_module,
-            producer_run_ref=producer_run_ref,
-            idempotency_key=idempotency_key,
-            request_hash=request_hash,
-            datasource_id=datasource.id,
-            materialization_id=materialization.id,
-            status="published",
-            details_={
-                "fingerprint": fingerprint,
-                "row_count": len(normalized),
-                "set_current": request.set_current,
-                "frozen": request.freeze,
-            },
-            created_by=_actor_id(api_principal),
+            actor_id=actor_id,
+            datasource=datasource,
+            materialization=materialization,
+            request=request,
+            prepared=prepared,
         )
-        db.add(publication)
-        db.flush()
         return DatasourcePublicationResult(
             ref=_publication_ref(publication.id),
             status=publication.status,
@@ -293,13 +206,12 @@ class SqlDatasourceProvider:
         offset = max(0, int(request.offset))
         columns = tuple(dict.fromkeys(request.columns))
 
-        direct_live = request.consistency == "live"
-        implicit_live = (
+        read_live = request.consistency == "live" or (
             item.mode == "live"
             and not request.materialization_ref
             and request.consistency == "current"
         )
-        if direct_live or implicit_live:
+        if read_live:
             if not item.provider_ref:
                 raise DatasourceUnavailableError(
                     "This datasource has no live origin."
@@ -322,10 +234,29 @@ class SqlDatasourceProvider:
                 )
             return result
 
-        materialization = _selected_materialization(
+        return self._read_materialized(
             db,
             item=item,
             request=request,
+            limit=limit,
+            offset=offset,
+            columns=columns,
+        )
+
+    def _read_materialized(
+        self,
+        session: Session,
+        *,
+        item: DatasourceRecord,
+        request: DatasourceReadRequest,
+        limit: int,
+        offset: int,
+        columns: tuple[str, ...],
+    ) -> DatasourceReadResult:
+        materialization = _selected_materialization(
+            session,
+            item=item,
+            request=request,
         )
         if materialization is None:
             message = (
@@ -342,7 +273,18 @@ class SqlDatasourceProvider:
                 "The datasource fingerprint changed; refresh the consuming definition."
             )
         _validate_columns(materialization.schema_, columns)
-        window = materialization.rows[offset : offset + limit]
+        payload = payload_for_materialization(session, materialization)
+        if payload is None:
+            window = materialization.rows[offset : offset + limit]
+        else:
+            backend = self._payload_backends.require(payload.backend)
+            backend.verify(session, payload)
+            window = backend.read_rows(
+                session,
+                payload,
+                offset=offset,
+                limit=limit,
+            )
         rows = tuple(_select_columns(row, columns) for row in window)
         descriptor = _datasource_dto(item)
         return DatasourceReadResult(
@@ -698,10 +640,11 @@ class SqlDatasourceProvider:
             raise DatasourceUnavailableError(
                 "The datasource has no current state to freeze."
             )
+        current_payload = payload_for_materialization(db, current)
         materialization = _append_materialization(
             db,
             datasource=item,
-            rows=current.rows,
+            rows=current.rows if current_payload is None else (),
             schema=current.schema_,
             fingerprint=current.fingerprint,
             byte_count=current.byte_count,
@@ -715,6 +658,7 @@ class SqlDatasourceProvider:
             },
             metadata=dict(current.metadata_),
             set_current=False,
+            reusable_payload=current_payload,
         )
         return _materialization_dto(materialization)
 
@@ -918,21 +862,28 @@ def _append_materialization(
     provenance: Mapping[str, object] | None = None,
     metadata: Mapping[str, object] | None = None,
     set_current: bool,
+    reusable_payload: DatasourcePayloadRecord | None = None,
 ) -> DatasourceMaterializationRecord:
-    revision = (
-        session.scalar(
-            select(func.max(DatasourceMaterializationRecord.revision)).where(
-                DatasourceMaterializationRecord.datasource_id == datasource.id
-            )
+    datasource = _lock_datasource_for_materialization(session, datasource)
+    revision = _allocate_materialization_revision(session, datasource)
+    payload = reusable_payload or create_database_rows_payload(
+        session,
+        tenant_id=datasource.tenant_id,
+        rows=rows,
+        actor_id=actor_id,
+        metadata={
+            "datasource_id": datasource.id,
+            "fingerprint": fingerprint,
+        },
+    )
+    if payload.tenant_id != datasource.tenant_id or payload.state != "ready":
+        raise DatasourceValidationError(
+            "Only a ready payload from the same tenant can be materialized."
         )
-        or 0
-    ) + 1
-    schema_payload = [dict(field) for field in schema]
-    schema_changed = datasource.schema_ != schema_payload
-    schema_version = (
-        datasource.schema_version + 1
-        if schema_changed
-        else datasource.schema_version
+    validate_payload_size(payload, expected_byte_count=byte_count)
+    schema_payload, schema_version = _materialization_schema(
+        datasource,
+        schema,
     )
     materialization = DatasourceMaterializationRecord(
         tenant_id=datasource.tenant_id,
@@ -941,10 +892,12 @@ def _append_materialization(
         state="published",
         schema_version=max(1, int(schema_version or 1)),
         schema_=schema_payload,
-        rows=[dict(row) for row in rows],
+        payload_id=payload.id,
+        payload_checksum=payload.checksum,
+        rows=[],
         fingerprint=fingerprint,
-        row_count=len(rows),
-        byte_count=byte_count,
+        row_count=payload.row_count,
+        byte_count=payload.byte_count,
         frozen_at=utcnow() if frozen else None,
         frozen_label=_clean_optional(frozen_label),
         source_timestamp=source_timestamp,
@@ -955,17 +908,85 @@ def _append_materialization(
     session.add(materialization)
     session.flush()
     if set_current:
-        datasource.current_materialization_id = materialization.id
-        datasource.schema_ = schema_payload
-        datasource.schema_version = schema_version
-        datasource.fingerprint = fingerprint
-        datasource.row_count = materialization.row_count
-        datasource.byte_count = materialization.byte_count
-        datasource.updated_by = actor_id
+        _apply_current_materialization(
+            datasource,
+            materialization=materialization,
+            schema=schema_payload,
+            schema_version=schema_version,
+            actor_id=actor_id,
+        )
     session.flush()
     return materialization
 
 
+def _lock_datasource_for_materialization(
+    session: Session,
+    datasource: DatasourceRecord,
+) -> DatasourceRecord:
+    locked = session.scalar(
+        select(DatasourceRecord)
+        .where(
+            DatasourceRecord.id == datasource.id,
+            DatasourceRecord.tenant_id == datasource.tenant_id,
+            DatasourceRecord.deleted_at.is_(None),
+        )
+        .with_for_update()
+        .execution_options(populate_existing=True)
+    )
+    if locked is None:
+        raise DatasourceUnavailableError(
+            "The datasource is no longer available for materialization."
+        )
+    return locked
+
+
+def _allocate_materialization_revision(
+    session: Session,
+    datasource: DatasourceRecord,
+) -> int:
+    """Allocate under the datasource row lock held by the caller."""
+
+    return int(
+        session.scalar(
+            select(func.max(DatasourceMaterializationRecord.revision)).where(
+                DatasourceMaterializationRecord.datasource_id == datasource.id
+            )
+        )
+        or 0
+    ) + 1
+
+
+def _materialization_schema(
+    datasource: DatasourceRecord,
+    schema: Sequence[Mapping[str, object]],
+) -> tuple[list[dict[str, object]], int]:
+    schema_payload = [dict(field) for field in schema]
+    schema_changed = datasource.schema_ != schema_payload
+    schema_version = (
+        datasource.schema_version + 1
+        if schema_changed
+        else datasource.schema_version
+    )
+    return schema_payload, int(schema_version or 1)
+
+
+def _apply_current_materialization(
+    datasource: DatasourceRecord,
+    *,
+    materialization: DatasourceMaterializationRecord,
+    schema: list[dict[str, object]],
+    schema_version: int,
+    actor_id: str | None,
+) -> None:
+    datasource.current_materialization_id = materialization.id
+    datasource.schema_ = schema
+    datasource.schema_version = schema_version
+    datasource.fingerprint = materialization.fingerprint
+    datasource.row_count = materialization.row_count
+    datasource.byte_count = materialization.byte_count
+    datasource.updated_by = actor_id
+
+
 def _selected_materialization(
     session: Session,
     *,
@@ -1295,6 +1316,208 @@ def _valid_source_name(value: str) -> bool:
     return bool(re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]{0,119}", value))
 
 
+def _prepare_publication(
+    request: DatasourcePublicationRequest,
+) -> _PreparedPublication:
+    producer_module = request.producer_module.strip()
+    producer_run_ref = request.producer_run_ref.strip()
+    idempotency_key = request.idempotency_key.strip()
+    _validate_publication_identity(
+        producer_module=producer_module,
+        producer_run_ref=producer_run_ref,
+        idempotency_key=idempotency_key,
+    )
+    normalized = normalize_rows(request.rows)
+    schema = infer_schema(normalized)
+    fingerprint = fingerprint_rows(normalized, schema)
+    return _PreparedPublication(
+        producer_module=producer_module,
+        producer_run_ref=producer_run_ref,
+        idempotency_key=idempotency_key,
+        rows=normalized,
+        schema=schema,
+        fingerprint=fingerprint,
+        byte_count=encoded_size(normalized),
+        request_hash=_publication_request_hash(
+            request,
+            normalized=normalized,
+            fingerprint=fingerprint,
+        ),
+    )
+
+
+def _validate_publication_identity(
+    *,
+    producer_module: str,
+    producer_run_ref: str,
+    idempotency_key: str,
+) -> None:
+    values = (
+        (producer_module, 100, "A producer module"),
+        (producer_run_ref, 500, "A producer run reference"),
+        (idempotency_key, 255, "An idempotency key"),
+    )
+    for value, maximum, label in values:
+        if not value or len(value) > maximum:
+            raise DatasourceValidationError(
+                f"{label} of at most {maximum} characters is required."
+            )
+
+
+def _existing_publication_result(
+    session: Session,
+    *,
+    tenant_id: str,
+    producer_module: str,
+    idempotency_key: str,
+    request_hash: str,
+) -> DatasourcePublicationResult | None:
+    publication = session.scalar(
+        select(DatasourcePublicationRecord).where(
+            DatasourcePublicationRecord.tenant_id == tenant_id,
+            DatasourcePublicationRecord.producer_module == producer_module,
+            DatasourcePublicationRecord.idempotency_key == idempotency_key,
+        )
+    )
+    if publication is None:
+        return None
+    if publication.request_hash != request_hash:
+        raise DatasourceValidationError(
+            "The publication idempotency key was already used with different output."
+        )
+    datasource = session.get(DatasourceRecord, publication.datasource_id)
+    materialization = session.get(
+        DatasourceMaterializationRecord,
+        publication.materialization_id,
+    )
+    if (
+        datasource is None
+        or materialization is None
+        or datasource.tenant_id != tenant_id
+        or materialization.tenant_id != tenant_id
+    ):
+        raise DatasourceUnavailableError(
+            "The prior publication result is no longer available."
+        )
+    return DatasourcePublicationResult(
+        ref=_publication_ref(publication.id),
+        status=publication.status,
+        datasource=_datasource_dto(datasource),
+        materialization=_materialization_dto(materialization),
+        replayed=True,
+    )
+
+
+def _publication_target(
+    session: Session,
+    *,
+    tenant_id: str,
+    actor_id: str | None,
+    request: DatasourcePublicationRequest,
+    prepared: _PreparedPublication,
+) -> DatasourceRecord:
+    if request.target_datasource_ref:
+        datasource = _required_datasource(
+            session,
+            tenant_id=tenant_id,
+            datasource_ref=request.target_datasource_ref,
+        )
+        if datasource.mode == "live" or datasource.shape != "tabular":
+            raise DatasourceValidationError(
+                "Produced rows require a static or cached tabular datasource."
+            )
+        return datasource
+    name = str(request.name or "").strip()
+    source_name = str(request.source_name or "").strip()
+    if not name:
+        raise DatasourceValidationError(
+            "A datasource name is required for a new publication target."
+        )
+    if not _valid_source_name(source_name):
+        raise DatasourceValidationError(
+            "Datasource keys must start with a letter or underscore and contain "
+            "only letters, numbers, and underscores."
+        )
+    _ensure_source_name_available(
+        session,
+        tenant_id=tenant_id,
+        source_name=source_name,
+    )
+    datasource = DatasourceRecord(
+        tenant_id=tenant_id,
+        source_name=source_name,
+        name=name,
+        description=_clean_optional(request.description),
+        kind="custom",
+        mode="static",
+        shape="tabular",
+        status="active",
+        provider=prepared.producer_module,
+        provider_ref=prepared.producer_run_ref,
+        schema_version=1,
+        schema_=[field_payload(field) for field in prepared.schema],
+        fingerprint=prepared.fingerprint,
+        row_count=len(prepared.rows),
+        byte_count=prepared.byte_count,
+        provenance_={
+            **dict(request.provenance),
+            "producer_module": prepared.producer_module,
+            "producer_run_ref": prepared.producer_run_ref,
+        },
+        metadata_=dict(request.metadata),
+        created_by=actor_id,
+        updated_by=actor_id,
+    )
+    session.add(datasource)
+    session.flush()
+    return datasource
+
+
+def _publication_provenance(
+    request: DatasourcePublicationRequest,
+    prepared: _PreparedPublication,
+) -> dict[str, object]:
+    return {
+        **dict(request.provenance),
+        "producer_module": prepared.producer_module,
+        "producer_run_ref": prepared.producer_run_ref,
+        "idempotency_key": prepared.idempotency_key,
+        "published_at": utcnow().isoformat(),
+    }
+
+
+def _create_publication_record(
+    session: Session,
+    *,
+    tenant_id: str,
+    actor_id: str | None,
+    datasource: DatasourceRecord,
+    materialization: DatasourceMaterializationRecord,
+    request: DatasourcePublicationRequest,
+    prepared: _PreparedPublication,
+) -> DatasourcePublicationRecord:
+    publication = DatasourcePublicationRecord(
+        tenant_id=tenant_id,
+        producer_module=prepared.producer_module,
+        producer_run_ref=prepared.producer_run_ref,
+        idempotency_key=prepared.idempotency_key,
+        request_hash=prepared.request_hash,
+        datasource_id=datasource.id,
+        materialization_id=materialization.id,
+        status="published",
+        details_={
+            "fingerprint": prepared.fingerprint,
+            "row_count": len(prepared.rows),
+            "set_current": request.set_current,
+            "frozen": request.freeze,
+        },
+        created_by=actor_id,
+    )
+    session.add(publication)
+    session.flush()
+    return publication
+
+
 def _publication_request_hash(
     request: DatasourcePublicationRequest,
     *,