from datetime import datetime, timezone import importlib import unittest from alembic.migration import MigrationContext from alembic.operations import Operations from sqlalchemy import ( JSON, Column, MetaData, String, Table, create_engine, inspect, select, text, ) from sqlalchemy.orm import Session from govoplan_access.backend.db import models as _access_models # register FK metadata from govoplan_core.core.change_sequence import ChangeSequenceEntry from govoplan_files.backend.db.models import FileAsset from govoplan_files.backend.storage.common import FileSourceConflict from govoplan_files.backend.storage.files import find_asset_by_source from govoplan_files.backend.storage.provenance import source_identity_hash def provenance(external="remote-1"): return {"connector_id": " connector ", "provider": "s3", "external_id": external} class SourceIdentityIndexTests(unittest.TestCase): def setUp(self): self.engine = create_engine("sqlite:///:memory:") for table in ( _access_models.Account.__table__, _access_models.User.__table__, _access_models.Group.__table__, FileAsset.__table__, ChangeSequenceEntry.__table__, ): table.create(self.engine) self.session = Session(self.engine) def tearDown(self): self.session.close() self.engine.dispose() def asset(self, name, *, tenant="tenant", owner="owner", source=None, group=False): row = FileAsset( id=name, tenant_id=tenant, owner_type="group" if group else "user", owner_user_id=None if group else owner, owner_group_id=owner if group else None, display_path=f"{name}.txt", filename=f"{name}.txt", metadata_={"source_provenance": source or provenance()}, ) self.session.add(row) self.session.flush() return row def find(self, source=None, *, tenant="tenant", owner="owner", group=False): return find_asset_by_source( self.session, tenant_id=tenant, owner_type="group" if group else "user", owner_id=owner, source_provenance=source or provenance(), ) def test_index_lookup_preserves_tenant_owner_provider_and_external_identity(self): row = self.asset("wanted") self.asset("other-tenant", tenant="another") self.asset("other-owner", owner="another") self.asset("group", group=True) self.asset("other-provider", source={**provenance(), "provider": "webdav"}) for index in range(100): self.asset(f"unrelated-{index}", source=provenance(f"remote-{index + 2}")) self.assertIs(row, self.find()) self.assertIs(row, self.find()) # retry self.assertEqual("group", self.find(group=True).id) statement = ( select(FileAsset.id) .where( FileAsset.tenant_id == "tenant", FileAsset.owner_type == "user", FileAsset.owner_user_id == "owner", FileAsset.source_identity_hash == source_identity_hash(provenance()), FileAsset.deleted_at.is_(None), ) .limit(2) ) sql = str( statement.compile(self.engine, compile_kwargs={"literal_binds": True}) ) plan = self.session.execute(text("EXPLAIN QUERY PLAN " + sql)).all() self.assertIn("ix_file_assets_user_source", str(plan)) self.assertIn("LIMIT 2", sql) def test_metadata_copy_delete_restore_and_duplicates_never_choose_a_winner(self): row = self.asset("original") copy = self.asset("copy", source=dict(row.metadata_["source_provenance"])) self.assertEqual(row.source_identity_hash, copy.source_identity_hash) with self.assertRaisesRegex(FileSourceConflict, "Multiple active"): self.find() copy.deleted_at = datetime.now(timezone.utc) self.session.flush() self.assertIs(row, self.find()) copy.deleted_at = None self.session.flush() with self.assertRaises(FileSourceConflict): self.find() copy.metadata_ = {"source_provenance": provenance("changed")} self.session.flush() self.assertIs(row, self.find()) self.assertIs(copy, self.find(provenance("changed"))) copy.metadata_ = {} self.session.flush() self.assertIsNone(copy.source_identity_hash) self.assertIsNone(self.find(provenance("changed"))) def test_backfill_retains_all_legacy_records_and_metadata_in_bounded_batches(self): migration = importlib.import_module( "govoplan_files.backend.migrations.versions.a2b3c4d5e701_file_source_identity_index" ) engine = create_engine("sqlite:///:memory:") metadata = MetaData() old = Table( "file_assets", metadata, Column("id", String, primary_key=True), Column("tenant_id", String), Column("owner_type", String), Column("owner_user_id", String), Column("owner_group_id", String), Column("metadata", JSON), ) metadata.create_all(engine) rows = [ { "id": f"{index:04d}", "tenant_id": "tenant", "owner_type": "user", "owner_user_id": "owner", "metadata": { "source_provenance": provenance( "duplicate" if index < 2 else str(index) ) }, } for index in range(503) ] rows.append( { "id": "no-source", "tenant_id": "tenant", "owner_type": "user", "owner_user_id": "owner", "metadata": {"unrelated": [1, 2]}, } ) try: with engine.begin() as connection: connection.execute(old.insert(), rows) with Operations.context(MigrationContext.configure(connection)): migration.upgrade() new = Table("file_assets", MetaData(), autoload_with=connection) actual = ( connection.execute(select(new).order_by(new.c.id)).mappings().all() ) self.assertEqual(len(rows), len(actual)) self.assertEqual( [row["metadata"] for row in rows], [row["metadata"] for row in actual], ) self.assertEqual( actual[0]["source_identity_hash"], actual[1]["source_identity_hash"] ) for row in actual: self.assertEqual( source_identity_hash(row["metadata"].get("source_provenance")), row["source_identity_hash"], ) indexes = inspect(connection).get_indexes("file_assets") self.assertEqual(2, len(indexes)) self.assertFalse(any(index["unique"] for index in indexes)) with Operations.context(MigrationContext.configure(connection)): migration.downgrade() self.assertEqual( len(rows), len(connection.execute(select(old.c.id)).all()) ) finally: engine.dispose() if __name__ == "__main__": unittest.main()