worldquant-alpha-system/backend/tests/preparations_postgres.py

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