vent_agent_2.0/extensions/main_fan/gateway_queries.py

"""网关查询共用的窗口校验、服务端身份和可信文件归档。"""
import asyncio
from datetime import datetime, timedelta

from pydantic import BaseModel, ConfigDict, Field

import config
from algorithms.main_fan.analysis import LOCAL_ZONE, local_time
from db import db_conn
from extensions import get_service
from extensions.main_fan import exploration
from extensions.registry import load_extensions
from runtime.python_sandbox import digest, workspace_file


class CurveQuery(BaseModel):
    model_config = ConfigDict(extra='forbid', str_strip_whitespace=True)
    node_id: str = Field(min_length=1, max_length=128, pattern=r'^[A-Za-z0-9_-]+$')
    dept_id: str = Field(min_length=1, max_length=64, pattern=r'^[A-Za-z0-9_-]+$')
    start_time: str = Field(min_length=1, max_length=64)
    end_time: str = Field(min_length=1, max_length=64)


def parse_time(value):
    if not isinstance(value, str) or not ('T' in value or ' ' in value):
        raise ValueError('时间须包含日期和时分秒，例如 2026-10-19 11:00:00 或带时区的 ISO 时间')
    try:
        return local_time(datetime.fromisoformat(value))
    except ValueError as exc:
        raise ValueError('无法解析时间，请使用日期和时分秒，可附带时区') from exc


def curve_parameters(node_id, dept_id, start_time, end_time, *, node_key='nodeIds'):
    query = CurveQuery(node_id=node_id, dept_id=dept_id, start_time=start_time, end_time=end_time)
    start, end = parse_time(query.start_time), parse_time(query.end_time)
    if start.microsecond or end.microsecond:
        raise ValueError('查询起止时间须精确到秒，不支持小数秒')
    if start >= end:
        raise ValueError('查询结束时间必须晚于开始时间')
    maximum = config.get_int('MAIN_FAN_SERIES_MAX_HOURS', 24)
    if maximum <= 0 or end - start > timedelta(hours=maximum):
        raise ValueError(f'单次查询窗口最多 {maximum} 小时，请分段查询')
    return {node_key: query.node_id, 'deptId': query.dept_id,
            'startTime': start.strftime('%Y-%m-%d %H:%M:%S'),
            'endTime': end.strftime('%Y-%m-%d %H:%M:%S')}, start, end


async def fetch(identifier, user, session, path, params):
    workflow = await asyncio.to_thread(exploration._load, identifier, user, session)
    await asyncio.to_thread(exploration.validate_inputs, workflow, user, session)
    bound_dept = workflow['context']['task_input'].get('dept_id')
    actual_dept = str(dict(params).get('deptId', ''))
    if bound_dept and actual_dept != bound_dept:
        raise ValueError('查询部门不属于本次已绑定煤矿')
    if len(workflow['stages']) >= 128:
        raise ValueError('探索操作已达上限，请先提交已有结果')
    if not config.get_bool('SHANXI_GATEWAY_ENABLED', False):
        raise ValueError('取数需要启用 SHANXI_GATEWAY_ENABLED 并配置网关地址和凭证')
    try:
        gateway = get_service('shaanxi_gateway')
    except RuntimeError as exc:
        raise ValueError('取数网关尚未启动，请在启用网关的 RuntimeEnvironment 中调用') from exc
    spec = load_extensions().business('main_fan')
    body = await gateway.proxy(spec.gateway.agent_code, path, params=params)
    request = {'method': 'GET', 'path': path, 'params': params,
               'agent_code': spec.gateway.agent_code, 'fetched_at': datetime.now(LOCAL_ZONE).isoformat()}
    return workflow, request, body


def archive(workflow, user, session, request, body, data, summary, *, kind, collection, operation, filename):
    exploration.validate_inputs(workflow, user, session)
    query_id = db_conn.new_id()
    prefix = f"/workspace/exploration/{workflow['id']}/queries/{query_id}"
    paths = {'request_path': prefix + '/request.json', 'raw_path': prefix + '/response.json',
             'data_path': prefix + '/' + filename}
    data = {'schema_version': 1, 'query_id': query_id, 'query': request, **data}
    exports = {}
    for key, value in (('request_path', request), ('raw_path', body), ('data_path', data)):
        path = workspace_file(user, session, paths[key])
        exploration._write(path, value)
        exports[paths[key]] = digest(path)
    params = dict(request['params'])
    receipt = {'ok': True, 'query_id': query_id, 'kind': kind, 'dept_id': params['deptId'],
               **summary, **paths,
               'python_paths': {key.removesuffix('_path'): value.removeprefix('/workspace/')
                                for key, value in paths.items()}, 'sha256': exports[paths['data_path']]}
    if 'startTime' in params:
        receipt.update(node_id=params.get('nodeId', params.get('nodeIds')), requested_range={
            'start_time': params['startTime'], 'end_time': params['endTime'], 'timezone': 'Asia/Shanghai'})
    context = workflow['context']
    context.setdefault(collection, []).append(receipt)
    context.setdefault('query_exports', {}).update(exports)
    exploration._persist(workflow, user, session, operation, receipt)
    return receipt