refactor: unify data preparations and research input snapshots

2026-09-12 01:24:02 +08:00
parent 849f86fef7
commit 394438e753
82 changed files with 4146 additions and 2076 deletions
@@ -0,0 +1,139 @@
+"""Isolated PostgreSQL migration and concurrency acceptance, using synthetic upstream only."""
+
+import asyncio
+import os
+
+from alembic import command
+from alembic.config import Config
+from cryptography.fernet import Fernet
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import create_async_engine
+
+URL = "postgresql+asyncpg://postgres:preparations-test-only@127.0.0.1:18437/preparations_test"
+os.environ.update(
+    DATABASE_URL=URL,
+    ADMIN_PASSWORD="migration-test-only",
+    ENCRYPTION_KEY=Fernet.generate_key().decode(),
+    WQ_EMAIL="",
+    WQ_PASSWORD="",
+)
+
+
+async def sql(query):
+    engine = create_async_engine(URL)
+    try:
+        async with engine.begin() as db:
+            result = await db.execute(text(query))
+            return result.fetchall() if result.returns_rows else None
+    finally:
+        await engine.dispose()
+
+
+async def acceptance():
+    import httpx
+
+    from app.catalog.contracts import CatalogJobInput, Scope
+    from app.catalog.service import Catalog
+    from app.config import Settings
+    from app.main import create_app
+    from app.worldquant import WqClient
+    from tests.catalog_fake import catalog_response
+    from tests.test_catalog import SCOPE, prepare, search, sync
+
+    def upstream(request):
+        if request.url.path == "/authentication":
+            return httpx.Response(201, json={"token": {"expiry": 14400}})
+        if request.url.path == "/users/self":
+            return httpx.Response(200, json={"id": "PG_TEST_USER"})
+        from tests.research_metadata_fake import response
+
+        metadata = response(request)
+        if metadata is not None:
+            return metadata
+        assert request.method == "GET"
+        return catalog_response(request) or httpx.Response(404)
+
+    settings = Settings(_env_file=None, enable_runner=False, public_origin="http://testserver")
+    app = create_app(settings, WqClient(settings, transport=httpx.MockTransport(upstream)))
+    async with app.router.lifespan_context(app):
+        async with httpx.AsyncClient(
+            transport=httpx.ASGITransport(app=app),
+            base_url="http://testserver",
+            headers={"X-WQ-Request": "1"},
+        ) as client:
+            assert (
+                await client.post(
+                    "/api/v1/auth/login", json={"username": "admin", "password": "migration-test-only"}
+                )
+            ).status_code == 200
+            await client.put(
+                "/api/v1/account/credentials",
+                json={"email": "synthetic@example.com", "password": "synthetic-only"},
+            )
+            job = (await client.post("/api/v1/account/connect")).json()
+            await app.state.runner.execute(job["id"])
+            fixture = (client, app.state.runner, {})
+            await sync(fixture)
+            await sync(fixture, "TEST_FIN")
+            snapshot = (
+                await prepare(
+                    client, (await search(client, "/datasets/TEST_FIN/fields"))["collection_version"]
+                )
+            ).json()
+            assert len(snapshot["field_ids"]) == 123
+            fields = await client.get("/api/v1/catalog/fields", params={**SCOPE, "category": "基本面", "limit": 2, "offset": 2})
+            assert fields.status_code == 200, fields.text
+            assert fields.json()["total"] == 123 and len(fields.json()["items"]) == 2
+            assert fields.json()["items"][0]["category"] == "基本面"
+
+
+            async def enqueue():
+                async with app.state.sessions.begin() as db:
+                    return (await Catalog(db).create_job(CatalogJobInput(scope=Scope(**SCOPE)), full=True)).id
+
+            jobs = await asyncio.gather(*(enqueue() for _ in range(5)))
+            assert len(set(jobs)) == 1, jobs
+            collection = (await client.get("/api/v1/data-preparations")).json()["items"][0]
+            ref = {"id": collection["id"], "version": collection["version"]}
+
+            async def freeze():
+                response = await client.post("/api/v1/data-preparations/freeze", json={"items": [ref]})
+                assert response.status_code == 201, response.text
+                return response.json()["items"][0]["id"]
+
+            assert len(set(await asyncio.gather(*(freeze() for _ in range(5))))) == 1
+            await client.delete(f"/api/v1/data-preparations/{ref['id']}?version={ref['version']}")
+            assert (await client.get(f"/api/v1/research/input-snapshots/{snapshot['id']}")).json() == snapshot
+
+
+if __name__ == "__main__":
+    assert not asyncio.run(sql("SELECT tablename FROM pg_tables WHERE schemaname='public'")), (
+        "Requires an empty isolated database"
+    )
+    config = Config("alembic.ini")
+    command.upgrade(config, "0014")
+    # Existing catalog checkpoints survive the structural change; no research data migration.
+    asyncio.run(
+        sql(
+            "INSERT INTO sync_jobs(id,kind,status,payload,checkpoint,total,processed,failed,cancel_requested,created_at,updated_at) VALUES ('checkpoint-test','catalog_sync','failed','{}','{\"offset\": 100}',0,100,0,false,now(),now())"
+        )
+    )
+    asyncio.run(sql("INSERT INTO catalog_scopes(key,scope) VALUES ('checkpoint-scope','{}')"))
+    asyncio.run(
+        sql(
+            "INSERT INTO catalog_batches(id,scope_key,dataset_id,complete,count) VALUES ('checkpoint-test','checkpoint-scope',NULL,false,100)"
+        )
+    )
+    command.upgrade(config, "head")
+    command.check(config)
+    assert asyncio.run(
+        sql("SELECT job_id, catalog_batches.\"offset\" FROM catalog_batches WHERE id='checkpoint-test'")
+    ) == [("checkpoint-test", 100)]
+    assert asyncio.run(sql("SELECT to_regclass('template_inputs')")) == [(None,)]
+    asyncio.run(sql("DELETE FROM catalog_batches WHERE id='checkpoint-test'"))
+    asyncio.run(sql("DELETE FROM catalog_scopes WHERE key='checkpoint-scope'"))
+    asyncio.run(sql("DELETE FROM sync_jobs WHERE id='checkpoint-test'"))
+    asyncio.run(acceptance())
+    print(
+        "PostgreSQL 17: 0014 → 0015 metadata, retained catalog checkpoint, concurrent job deduplication/freeze and independent snapshot passed"
+    )