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