123 lines
4.2 KiB
Python
123 lines
4.2 KiB
Python
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()
|