451 lines
15 KiB
Python
451 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
from govoplan_core.auth import ApiPrincipal
|
|
from govoplan_core.core.access import PrincipalRef
|
|
from govoplan_core.core.datasources import (
|
|
CAPABILITY_DATASOURCE_ORIGINS,
|
|
DatasourceAccessError,
|
|
DatasourceField,
|
|
DatasourceOrigin,
|
|
DatasourceOriginReadRequest,
|
|
DatasourceOriginReadResult,
|
|
DatasourcePublicationRequest,
|
|
DatasourceReadRequest,
|
|
DatasourceStageInput,
|
|
DatasourceValidationError,
|
|
)
|
|
from govoplan_core.db.base import Base, utcnow
|
|
from govoplan_datasources.backend.db.models import (
|
|
DatasourceMaterializationRecord,
|
|
DatasourcePublicationRecord,
|
|
DatasourceRecord,
|
|
DatasourceStageRecord,
|
|
)
|
|
from govoplan_datasources.backend.service import (
|
|
CATALOGUE_READ_SCOPE,
|
|
SOURCE_WRITE_SCOPE,
|
|
STAGE_WRITE_SCOPE,
|
|
SqlDatasourceProvider,
|
|
)
|
|
|
|
|
|
def principal(
|
|
tenant_id: str = "tenant-1",
|
|
*,
|
|
scopes: tuple[str, ...] = (
|
|
CATALOGUE_READ_SCOPE,
|
|
SOURCE_WRITE_SCOPE,
|
|
STAGE_WRITE_SCOPE,
|
|
),
|
|
) -> ApiPrincipal:
|
|
return ApiPrincipal(
|
|
principal=PrincipalRef(
|
|
account_id="account-1",
|
|
membership_id="membership-1",
|
|
tenant_id=tenant_id,
|
|
scopes=frozenset(scopes),
|
|
),
|
|
account=object(),
|
|
user=object(),
|
|
)
|
|
|
|
|
|
class FakeOriginProvider:
|
|
def __init__(self) -> None:
|
|
self.rows = [{"id": 1, "name": "Initial"}]
|
|
|
|
def _origin(self) -> DatasourceOrigin:
|
|
return DatasourceOrigin(
|
|
ref="snapshot:origin-1",
|
|
source_name="connector_cases",
|
|
name="Connector cases",
|
|
kind="database",
|
|
shape="tabular",
|
|
supported_modes=("live", "cached"),
|
|
provider="connectors.test",
|
|
schema=(
|
|
DatasourceField(name="id", data_type="integer", nullable=False),
|
|
DatasourceField(name="name", data_type="string", nullable=False),
|
|
),
|
|
schema_version="1",
|
|
fingerprint=f"version-{len(self.rows)}-{self.rows[-1]['name']}",
|
|
row_count=len(self.rows),
|
|
updated_at=utcnow(),
|
|
)
|
|
|
|
def list_origins(
|
|
self,
|
|
_session: object,
|
|
_principal: object,
|
|
*,
|
|
query: str = "",
|
|
limit: int = 100,
|
|
):
|
|
origin = self._origin()
|
|
return (origin,) if query.lower() in origin.name.lower() else ()
|
|
|
|
def get_origin(
|
|
self,
|
|
_session: object,
|
|
_principal: object,
|
|
*,
|
|
origin_ref: str,
|
|
):
|
|
return self._origin() if origin_ref == "snapshot:origin-1" else None
|
|
|
|
def read_origin(
|
|
self,
|
|
_session: object,
|
|
_principal: object,
|
|
*,
|
|
request: DatasourceOriginReadRequest,
|
|
) -> DatasourceOriginReadResult:
|
|
origin = self.get_origin(None, None, origin_ref=request.origin_ref)
|
|
if origin is None:
|
|
raise AssertionError("Unexpected origin")
|
|
rows = self.rows[request.offset : request.offset + request.limit]
|
|
return DatasourceOriginReadResult(
|
|
origin=origin,
|
|
rows=tuple(dict(row) for row in rows),
|
|
total_rows=len(self.rows),
|
|
truncated=request.offset + len(rows) < len(self.rows),
|
|
)
|
|
|
|
|
|
class FakeRegistry:
|
|
def __init__(self, origin_provider: FakeOriginProvider) -> None:
|
|
self.origin_provider = origin_provider
|
|
|
|
def has_capability(self, name: str) -> bool:
|
|
return name == CAPABILITY_DATASOURCE_ORIGINS
|
|
|
|
def capability(self, name: str) -> object:
|
|
if not self.has_capability(name):
|
|
raise KeyError(name)
|
|
return self.origin_provider
|
|
|
|
|
|
class DatasourceLifecycleTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.engine = create_engine("sqlite:///:memory:")
|
|
Base.metadata.create_all(
|
|
self.engine,
|
|
tables=[
|
|
DatasourceRecord.__table__,
|
|
DatasourceMaterializationRecord.__table__,
|
|
DatasourceStageRecord.__table__,
|
|
DatasourcePublicationRecord.__table__,
|
|
],
|
|
)
|
|
self.Session = sessionmaker(bind=self.engine)
|
|
self.session = self.Session()
|
|
self.origins = FakeOriginProvider()
|
|
self.provider = SqlDatasourceProvider(
|
|
registry=FakeRegistry(self.origins),
|
|
)
|
|
|
|
def tearDown(self) -> None:
|
|
self.session.close()
|
|
Base.metadata.drop_all(
|
|
self.engine,
|
|
tables=[
|
|
DatasourceStageRecord.__table__,
|
|
DatasourcePublicationRecord.__table__,
|
|
DatasourceMaterializationRecord.__table__,
|
|
DatasourceRecord.__table__,
|
|
],
|
|
)
|
|
self.engine.dispose()
|
|
|
|
def test_static_stage_promote_update_and_frozen_read(self) -> None:
|
|
first_stage = self.provider.create_stage(
|
|
self.session,
|
|
principal(),
|
|
stage=DatasourceStageInput(
|
|
name="Monthly cases",
|
|
source_name="monthly_cases",
|
|
kind="upload",
|
|
mode="static",
|
|
shape="tabular",
|
|
rows=({"id": 1, "status": "new"},),
|
|
),
|
|
)
|
|
datasource, first = self.provider.promote_stage(
|
|
self.session,
|
|
principal(),
|
|
stage_ref=first_stage.ref,
|
|
)
|
|
frozen = self.provider.freeze_datasource(
|
|
self.session,
|
|
principal(),
|
|
datasource_ref=datasource.ref,
|
|
label="Import evidence",
|
|
)
|
|
second_stage = self.provider.create_stage(
|
|
self.session,
|
|
principal(),
|
|
stage=DatasourceStageInput(
|
|
name="Monthly cases",
|
|
source_name="monthly_cases",
|
|
kind="upload",
|
|
mode="static",
|
|
shape="tabular",
|
|
target_datasource_ref=datasource.ref,
|
|
rows=({"id": 2, "status": "checked", "note": "schema changed"},),
|
|
),
|
|
)
|
|
updated, second = self.provider.promote_stage(
|
|
self.session,
|
|
principal(),
|
|
stage_ref=second_stage.ref,
|
|
)
|
|
self.session.commit()
|
|
|
|
current = self.provider.read_datasource(
|
|
self.session,
|
|
principal(),
|
|
request=DatasourceReadRequest(datasource_ref=datasource.ref),
|
|
)
|
|
frozen_result = self.provider.read_datasource(
|
|
self.session,
|
|
principal(),
|
|
request=DatasourceReadRequest(
|
|
datasource_ref=datasource.ref,
|
|
consistency="frozen",
|
|
),
|
|
)
|
|
|
|
self.assertEqual(1, first.revision)
|
|
self.assertEqual(2, frozen.revision)
|
|
self.assertEqual(3, second.revision)
|
|
self.assertEqual(
|
|
[{"id": 2, "status": "checked", "note": "schema changed"}],
|
|
list(current.rows),
|
|
)
|
|
self.assertEqual([{"id": 1, "status": "new"}], list(frozen_result.rows))
|
|
self.assertEqual(second.ref, updated.current_materialization_ref)
|
|
self.assertEqual("2", updated.schema_version)
|
|
self.assertEqual(
|
|
"2",
|
|
current.datasource.schema_version,
|
|
)
|
|
self.assertEqual(frozen.ref, frozen_result.materialization.ref)
|
|
|
|
def test_live_reads_origin_and_cached_refresh_is_explicit(self) -> None:
|
|
live = self.provider.register_origin(
|
|
self.session,
|
|
principal(),
|
|
origin_ref="snapshot:origin-1",
|
|
name="Live cases",
|
|
source_name="live_cases",
|
|
mode="live",
|
|
)
|
|
cached = self.provider.register_origin(
|
|
self.session,
|
|
principal(),
|
|
origin_ref="snapshot:origin-1",
|
|
name="Cached cases",
|
|
source_name="cached_cases",
|
|
mode="cached",
|
|
)
|
|
self.session.commit()
|
|
|
|
self.origins.rows = [
|
|
{"id": 1, "name": "Initial"},
|
|
{"id": 2, "name": "New"},
|
|
]
|
|
live_read = self.provider.read_datasource(
|
|
self.session,
|
|
principal(),
|
|
request=DatasourceReadRequest(datasource_ref=live.ref),
|
|
)
|
|
cached_before = self.provider.read_datasource(
|
|
self.session,
|
|
principal(),
|
|
request=DatasourceReadRequest(datasource_ref=cached.ref),
|
|
)
|
|
cached_live = self.provider.read_datasource(
|
|
self.session,
|
|
principal(),
|
|
request=DatasourceReadRequest(
|
|
datasource_ref=cached.ref,
|
|
consistency="live",
|
|
),
|
|
)
|
|
refreshed, materialization = self.provider.refresh_datasource(
|
|
self.session,
|
|
principal(),
|
|
datasource_ref=cached.ref,
|
|
)
|
|
cached_after = self.provider.read_datasource(
|
|
self.session,
|
|
principal(),
|
|
request=DatasourceReadRequest(datasource_ref=cached.ref),
|
|
)
|
|
|
|
self.assertEqual(2, live_read.total_rows)
|
|
self.assertEqual(1, cached_before.total_rows)
|
|
self.assertEqual(2, cached_live.total_rows)
|
|
self.assertEqual(2, cached_after.total_rows)
|
|
self.assertEqual(materialization.ref, refreshed.current_materialization_ref)
|
|
|
|
def test_tenant_and_scope_isolation(self) -> None:
|
|
stage = self.provider.create_stage(
|
|
self.session,
|
|
principal(),
|
|
stage=DatasourceStageInput(
|
|
name="Private",
|
|
source_name="private",
|
|
kind="upload",
|
|
mode="static",
|
|
shape="tabular",
|
|
rows=({"id": 1},),
|
|
),
|
|
)
|
|
datasource, _ = self.provider.promote_stage(
|
|
self.session,
|
|
principal(),
|
|
stage_ref=stage.ref,
|
|
)
|
|
self.session.commit()
|
|
|
|
self.assertEqual(
|
|
(),
|
|
self.provider.list_datasources(self.session, principal("tenant-2")),
|
|
)
|
|
self.assertIsNone(
|
|
self.provider.get_datasource(
|
|
self.session,
|
|
principal("tenant-2"),
|
|
datasource_ref=datasource.ref,
|
|
)
|
|
)
|
|
with self.assertRaises(DatasourceAccessError):
|
|
self.provider.list_datasources(
|
|
self.session,
|
|
principal(scopes=()),
|
|
)
|
|
|
|
def test_stage_writer_can_list_stages_without_catalogue_scope(self) -> None:
|
|
writer = principal(scopes=(STAGE_WRITE_SCOPE,))
|
|
stage = self.provider.create_stage(
|
|
self.session,
|
|
writer,
|
|
stage=DatasourceStageInput(
|
|
name="Pending import",
|
|
source_name="pending_import",
|
|
kind="upload",
|
|
mode="static",
|
|
shape="tabular",
|
|
rows=({"id": 1},),
|
|
),
|
|
)
|
|
|
|
self.assertEqual(
|
|
(stage.ref,),
|
|
tuple(item.ref for item in self.provider.list_stages(self.session, writer)),
|
|
)
|
|
|
|
def test_producer_publication_is_atomic_idempotent_and_addressable(self) -> None:
|
|
producer = principal(scopes=(SOURCE_WRITE_SCOPE,))
|
|
request = DatasourcePublicationRequest(
|
|
producer_module="dataflow",
|
|
producer_run_ref="dataflow-run:run-1",
|
|
idempotency_key="monthly-output-2026-07",
|
|
name="Monthly comparison result",
|
|
source_name="monthly_comparison_result",
|
|
rows=(
|
|
{"case_id": "A-1", "result": "match"},
|
|
{"case_id": "A-2", "result": "review"},
|
|
),
|
|
freeze=True,
|
|
frozen_label="July 2026",
|
|
provenance={"pipeline_revision": 3},
|
|
)
|
|
|
|
first = self.provider.publish_rows(
|
|
self.session,
|
|
producer,
|
|
request=request,
|
|
)
|
|
replay = self.provider.publish_rows(
|
|
self.session,
|
|
producer,
|
|
request=request,
|
|
)
|
|
second = self.provider.publish_rows(
|
|
self.session,
|
|
producer,
|
|
request=DatasourcePublicationRequest(
|
|
producer_module="dataflow",
|
|
producer_run_ref="dataflow-run:run-2",
|
|
idempotency_key="monthly-output-2026-08",
|
|
target_datasource_ref=first.datasource.ref,
|
|
rows=({"case_id": "A-3", "result": "match"},),
|
|
),
|
|
)
|
|
self.session.commit()
|
|
|
|
self.assertFalse(first.replayed)
|
|
self.assertTrue(replay.replayed)
|
|
self.assertEqual(first.ref, replay.ref)
|
|
self.assertEqual(first.materialization.ref, replay.materialization.ref)
|
|
self.assertIsNotNone(first.materialization.frozen_at)
|
|
self.assertEqual(1, first.materialization.revision)
|
|
self.assertEqual(2, second.materialization.revision)
|
|
self.assertEqual(first.datasource.ref, second.datasource.ref)
|
|
self.assertEqual(
|
|
2,
|
|
self.session.query(DatasourcePublicationRecord).count(),
|
|
)
|
|
|
|
def test_publication_idempotency_key_rejects_different_output(self) -> None:
|
|
producer = principal(scopes=(SOURCE_WRITE_SCOPE,))
|
|
base = DatasourcePublicationRequest(
|
|
producer_module="dataflow",
|
|
producer_run_ref="dataflow-run:run-1",
|
|
idempotency_key="stable-key",
|
|
name="Result",
|
|
source_name="result",
|
|
rows=({"id": 1},),
|
|
)
|
|
self.provider.publish_rows(self.session, producer, request=base)
|
|
|
|
with self.assertRaises(DatasourceValidationError):
|
|
self.provider.publish_rows(
|
|
self.session,
|
|
producer,
|
|
request=DatasourcePublicationRequest(
|
|
producer_module="dataflow",
|
|
producer_run_ref="dataflow-run:run-1",
|
|
idempotency_key="stable-key",
|
|
name="Result",
|
|
source_name="result",
|
|
rows=({"id": 2},),
|
|
),
|
|
)
|
|
|
|
def test_publication_requires_source_write_scope(self) -> None:
|
|
with self.assertRaises(DatasourceAccessError):
|
|
self.provider.publish_rows(
|
|
self.session,
|
|
principal(scopes=()),
|
|
request=DatasourcePublicationRequest(
|
|
producer_module="dataflow",
|
|
producer_run_ref="dataflow-run:run-1",
|
|
idempotency_key="denied",
|
|
name="Result",
|
|
source_name="result",
|
|
rows=({"id": 1},),
|
|
),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|