ai-agent-admin/backend-fastapi/online_dev/report_manager/dataset_bridge.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""报表数据集桥接：版本 ↔ core_data_source"""
import asyncio
import logging
from typing import Any, Dict, List, Optional, Tuple

from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession

from app.database import AsyncSessionLocal
from core.data_source.model import DataSource
from core.data_source.service import DataSourceService
from online_dev.report_manager.engine.convert_lookup import ConvertLookupCache
from online_dev.report_manager.engine.dataset_transform import transform_dataset_rows
from online_dev.report_manager.engine.sort_apply import (
    apply_sort_to_rows,
    get_sort_rules_for_alias,
)
from online_dev.report_manager.model import ReportDataset

logger = logging.getLogger(__name__)


class ReportDatasetBridge:
    @staticmethod
    async def list_by_version(db: AsyncSession, version_id: str) -> List[Dict[str, Any]]:
        stmt = (
            select(ReportDataset, DataSource.code, DataSource.name)
            .join(DataSource, DataSource.id == ReportDataset.data_source_id)
            .where(
                ReportDataset.version_id == version_id,
                ReportDataset.is_deleted == False,
            )
            .order_by(ReportDataset.sort)
        )
        rows = await db.execute(stmt)
        result = []
        for ds, code, name in rows:
            result.append({
                "id": ds.id,
                "version_id": ds.version_id,
                "data_source_id": ds.data_source_id,
                "data_source_code": code,
                "data_source_name": name,
                "alias": ds.alias,
                "field_mapping": ds.field_mapping or {},
                "convert_config": ds.convert_config or {},
                "sort": ds.sort or 0,
            })
        return result

    @staticmethod
    async def sync_datasets(
        db: AsyncSession,
        version_id: str,
        data_set_list: List[Dict[str, Any]],
    ) -> None:
        stmt = update(ReportDataset).where(
            ReportDataset.version_id == version_id,
            ReportDataset.is_deleted == False,
        ).values(is_deleted=True)
        await db.execute(stmt)

        for idx, item in enumerate(data_set_list):
            ds_id = item.get("data_source_id") or item.get("dataSourceId")
            if not ds_id:
                continue
            alias = item.get("alias") or item.get("name") or f"ds_{idx}"
            record = ReportDataset(
                version_id=version_id,
                data_source_id=ds_id,
                alias=alias,
                field_mapping=item.get("field_mapping") or item.get("fieldMapping") or {},
                convert_config=item.get("convert_config") or item.get("convertConfig") or {},
                sort=item.get("sort", idx),
            )
            db.add(record)
        await db.flush()

    @staticmethod
    async def delete_by_version(db: AsyncSession, version_id: str) -> None:
        stmt = update(ReportDataset).where(
            ReportDataset.version_id == version_id,
        ).values(is_deleted=True)
        await db.execute(stmt)

    @staticmethod
    async def copy_datasets(db: AsyncSession, src_version_id: str, dst_version_id: str) -> None:
        stmt = select(ReportDataset).where(
            ReportDataset.version_id == src_version_id,
            ReportDataset.is_deleted == False,
        )
        for src in (await db.execute(stmt)).scalars().all():
            db.add(ReportDataset(
                version_id=dst_version_id,
                data_source_id=src.data_source_id,
                alias=src.alias,
                field_mapping=src.field_mapping,
                convert_config=src.convert_config,
                sort=src.sort,
            ))

    @staticmethod
    async def _fetch_one_dataset(
        ds: Dict[str, Any],
        params: Dict[str, Any],
        max_rows: int,
        sort_list: List[Any],
        version_convert: Any,
        lookup: Optional[ConvertLookupCache] = None,
    ) -> Tuple[str, List[Any]]:
        async with AsyncSessionLocal() as session:
            source = await DataSourceService.get_by_id(session, ds["data_source_id"])
            if not source:
                return ds["alias"], []
            data = await DataSourceService.execute_by_id(
                session, ds["data_source_id"], params, max_rows=max_rows
            )
            rows: List[Any]
            if isinstance(data, list):
                rows = data
            elif isinstance(data, dict):
                rows = [data]
            else:
                rows = [{"value": data}] if data is not None else []
            rules = get_sort_rules_for_alias(sort_list or [], ds["alias"])
            if rules and rows and isinstance(rows[0], dict):
                rows = apply_sort_to_rows(rows, rules)
            rows = transform_dataset_rows(
                rows,
                field_mapping=ds.get("field_mapping"),
                dataset_convert=ds.get("convert_config"),
                version_convert=version_convert,
                alias=ds["alias"],
                lookup=lookup,
            )
            return ds["alias"], rows

    @staticmethod
    async def fetch_all(
        db: AsyncSession,
        version_id: str,
        params: Dict[str, Any] = None,
        max_rows: int = 50000,
        sort_list: List[Any] = None,
        version_convert: Any = None,
        lookup: Optional[ConvertLookupCache] = None,
    ) -> Dict[str, List[Any]]:
        """拉取版本关联的全部数据集（含 field_mapping / convertConfig 变换）"""
        datasets = await ReportDatasetBridge.list_by_version(db, version_id)
        params = params or {}
        if not datasets:
            return {}
        if len(datasets) == 1:
            alias, rows = await ReportDatasetBridge._fetch_one_dataset(
                datasets[0], params, max_rows, sort_list, version_convert, lookup
            )
            return {alias: rows}
        pairs = await asyncio.gather(
            *[
                ReportDatasetBridge._fetch_one_dataset(
                    ds, params, max_rows, sort_list, version_convert, lookup
                )
                for ds in datasets
            ]
        )
        return dict(pairs)