AI问数二期

2026-07-16 13:55:21 +08:00
parent 7be72a50ad
commit a84173066e
10 changed files with 1731 additions and 302 deletions
@@ -0,0 +1,314 @@
+import asyncio
+import hashlib
+import json
+import logging
+from datetime import datetime, timezone
+from typing import Any
+from uuid import uuid4
+
+import httpx
+from fastapi import HTTPException
+
+from config.cubeReport import get_cube_report_settings
+from config.redis import redis_client
+
+CUBE_METADATA_CACHE_KEY = "cube-report:metadata:v1"
+CUBE_METADATA_LOCK_KEY = "cube-report:metadata:refresh-lock:v1"
+_TIME_GRANULARITIES = {
+    "day",
+    "hour",
+    "minute",
+    "month",
+    "quarter",
+    "second",
+    "week",
+    "year",
+}
+
+logger = logging.getLogger(__name__)
+_refresh_lock = asyncio.Lock()
+_background_tasks: set[asyncio.Task[Any]] = set()
+
+
+def _utc_now() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+def _member_name(cube_name: str, value: Any) -> str:
+    name = str(value or "").strip()
+    if not name:
+        return ""
+    if "." not in name and cube_name:
+        return f"{cube_name}.{name}"
+    return name
+
+
+def _member_title(member: dict[str, Any], member_name: str) -> str:
+    for key in ("shortTitle", "short_title", "title"):
+        title = str(member.get(key) or "").strip()
+        if title:
+            return title
+    return member_name
+
+
+def normalize_cube_metadata(payload: Any) -> dict[str, Any]:
+    """将 Cube /meta 响应压缩为成员名到展示标题的映射。"""
+    if not isinstance(payload, dict):
+        raise ValueError("Cube 元数据响应不是 JSON 对象")
+    cubes = payload.get("cubes")
+    if not isinstance(cubes, list):
+        raise ValueError("Cube 元数据缺少 cubes 数组")
+
+    members: dict[str, str] = {}
+    cube_count = 0
+    for cube in cubes:
+        if not isinstance(cube, dict):
+            continue
+        cube_name = str(cube.get("name") or "").strip()
+        if not cube_name:
+            continue
+        cube_count += 1
+        for collection_name in ("dimensions", "measures", "segments"):
+            collection = cube.get(collection_name)
+            if not isinstance(collection, list):
+                continue
+            for member in collection:
+                if not isinstance(member, dict):
+                    continue
+                name = _member_name(cube_name, member.get("name"))
+                if name:
+                    members[name] = _member_title(member, name)
+
+    if not members:
+        raise ValueError("Cube 元数据中没有可用成员")
+
+    schema_json = json.dumps(
+        members,
+        ensure_ascii=False,
+        separators=(",", ":"),
+        sort_keys=True,
+    )
+    return {
+        "version": 1,
+        "members": members,
+        "cubeCount": cube_count,
+        "memberCount": len(members),
+        "refreshedAt": _utc_now().isoformat(),
+        "schemaHash": hashlib.sha256(schema_json.encode("utf-8")).hexdigest(),
+    }
+
+
+async def _read_cache() -> dict[str, Any] | None:
+    try:
+        value = await asyncio.to_thread(
+            redis_client.get_value, CUBE_METADATA_CACHE_KEY
+        )
+    except Exception as error:
+        logger.warning("读取 Cube 元数据缓存失败：%s", error)
+        return None
+    if not isinstance(value, dict) or not isinstance(value.get("members"), dict):
+        return None
+    return value
+
+
+async def _write_cache(value: dict[str, Any]) -> None:
+    await asyncio.to_thread(
+        redis_client.set_value,
+        CUBE_METADATA_CACHE_KEY,
+        value,
+    )
+
+
+async def _cache_is_refreshing() -> bool:
+    try:
+        return bool(
+            await asyncio.to_thread(
+                redis_client.redis.exists, CUBE_METADATA_LOCK_KEY
+            )
+        )
+    except Exception:
+        return _refresh_lock.locked()
+
+
+def _is_stale(cache: dict[str, Any]) -> bool:
+    value = cache.get("refreshedAt")
+    try:
+        refreshed_at = datetime.fromisoformat(str(value).replace("Z", "+00:00"))
+        if refreshed_at.tzinfo is None:
+            refreshed_at = refreshed_at.replace(tzinfo=timezone.utc)
+    except (TypeError, ValueError):
+        return True
+    max_age = max(get_cube_report_settings().metadata_refresh_seconds, 60)
+    return (_utc_now() - refreshed_at).total_seconds() >= max_age
+
+
+def _status(
+    cache: dict[str, Any] | None,
+    *,
+    changed: bool | None = None,
+    refreshing: bool = False,
+) -> dict[str, Any]:
+    result: dict[str, Any] = {
+        "available": bool(cache),
+        "cubeCount": int(cache.get("cubeCount") or 0) if cache else 0,
+        "memberCount": int(cache.get("memberCount") or 0) if cache else 0,
+        "refreshedAt": cache.get("refreshedAt") if cache else None,
+        "refreshing": refreshing,
+        "stale": _is_stale(cache) if cache else True,
+    }
+    if changed is not None:
+        result["changed"] = changed
+    return result
+
+
+async def _fetch_cube_metadata() -> dict[str, Any]:
+    settings = get_cube_report_settings()
+    headers = (
+        {"Authorization": f"Bearer {settings.cube_api_token}"}
+        if settings.cube_api_token
+        else {}
+    )
+    timeout = httpx.Timeout(settings.request_timeout_seconds, connect=15.0)
+    async with httpx.AsyncClient(timeout=timeout) as client:
+        response = await client.get(
+            f"{settings.cube_api_base}/meta",
+            headers=headers,
+        )
+    if response.is_error:
+        raise HTTPException(
+            status_code=502,
+            detail=f"Cube 元数据获取失败：{response.text[:500]}",
+        )
+    try:
+        return normalize_cube_metadata(response.json())
+    except (TypeError, ValueError) as error:
+        raise HTTPException(status_code=502, detail=str(error)) from error
+
+
+async def _acquire_distributed_lock(token: str) -> bool:
+    settings = get_cube_report_settings()
+    expires = max(int(settings.request_timeout_seconds) + 30, 60)
+    try:
+        return bool(
+            await asyncio.to_thread(
+                redis_client.redis.set,
+                CUBE_METADATA_LOCK_KEY,
+                token,
+                nx=True,
+                ex=expires,
+            )
+        )
+    except Exception as error:
+        logger.warning("Cube 元数据刷新锁不可用，将使用进程内锁：%s", error)
+        return True
+
+
+async def _release_distributed_lock(token: str) -> None:
+    script = (
+        "if redis.call('get', KEYS[1]) == ARGV[1] then "
+        "return redis.call('del', KEYS[1]) else return 0 end"
+    )
+    try:
+        await asyncio.to_thread(
+            redis_client.redis.eval,
+            script,
+            1,
+            CUBE_METADATA_LOCK_KEY,
+            token,
+        )
+    except Exception as error:
+        logger.warning("释放 Cube 元数据刷新锁失败：%s", error)
+
+
+async def refresh_cube_metadata_cache() -> dict[str, Any]:
+    """从 Cube 重新加载元数据，并以单个 Redis 值原子替换旧缓存。"""
+    async with _refresh_lock:
+        token = uuid4().hex
+        acquired = await _acquire_distributed_lock(token)
+        if not acquired:
+            return _status(
+                await _read_cache(),
+                changed=False,
+                refreshing=True,
+            )
+        try:
+            previous = await _read_cache()
+            current = await _fetch_cube_metadata()
+            await _write_cache(current)
+            return _status(
+                current,
+                changed=(
+                    not previous
+                    or previous.get("schemaHash") != current.get("schemaHash")
+                ),
+            )
+        finally:
+            await _release_distributed_lock(token)
+
+
+async def get_cube_metadata_status() -> dict[str, Any]:
+    return _status(
+        await _read_cache(),
+        refreshing=await _cache_is_refreshing(),
+    )
+
+
+def _track_background_task(task: asyncio.Task[Any]) -> None:
+    _background_tasks.add(task)
+    task.add_done_callback(_background_tasks.discard)
+
+
+async def _background_refresh() -> None:
+    try:
+        await refresh_cube_metadata_cache()
+    except Exception as error:
+        logger.warning("后台刷新 Cube 元数据失败：%s", error)
+
+
+def _schedule_background_refresh() -> None:
+    if any(not task.done() for task in _background_tasks):
+        return
+    _track_background_task(asyncio.create_task(_background_refresh()))
+
+
+async def get_cube_member_titles() -> dict[str, str]:
+    cache = await _read_cache()
+    if not cache:
+        try:
+            await refresh_cube_metadata_cache()
+            cache = await _read_cache()
+        except Exception as error:
+            logger.warning("Cube 元数据缓存回源失败：%s", error)
+            return {}
+    elif _is_stale(cache):
+        _schedule_background_refresh()
+
+    members = cache.get("members") if cache else None
+    if not isinstance(members, dict):
+        return {}
+    return {
+        str(key): str(value)
+        for key, value in members.items()
+        if key and value
+    }
+
+
+def cube_member_candidates(member: str, column: str) -> list[str]:
+    """生成表头查找候选项，兼容 SQL 别名和时间粒度成员。"""
+    candidates: list[str] = []
+    for value in (member, member.replace("__", "."), column.replace("__", ".")):
+        if value and value not in candidates:
+            candidates.append(value)
+        parts = value.rsplit(".", 1)
+        if len(parts) == 2 and parts[1].lower() in _TIME_GRANULARITIES:
+            if parts[0] not in candidates:
+                candidates.append(parts[0])
+    return candidates
+
+
+async def warm_cube_metadata_cache() -> None:
+    """应用启动时预热；失败只记录日志，不影响主服务可用性。"""
+    try:
+        await refresh_cube_metadata_cache()
+    except Exception as error:
+        logger.warning("启动预热 Cube 元数据失败：%s", error)