govoplan-files/tests/test_source_identity_index.py

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()