Serialize concurrent datasource publications

tests/test_postgres_materialization_concurrency.py

@@ -9,14 +9,23 @@ from concurrent.futures import ThreadPoolExecutor
 from sqlalchemy import create_engine, text
 from sqlalchemy.orm import Session
 
+from govoplan_core.auth import ApiPrincipal
+from govoplan_core.core.access import PrincipalRef
+from govoplan_core.core.change_sequence import ChangeSequenceEntry
+from govoplan_core.core.datasources import DatasourcePublicationRequest
 from govoplan_core.db.base import Base
 from govoplan_datasources.backend.db.models import (
     DatasourceMaterializationRecord,
     DatasourcePayloadRecord,
     DatasourcePayloadRowRecord,
+    DatasourcePublicationRecord,
     DatasourceRecord,
 )
-from govoplan_datasources.backend.service import _append_materialization
+from govoplan_datasources.backend.service import (
+    SOURCE_WRITE_SCOPE,
+    SqlDatasourceProvider,
+    _append_materialization,
+)
 from govoplan_datasources.backend.tabular import (
     encoded_size,
     field_payload,
@@ -45,6 +54,8 @@ class DatasourceMaterializationPostgresTests(unittest.TestCase):
             DatasourcePayloadRecord.__table__,
             DatasourcePayloadRowRecord.__table__,
             DatasourceMaterializationRecord.__table__,
+            DatasourcePublicationRecord.__table__,
+            ChangeSequenceEntry.__table__,
         ]
         Base.metadata.create_all(self.engine, tables=self.tables)
         with Session(self.engine) as session:
@@ -117,6 +128,55 @@ class DatasourceMaterializationPostgresTests(unittest.TestCase):
                 ),
             )
 
+    def test_same_publication_identity_is_serialized_and_replayed(self) -> None:
+        barrier = threading.Barrier(2)
+        provider = SqlDatasourceProvider()
+        request = DatasourcePublicationRequest(
+            producer_module="dataflow",
+            producer_run_ref="dataflow-run:concurrent",
+            idempotency_key="concurrent-publication",
+            name="Concurrent publication",
+            source_name="concurrent_publication",
+            rows=({"id": 1, "result": "match"},),
+        )
+
+        def publish(worker: int) -> tuple[str, str, bool]:
+            api_principal = ApiPrincipal(
+                principal=PrincipalRef(
+                    account_id=f"worker-{worker}",
+                    membership_id=f"membership-{worker}",
+                    tenant_id="tenant-1",
+                    scopes=frozenset((SOURCE_WRITE_SCOPE,)),
+                ),
+                account=object(),
+                user=object(),
+            )
+            with Session(self.engine) as session:
+                barrier.wait(timeout=10)
+                result = provider.publish_rows(
+                    session,
+                    api_principal,
+                    request=request,
+                )
+                session.commit()
+                return result.ref, result.materialization.ref, result.replayed
+
+        with ThreadPoolExecutor(max_workers=2) as executor:
+            results = tuple(executor.map(publish, (1, 2)))
+
+        self.assertEqual(1, len({item[0] for item in results}))
+        self.assertEqual(1, len({item[1] for item in results}))
+        self.assertEqual([False, True], sorted(item[2] for item in results))
+        with Session(self.engine) as session:
+            self.assertEqual(
+                1,
+                session.query(DatasourcePublicationRecord).count(),
+            )
+            self.assertEqual(
+                1,
+                session.query(DatasourceMaterializationRecord).count(),
+            )
+
 
 if __name__ == "__main__":
     unittest.main()