worldquant-alpha-system/backend/app/catalog/research_metadata.py

"""Bounded metadata reads with all-or-nothing publication and separate annotations."""

import hashlib
import json

from fastapi import HTTPException
from sqlalchemy import select, update
from sqlalchemy.exc import IntegrityError

from ..models import CatalogResource, OperatorNote, now
from ..research.serialization import encode_snapshot as jsonable_encoder
from ..worldquant import WqError


async def upstream(operation):
    try:
        return await operation
    except WqError as exc:
        raise HTTPException(
            409 if exc.code in ("disconnected", "verification_required") else 502, str(exc)
        ) from None


def availability_key(field_id, scope):
    return "field:" + hashlib.sha256(f"{field_id}|{scope.key()}".encode()).hexdigest()


def setting_rows(data):
    """Decode hierarchical OPTIONS choices; incomplete options are not invented."""
    try:
        children = data["actions"]["POST"]["settings"]["children"]

        def choices(key, instrument=None, region=None, delay=None, universe=None):
            value = children[key]["choices"]
            if isinstance(value, dict) and "instrumentType" in value:
                value = value["instrumentType"][instrument]
            elif isinstance(value, dict) and instrument in value:
                value = value[instrument]
            if isinstance(value, dict) and "region" in value:
                value = value["region"][region]
            for dimension, selected in (("delay", delay), ("universe", universe)):
                if isinstance(value, dict) and dimension in value:
                    branches = value[dimension]
                    if not isinstance(branches, dict):
                        raise ValueError()
                    value = branches.get(str(selected), branches.get(selected))
            if not isinstance(value, list):
                raise ValueError()
            return [item["value"] for item in value]

        rows = []
        for instrument in choices("instrumentType"):
            if instrument != "EQUITY":
                continue
            for region in choices("region", instrument):
                for delay in choices("delay", instrument, region):
                    if type(delay) is not int or delay not in (0, 1):
                        continue
                    for universe in choices("universe", instrument, region, delay):
                        neutralizations = (
                            choices("neutralization", instrument, region, delay, universe)
                            if "neutralization" in children
                            else []
                        )
                        fields = {}
                        for key in (
                            "decay", "truncation", "pasteurization", "unitHandling",
                            "nanHandling", "language", "visualization", "maxTrade", "maxPosition",
                        ):
                            definition = children.get(key)
                            if not isinstance(definition, dict):
                                continue
                            field = {}
                            if "choices" in definition:
                                field["choices"] = choices(key, instrument, region, delay, universe)
                                if any(type(v) not in (str, int, float, bool) for v in field["choices"]):
                                    raise ValueError()
                            for source, target in (("min_value", "minimum"), ("max_value", "maximum")):
                                if type(definition.get(source)) in (int, float):
                                    field[target] = definition[source]
                            fields[key] = field
                        rows.append(
                            {
                                "instrument_type": instrument,
                                "region": region,
                                "universe": universe,
                                "delay": delay,
                                "neutralizations": neutralizations,
                                "fields": fields,
                            }
                        )
        if not rows:
            raise ValueError()
        return rows
    except (KeyError, TypeError, ValueError):
        raise HTTPException(502, "平台设置结构无法识别，未发布新快照") from None


def normalize_availability(data, *, instrument_type=None):
    """Use the request's instrument only for the platform field-detail `data` form.

    Legacy availability rows must still state their own instrument. Missing market,
    delay or universe never inherits the requested scope.
    """
    raw = data.get("availability")
    detail_form = raw is None and isinstance(data.get("data"), list)
    if detail_form:
        raw = data["data"]
    if not isinstance(raw, list):
        return {"status": "needs_review", "items": [], "reason": "平台未提供可识别的 availability 列表"}
    rows, malformed = [], False
    for item in raw:
        if not isinstance(item, dict):
            malformed = True
            continue
        universes = item.get("universe", item.get("universes"))
        universes = universes if isinstance(universes, list) else [universes]
        for universe in universes:
            if (
                item.get("instrumentType", instrument_type if detail_form else None) == "EQUITY"
                and type(item.get("delay")) is int
                and item["delay"] in (0, 1)
                and isinstance(item.get("region"), str)
                and isinstance(universe, str)
            ):
                rows.append(
                    {
                        "instrument_type": "EQUITY",
                        "region": item["region"],
                        "delay": item["delay"],
                        "universe": universe,
                    }
                )
            else:
                malformed = True
    return {
        "status": "available" if rows and not malformed else "needs_review",
        "items": rows,
        "reason": "可用性列表包含不完整项" if malformed else "",
    }


class ResearchMetadata:
    def __init__(self, db, client=None):
        self.db, self.client = db, client

    async def publish(self, key, kind, content):
        row = await self.db.get(CatalogResource, key)
        if row:
            row.content, row.fetched_at = content, now()
        else:
            row = CatalogResource(key=key, kind=kind, content=content)
            self.db.add(row)
        await self.db.flush()
        return self.output(row)

    @staticmethod
    def output(row):
        return jsonable_encoder({"key": row.key, "content": row.content, "fetched_at": row.fetched_at})

    async def get(self, key):
        row = await self.db.get(CatalogResource, key)
        return self.output(row) if row else {"key": key, "content": {}, "fetched_at": None}

    async def refresh_operators(self):
        items, seen = [], set()
        for offset in range(0, 10000, 100):
            page = await upstream(self.client.operators(offset))
            values = page if isinstance(page, list) else page.get("results")
            if not isinstance(values, list):
                raise HTTPException(502, "算子目录格式无法识别，保留原快照")
            for item in values:
                if not isinstance(item, dict) or not isinstance(item.get("name"), str):
                    raise HTTPException(502, "算子目录缺少名称，保留原快照")
                if item["name"] in seen:
                    raise HTTPException(502, "算子分页重复，未发布不完整目录")
                seen.add(item["name"])
                items.append(
                    {
                        key: item.get(key)
                        for key in (
                            "name",
                            "category",
                            "description",
                            "definition",
                            "example",
                            "scope",
                            "type",
                            "parameters",
                        )
                    }
                )
            if (
                isinstance(page, list)
                or (isinstance(page.get("count"), int) and offset + len(values) >= page["count"])
                or (not page.get("next") and len(values) < 100)
            ):
                return await self.publish("operators", "operators", {"items": items})
            if not values:
                raise HTTPException(502, "算子分页提前结束")
        raise HTTPException(502, "算子分页超过本地限制，未发布新快照")

    async def operators(self, q="", category=None, favorite=False, limit=25, offset=0):
        snapshot = await self.get("operators")
        notes = {r.name: r for r in await self.db.scalars(select(OperatorNote))}
        rows = []
        for item in snapshot["content"].get("items", []):
            note = notes.get(item["name"])
            if q.lower() not in json.dumps(item, ensure_ascii=False).lower() or (
                category and item["category"] != category
            ):
                continue
            if favorite and not (note and note.favorite):
                continue
            rows.append(
                {
                    **item,
                    "local": {
                        "note": note.note if note else "",
                        "favorite": note.favorite if note else False,
                        "version": note.version if note else 0,
                    },
                }
            )
        return {
            "items": rows[offset : offset + limit],
            "total": len(rows),
            "limit": limit,
            "offset": offset,
            "fetched_at": snapshot["fetched_at"],
            "categories": sorted(
                {str(i.get("category")) for i in snapshot["content"].get("items", []) if i.get("category")}
            ),
        }

    async def annotate(self, name, body):
        snapshot = await self.get("operators")
        if name not in {i["name"] for i in snapshot["content"].get("items", [])}:
            raise HTTPException(404, "算子不在已同步目录中")
        if body.version == 0:
            if await self.db.get(OperatorNote, name):
                raise HTTPException(409, "备注已变化")
            self.db.add(OperatorNote(name=name, note=body.note, favorite=body.favorite))
            try:
                await self.db.flush()
            except IntegrityError:
                raise HTTPException(409, "备注已变化，请保留草稿并刷新") from None
        else:
            result = await self.db.execute(
                update(OperatorNote)
                .where(OperatorNote.name == name, OperatorNote.version == body.version)
                .values(note=body.note, favorite=body.favorite, version=body.version + 1)
            )
            if result.rowcount != 1:
                raise HTTPException(409, "备注已变化，请保留草稿并刷新")
        return {"ok": True, "version": body.version + 1}

    async def refresh_settings(self):
        data = await upstream(self.client.research_setting_options())
        return await self.publish("settings", "settings", {"items": setting_rows(data)})

    async def refresh_availability(self, body):
        data = await upstream(self.client.field_availability(body.field_id, body.scope))
        if data.get("id") is not None and data["id"] != body.field_id:
            raise HTTPException(502, "平台返回字段与请求不一致，保留原可用性快照")
        content = {
            **normalize_availability(data, instrument_type=body.scope.instrument_type),
            "field_id": body.field_id,
            "scope": body.scope.model_dump(),
        }
        return await self.publish(availability_key(body.field_id, body.scope), "availability", content)