vent_agent_2.0/runtime/test_lab/files.py

"""文件资料库：仅按归属和服务器生成的ID定位文件，不接受客户端磁盘路径。"""
import hashlib
import json
import mimetypes
from pathlib import Path
import re

import config
from db import db_conn, test_lab_store

ALLOWED = {'.pdf','.png','.jpg','.jpeg','.webp','.json','.xlsx','.csv'}
DOCUMENTS = {'.pdf','.png','.jpg','.jpeg','.webp'}
MAX_BYTES = 50*1024*1024


def root():
    path = Path(config.get_str('TEST_LAB_DATA_DIR', 'data/test_lab'))
    return path if path.is_absolute() else config.ROOT/path


def user_root(user):
    return root()/hashlib.sha256(user.encode()).hexdigest()[:24]


def samples():
    configured = config.get_str('TEST_LAB_SAMPLE_DIR')
    directory = Path(configured) if configured else config.ROOT.parent/'陕西局代码/测试数据'
    if not directory.is_absolute():
        directory = config.ROOT/directory
    if not directory.is_dir():
        return {}
    directory = directory.resolve()
    result = {}
    for path in sorted(directory.rglob('*')):
        if len(result) >= 100:
            break
        if path.is_file() and not path.is_symlink() and path.suffix.lower() in ALLOWED and path.resolve().is_relative_to(directory):
            key = 'sample-'+hashlib.sha256(path.relative_to(directory).as_posix().encode()).hexdigest()[:24]
            result[key] = path
    return result


def describe(identifier, filename, size, *, source, digest=''):
    suffix = Path(filename).suffix.lower()
    kind = 'pdf' if suffix=='.pdf' else 'image' if suffix in DOCUMENTS else 'json' if suffix=='.json' else 'excel'
    return {'id':identifier,'filename':filename,'size':size,'kind':kind,'source':source,'sha256':digest,
        'can_analyze':suffix in DOCUMENTS or suffix=='.json',
        'url':f'/api/test-lab/assets/{identifier}/content','preview_url':f'/api/test-lab/assets/{identifier}/preview'}


def list_assets(user):
    builtin = [describe(key,path.name,path.stat().st_size,source='sample') for key,path in samples().items()]
    uploaded = [describe(row['id'],row['filename'],row['size'],source='upload',digest=row['sha256']) for row in test_lab_store.assets(user)]
    return builtin+uploaded


def owned_file(identifier, user):
    if identifier.startswith('sample-'):
        path = samples().get(identifier)
        if path is None:
            raise ValueError('内置测试资料不存在')
        return path, describe(identifier,path.name,path.stat().st_size,source='sample')
    if not re.fullmatch(r'[0-9a-f]{32}', identifier):
        raise ValueError('测试资料不存在')
    row = test_lab_store.asset(identifier,user)
    path = user_root(user)/identifier/row['filename']
    if not path.is_file() or path.is_symlink() or not path.resolve().is_relative_to(user_root(user).resolve()):
        raise ValueError('测试资料文件不存在')
    return path, describe(identifier,row['filename'],row['size'],source='upload',digest=row['sha256'])


async def upload(file, user):
    name = re.sub(r'[^\w\u4e00-\u9fff.\-]+','_',Path((file.filename or '').replace('\\','/')).name)[:180]
    if Path(name).suffix.lower() not in ALLOWED:
        raise ValueError('支持PDF、PNG、JPEG、WebP、JSON、XLSX和CSV文件')
    identifier = db_conn.new_id()
    directory = user_root(user)/identifier
    directory.mkdir(parents=True,exist_ok=False)
    path = directory/name
    size, digest = 0, hashlib.sha256()
    try:
        with path.open('xb') as stream:
            while chunk := await file.read(1024*1024):
                size += len(chunk)
                if size>MAX_BYTES:
                    raise ValueError('单个文件不得超过50MB')
                stream.write(chunk)
                digest.update(chunk)
        if not size:
            raise ValueError('不能上传空文件')
        if path.suffix.lower()=='.json':
            if size>5*1024*1024:
                raise ValueError('JSON测试数据不得超过5MB')
            value = json.loads(path.read_text(encoding='utf-8-sig'))
            if not isinstance(value, dict):
                raise ValueError('JSON测试数据必须是一个对象')
        test_lab_store.save_asset(user,name,digest.hexdigest(),size,identifier)
    except BaseException:
        path.unlink(missing_ok=True)
        directory.rmdir()
        raise
    finally:
        await file.close()
    return describe(identifier,name,size,source='upload',digest=digest.hexdigest())


def media_type(path):
    return {'.pdf':'application/pdf','.json':'application/json','.xlsx':'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
        '.csv':'text/csv; charset=utf-8'}.get(path.suffix.lower(),mimetypes.guess_type(path.name)[0] or 'application/octet-stream')


def json_value(path):
    if path.stat().st_size>5*1024*1024:
        raise ValueError('JSON文件过大')
    return json.loads(path.read_text(encoding='utf-8-sig'))


def series_fragment(raw, source):
    """兼容项目已有的开停/风量mock JSON；设备归属和单位仍须显式给出。"""
    if raw.get('schema_version')!=1 or raw.get('data_kind')!='synthetic' or not isinstance(raw.get('series'),list):
        return None
    from algorithms.main_fan.models import Device, FanStateRecord, OperatingSeries
    common={'mine_code':raw['mine_code'],'mine_name':raw['mine_name'],
        'evidence':{'source':source,'method':'synthetic','label':'上传的测试时序','excerpt':raw.get('description','上传的测试时序')}}
    fragment={'devices':[],'fan_state_records':[],'operating_series':[]}
    for series in raw['series']:
        if series.get('state_scope')=='whole_fan':
            fragment['devices'].append(Device(device_code=series['fan_id'],name=series['fan_name'],**common).model_dump(mode='json'))
            for sample in series['samples']:
                if sample['value'] not in (0,1):
                    raise ValueError('整机开停曲线的值必须为0或1')
                fragment['fan_state_records'].append(FanStateRecord(fan_id=series['fan_id'],observed_at=sample['observed_at'],
                    state='running' if sample['value']==1 else 'stopped',complete_fan_state=True,**common).model_dump(mode='json'))
        else:
            fragment['operating_series'].append(OperatingSeries(fan_id=series['fan_id'],metric='flow',unit=series['unit'],
                sensor_id=series['sensor_id'],mapping_confirmed=series['mapping_confirmed'],samples=series['samples'],**common).model_dump(mode='json'))
    return fragment