feat: harden datasource materialization payloads

2026-07-29 19:38:34 +02:00
parent 5f0847b8ca
commit f3288dbd28
9 changed files with 1460 additions and 179 deletions
@@ -17,11 +17,14 @@ from govoplan_core.core.datasources import (
    DatasourcePublicationRequest,
    DatasourceReadRequest,
    DatasourceStageInput,
+    DatasourceUnavailableError,
    DatasourceValidationError,
 )
 from govoplan_core.db.base import Base, utcnow
 from govoplan_datasources.backend.db.models import (
    DatasourceMaterializationRecord,
+    DatasourcePayloadRecord,
+    DatasourcePayloadRowRecord,
    DatasourcePublicationRecord,
    DatasourceRecord,
    DatasourceStageRecord,
@@ -32,6 +35,12 @@ from govoplan_datasources.backend.service import (
    STAGE_WRITE_SCOPE,
    SqlDatasourceProvider,
 )
+from govoplan_datasources.backend.payloads import (
+    create_database_rows_payload,
+    finalize_payload_deletion,
+    mark_unreferenced_payload_for_deletion,
+    verify_payload_integrity,
+)


 def principal(
@@ -137,6 +146,8 @@ class DatasourceLifecycleTests(unittest.TestCase):
            self.engine,
            tables=[
                DatasourceRecord.__table__,
+                DatasourcePayloadRecord.__table__,
+                DatasourcePayloadRowRecord.__table__,
                DatasourceMaterializationRecord.__table__,
                DatasourceStageRecord.__table__,
                DatasourcePublicationRecord.__table__,
@@ -157,6 +168,8 @@ class DatasourceLifecycleTests(unittest.TestCase):
                DatasourceStageRecord.__table__,
                DatasourcePublicationRecord.__table__,
                DatasourceMaterializationRecord.__table__,
+                DatasourcePayloadRowRecord.__table__,
+                DatasourcePayloadRecord.__table__,
                DatasourceRecord.__table__,
            ],
        )
@@ -235,6 +248,18 @@ class DatasourceLifecycleTests(unittest.TestCase):
            current.datasource.schema_version,
        )
        self.assertEqual(frozen.ref, frozen_result.materialization.ref)
+        first_record = self.session.get(
+            DatasourceMaterializationRecord,
+            first.ref.removeprefix("materialization:"),
+        )
+        frozen_record = self.session.get(
+            DatasourceMaterializationRecord,
+            frozen.ref.removeprefix("materialization:"),
+        )
+        self.assertIsNotNone(first_record)
+        self.assertIsNotNone(frozen_record)
+        self.assertEqual(first_record.payload_id, frozen_record.payload_id)
+        self.assertEqual([], first_record.rows)

    def test_live_reads_origin_and_cached_refresh_is_explicit(self) -> None:
        live = self.provider.register_origin(
@@ -445,6 +470,114 @@ class DatasourceLifecycleTests(unittest.TestCase):
                ),
            )

+    def test_payload_preview_is_paged_and_metadata_mismatch_is_rejected(self) -> None:
+        stage = self.provider.create_stage(
+            self.session,
+            principal(),
+            stage=DatasourceStageInput(
+                name="Paged",
+                source_name="paged",
+                kind="upload",
+                mode="static",
+                shape="tabular",
+                rows=tuple({"id": index} for index in range(20)),
+            ),
+        )
+        datasource, materialization = self.provider.promote_stage(
+            self.session,
+            principal(),
+            stage_ref=stage.ref,
+        )
+        self.session.commit()
+
+        result = self.provider.read_datasource(
+            self.session,
+            principal(),
+            request=DatasourceReadRequest(
+                datasource_ref=datasource.ref,
+                offset=7,
+                limit=3,
+            ),
+        )
+        self.assertEqual([{"id": 7}, {"id": 8}, {"id": 9}], list(result.rows))
+        record = self.session.get(
+            DatasourceMaterializationRecord,
+            materialization.ref.removeprefix("materialization:"),
+        )
+        self.assertIsNotNone(record)
+        self.assertEqual([], record.rows)
+        self.assertEqual(
+            20,
+            self.session.query(DatasourcePayloadRowRecord)
+            .filter(DatasourcePayloadRowRecord.payload_id == record.payload_id)
+            .count(),
+        )
+
+        payload = self.session.get(DatasourcePayloadRecord, record.payload_id)
+        self.assertIsNotNone(payload)
+        verify_payload_integrity(self.session, payload)
+        payload.row_count += 1
+        self.session.flush()
+        with self.assertRaises(DatasourceUnavailableError):
+            self.provider.read_datasource(
+                self.session,
+                principal(),
+                request=DatasourceReadRequest(datasource_ref=datasource.ref),
+            )
+
+    def test_payload_deletion_is_staged_and_reference_safe(self) -> None:
+        payload = create_database_rows_payload(
+            self.session,
+            tenant_id="tenant-1",
+            rows=({"id": 1}, {"id": 2}),
+            actor_id="account-1",
+        )
+        self.session.commit()
+
+        self.assertTrue(
+            mark_unreferenced_payload_for_deletion(self.session, payload)
+        )
+        self.assertEqual("deleting", payload.state)
+        self.session.commit()
+
+        finalize_payload_deletion(self.session, payload)
+        self.session.commit()
+        self.assertEqual(0, self.session.query(DatasourcePayloadRecord).count())
+        self.assertEqual(
+            0,
+            self.session.query(DatasourcePayloadRowRecord).count(),
+        )
+
+    def test_rolled_back_materialization_leaves_no_payload_rows(self) -> None:
+        stage = self.provider.create_stage(
+            self.session,
+            principal(),
+            stage=DatasourceStageInput(
+                name="Rollback",
+                source_name="rollback",
+                kind="upload",
+                mode="static",
+                shape="tabular",
+                rows=({"id": 1},),
+            ),
+        )
+        self.provider.promote_stage(
+            self.session,
+            principal(),
+            stage_ref=stage.ref,
+        )
+        self.session.rollback()
+
+        self.assertEqual(0, self.session.query(DatasourcePayloadRecord).count())
+        self.assertEqual(
+            0,
+            self.session.query(DatasourcePayloadRowRecord).count(),
+        )
+        self.assertEqual(
+            0,
+            self.session.query(DatasourceMaterializationRecord).count(),
+        )
+

 if __name__ == "__main__":
    unittest.main()
@@ -24,13 +24,15 @@ class DatasourceMigrationTests(unittest.TestCase):
            try:
                with engine.connect() as connection:
                    self.assertIn(
-                        "c4e9f1a7d2b6",
+                        "d5f0a2b8c3e7",
                        set(MigrationContext.configure(connection).get_current_heads()),
                    )
                    self.assertEqual(
                        {
                            "datasource_catalogue",
                            "datasource_materializations",
+                            "datasource_payload_rows",
+                            "datasource_payloads",
                            "datasource_publications",
                            "datasource_stages",
                        },
@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import os
+import threading
+import unittest
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+
+from sqlalchemy import create_engine, text
+from sqlalchemy.orm import Session
+
+from govoplan_core.db.base import Base
+from govoplan_datasources.backend.db.models import (
+    DatasourceMaterializationRecord,
+    DatasourcePayloadRecord,
+    DatasourcePayloadRowRecord,
+    DatasourceRecord,
+)
+from govoplan_datasources.backend.service import _append_materialization
+from govoplan_datasources.backend.tabular import (
+    encoded_size,
+    field_payload,
+    fingerprint_rows,
+    infer_schema,
+)
+
+
+@unittest.skipUnless(
+    os.environ.get("GOVOPLAN_DATASOURCES_TEST_POSTGRES_URL"),
+    "set GOVOPLAN_DATASOURCES_TEST_POSTGRES_URL for PostgreSQL concurrency checks",
+)
+class DatasourceMaterializationPostgresTests(unittest.TestCase):
+    def setUp(self) -> None:
+        database_url = os.environ["GOVOPLAN_DATASOURCES_TEST_POSTGRES_URL"]
+        self.schema = f"datasource_revision_race_{uuid.uuid4().hex}"
+        self.admin_engine = create_engine(database_url)
+        with self.admin_engine.begin() as connection:
+            connection.execute(text(f'CREATE SCHEMA "{self.schema}"'))
+        self.engine = create_engine(
+            database_url,
+            connect_args={"options": f"-c search_path={self.schema}"},
+        )
+        self.tables = [
+            DatasourceRecord.__table__,
+            DatasourcePayloadRecord.__table__,
+            DatasourcePayloadRowRecord.__table__,
+            DatasourceMaterializationRecord.__table__,
+        ]
+        Base.metadata.create_all(self.engine, tables=self.tables)
+        with Session(self.engine) as session:
+            session.add(
+                DatasourceRecord(
+                    id="datasource-1",
+                    tenant_id="tenant-1",
+                    source_name="concurrent",
+                    name="Concurrent",
+                    kind="custom",
+                    mode="static",
+                    shape="tabular",
+                )
+            )
+            session.commit()
+
+    def tearDown(self) -> None:
+        try:
+            Base.metadata.drop_all(
+                self.engine,
+                tables=list(reversed(self.tables)),
+            )
+        finally:
+            self.engine.dispose()
+            with self.admin_engine.begin() as connection:
+                connection.execute(text(f'DROP SCHEMA "{self.schema}"'))
+            self.admin_engine.dispose()
+
+    def test_competing_publishers_receive_distinct_monotonic_revisions(self) -> None:
+        barrier = threading.Barrier(2)
+
+        def publish(value: int) -> int:
+            rows = ({"id": value},)
+            schema = infer_schema(rows)
+            with Session(self.engine) as session:
+                datasource = session.get(DatasourceRecord, "datasource-1")
+                self.assertIsNotNone(datasource)
+                barrier.wait(timeout=10)
+                materialization = _append_materialization(
+                    session,
+                    datasource=datasource,
+                    rows=rows,
+                    schema=[field_payload(field) for field in schema],
+                    fingerprint=fingerprint_rows(rows, schema),
+                    byte_count=encoded_size(rows),
+                    actor_id=f"worker-{value}",
+                    set_current=True,
+                )
+                revision = materialization.revision
+                session.commit()
+                return revision
+
+        with ThreadPoolExecutor(max_workers=2) as executor:
+            revisions = tuple(executor.map(publish, (1, 2)))
+
+        self.assertEqual((1, 2), tuple(sorted(revisions)))
+        with Session(self.engine) as session:
+            self.assertEqual(
+                [1, 2],
+                list(
+                    session.scalars(
+                        DatasourceMaterializationRecord.__table__.select()
+                        .with_only_columns(
+                            DatasourceMaterializationRecord.revision
+                        )
+                        .order_by(
+                            DatasourceMaterializationRecord.revision.asc()
+                        )
+                    )
+                ),
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()