"""Disposable PostgreSQL migration/concurrency acceptance; synthetic platform only. MCP_TEST_DATABASE_URL must identify a local database named wq_mcp_test. """ import asyncio import os from urllib.parse import urlsplit import httpx from alembic import command from alembic.config import Config from cryptography.fernet import Fernet from sqlalchemy import func, select from app.alphas import upsert_alpha from app.config import Settings from app.db import create_database from app.main import create_app from app.mcp_api.auth import SCOPES, authenticate, create_token from app.models import BacktestEvent, BacktestPreview, BacktestRun, Research, ResearchRequest from app.research_access.contracts import Control, Submit from app.research_access.service import ResearchAccess, ResearchError from app.worldquant import WqClient from tests.backtest_fake import Platform from tests.test_backtests import candidate, setup async def seed(settings): engine, sessions = create_database(settings.database_url) async with sessions.begin() as db: await upsert_alpha(db, {"id": "legacy", "regular": {"code": "rank(close)"}}) (await db.get(Research, "legacy")).note = "keep historical note" await engine.dispose() async def verify(settings): app = create_app(settings, WqClient(settings, transport=httpx.MockTransport(Platform()))) async with app.router.lifespan_context(app): await setup(app) async with app.state.sessions.begin() as db: _, secret = await create_token(db, "synthetic", SCOPES) principal = await authenticate(db, secret) args = Submit(name="concurrent", candidates=[candidate()], idempotency_key="same") async def submit(body): async with app.state.sessions.begin() as db: return await ResearchAccess(db, principal, app.state.runner.client, settings.public_origin).submit(body) a, b = await asyncio.gather(submit(args), submit(args)) assert a == b rid = a["backtest_run_id"] async with app.state.sessions() as db: assert await db.scalar(select(func.count()).select_from(BacktestRun)) == 1 assert await db.scalar(select(func.count()).select_from(BacktestPreview)) == 1 assert await db.scalar(select(func.count()).select_from(ResearchRequest)) == 1 assert (await db.get(Research, "legacy")).note == "keep historical note" try: await submit(args.model_copy(update={"name": "different"})) except ResearchError as exc: assert exc.data["code"] == "IDEMPOTENCY_CONFLICT" else: raise AssertionError("different payload accepted") control = Control(run_id=rid, action="pause", expected_version=1, idempotency_key="pause") async def pause(): async with app.state.sessions.begin() as db: return await ResearchAccess(db, principal, app.state.runner.client, settings.public_origin).control(control) x, y = await asyncio.gather(pause(), pause()) assert x == y async with app.state.sessions() as db: assert await db.scalar(select(func.count()).select_from(BacktestEvent).where(BacktestEvent.kind == "control")) == 1 # A fresh process-equivalent application returns the original response after restart. replacement = create_app(settings, WqClient(settings, transport=httpx.MockTransport(Platform()))) async with replacement.router.lifespan_context(replacement): async with replacement.state.sessions.begin() as db: principal = await authenticate(db, secret) replay = await ResearchAccess(db, principal, replacement.state.runner.client, settings.public_origin).submit(args) assert replay == a print("PostgreSQL: migration preserves history; concurrent submit/control and restart replay passed") if __name__ == "__main__": url = os.environ["MCP_TEST_DATABASE_URL"] parsed = urlsplit(url) if parsed.hostname not in {"127.0.0.1", "localhost"} or parsed.path != "/wq_mcp_test": raise SystemExit("Refusing a non-local/non-disposable database") key = Fernet.generate_key().decode() os.environ.update(DATABASE_URL=url, ADMIN_PASSWORD="synthetic-admin-only", ENCRYPTION_KEY=key, WQ_EMAIL="", WQ_PASSWORD="") settings = Settings(_env_file=None, database_url=url, admin_password="synthetic-admin-only", encryption_key=key, enable_runner=False, mcp_enabled=True, public_origin="http://testserver") config = Config("alembic.ini") command.upgrade(config, "0009") asyncio.run(seed(settings)) command.upgrade(config, "head") asyncio.run(verify(settings)) command.downgrade(config, "0009") command.upgrade(config, "head") print("PostgreSQL: additive migration downgrade/re-upgrade passed")