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

"""Isolated PostgreSQL stage-one migration/restore acceptance; fixed test URL only."""

import asyncio
import os
import subprocess
from pathlib import Path

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:research-test-only@127.0.0.1:18436/wq_research_stage1_test"
os.environ.update(
    DATABASE_URL=URL, ADMIN_PASSWORD="research-acceptance-only", ENCRYPTION_KEY=Fernet.generate_key().decode()
)


async def sql(statement):
    engine = create_async_engine(URL)
    async with engine.begin() as db:
        result = await db.execute(text(statement))
        rows = result.fetchall() if result.returns_rows else None
    await engine.dispose()
    return rows


async def acceptance():
    import httpx

    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.research_metadata_fake import response as metadata_response
    from tests.test_backtests import execute, setup, start
    from tests.test_catalog import prepare, sync
    from tests.test_research_workspace import expansion, template

    def upstream(request):
        if request.url.path == "/authentication":
            return httpx.Response(201, json={})
        if request.url.path == "/users/self":
            return httpx.Response(200, json={"id": "PG_RESEARCH_USER"})
        return metadata_response(request) or 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), base_url="http://testserver", headers={"X-WQ-Request": "1"}
        ) as client:
            assert (
                await client.post(
                    "/api/v1/auth/login", json={"username": "admin", "password": "research-acceptance-only"}
                )
            ).status_code == 200
            await client.put(
                "/api/v1/account/credentials",
                json={"email": "test@example.com", "password": "synthetic-only"},
            )
            job = (await client.post("/api/v1/account/connect")).json()
            await app.state.runner.execute(job["id"])
            catalog = (client, app.state.runner, {})
            await sync(catalog)
            version = (await sync(catalog, "TEST_FIN"))["id"]
            fixed = (await prepare(client, version)).json()
            assert (await client.post("/api/v1/catalog/operators/refresh")).status_code == 200
            assert (await client.post("/api/v1/catalog/setting-options/refresh")).status_code == 200
            saved = (
                await client.post("/api/v1/research/assets", json={"kind": "template", "content": template()})
            ).json()
            a, b = await asyncio.gather(
                *(
                    client.put(
                        "/api/v1/research/assets/" + saved["id"],
                        json={"kind": "template", "version": 1, "content": {**template(), "name": name}},
                    )
                    for name in ["A", "B"]
                )
            )
            assert sorted([a.status_code, b.status_code]) == [200, 409]
            old = (await client.get("/api/v1/research/assets/" + saved["id"] + "?version=1")).json()
            assert old["name"] == template()["name"]
            body = expansion(fixed["id"], asset_id=saved["id"], version=1)
            body.pop("template")
            experiment = (await client.post("/api/v1/research/experiments", json=body)).json()
            assert len(experiment["candidates"]) == 2
            fake, lane = await setup(app)
            preview = (
                await client.post("/api/v1/research/experiments/" + experiment["id"] + "/preview", json={})
            ).json()
            first, second = await asyncio.gather(
                start(client, preview, "research-confirm"), start(client, preview, "research-confirm")
            )
            assert first["backtest_run_id"] == second["backtest_run_id"]
            await execute(app, lane, first["backtest_run_id"])
            current = (await client.get("/api/v1/research/experiments/" + experiment["id"])).json()
            assert current["backtest_run_ids"] == [first["backtest_run_id"]]
    print(
        "PASS PostgreSQL: versions, concurrent CAS, fixed input → experiment → idempotent backtest → provenance"
    )


if __name__ == "__main__":
    config = Config("alembic.ini")
    if asyncio.run(sql("SELECT tablename FROM pg_tables WHERE schemaname='public'")):
        raise RuntimeError("Dedicated acceptance database must be empty")
    command.upgrade(config, "0005")
    asyncio.run(
        sql(
            "INSERT INTO alphas (id,hidden,settings,is_metrics,os_metrics,checks,synced_at,raw) VALUES ('OLD_RESEARCH',false,'{}','{}','{}','[]',now(),'{}')"
        )
    )
    asyncio.run(
        sql(
            "INSERT INTO research (alpha_id,note,tags,favorite,state,updated_at,version) VALUES ('OLD_RESEARCH','preserve note','[]',false,'inbox',now(),7)"
        )
    )
    command.upgrade(config, "head")
    command.check(config)
    assert asyncio.run(sql("SELECT note,version FROM research WHERE alpha_id='OLD_RESEARCH'")) == [
        ("preserve note", 7)
    ]
    asyncio.run(acceptance())
    dump = Path("/tmp/wq-research-stage1.dump")
    with dump.open("wb") as output:
        subprocess.run(
            [
                "docker",
                "exec",
                "wq-research-acceptance-pg",
                "pg_dump",
                "-U",
                "postgres",
                "-Fc",
                "wq_research_stage1_test",
            ],
            stdout=output,
            check=True,
        )
    subprocess.run(
        [
            "docker",
            "exec",
            "wq-research-acceptance-pg",
            "createdb",
            "-U",
            "postgres",
            "wq_research_restore_stage1",
        ],
        check=True,
    )
    with dump.open("rb") as input_file:
        subprocess.run(
            [
                "docker",
                "exec",
                "-i",
                "wq-research-acceptance-pg",
                "pg_restore",
                "-U",
                "postgres",
                "-d",
                "wq_research_restore_stage1",
            ],
            stdin=input_file,
            check=True,
        )
    query = "SELECT (SELECT count(*) FROM research_revisions),(SELECT count(*) FROM research_experiments),(SELECT count(*) FROM backtest_runs),(SELECT note FROM research WHERE alpha_id='OLD_RESEARCH')"
    original = subprocess.check_output(
        [
            "docker",
            "exec",
            "wq-research-acceptance-pg",
            "psql",
            "-U",
            "postgres",
            "-d",
            "wq_research_stage1_test",
            "-Atc",
            query,
        ]
    )
    restored = subprocess.check_output(
        [
            "docker",
            "exec",
            "wq-research-acceptance-pg",
            "psql",
            "-U",
            "postgres",
            "-d",
            "wq_research_restore_stage1",
            "-Atc",
            query,
        ]
    )
    assert original == restored
    print(
        "PASS PostgreSQL 17: 0005 → 0006, schema check, old notes preserved, pg_dump/pg_restore artifacts and provenance counts match"
    )