from __future__ import annotations import hashlib import json from datetime import datetime from typing import Any from ..db import get_connection LEVELS = ("轻微", "异常", "严重异常") SEVERITY_BY_LEVEL = {"轻微": 1, "异常": 2, "严重异常": 3} _SELECT_FIELDS = """ id, source_type, origin, target_type, target_id, target_key, channel_key, segment_key, unit, device_code, device_part, device_point, metric_key, metric_label, start_time, end_time, algorithm_start, algorithm_end, label_type, label_subtype, severity, confirm_status, usable_for_train, source_file, source_row, algorithm_version, features_json, note """ def _time_string(value: Any) -> str: if isinstance(value, datetime): return value.strftime("%Y-%m-%d %H:%M:%S") return str(value or "") def _parse_time(value: str) -> datetime: text = str(value or "").strip().replace("T", " ") for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M"): try: return datetime.strptime(text[:19] if fmt.endswith("%S") else text[:16], fmt) except ValueError: continue raise ValueError("时间格式不正确,应为 YYYY-MM-DD HH:MM:SS") def _snap_hour(value: datetime) -> datetime: return value.replace(minute=0, second=0, microsecond=0) def _row_key(row: dict[str, Any]) -> str: if row.get("target_type") == "pks": unit = str(row.get("unit") or "") metric = str(row.get("metric_key") or "") suffix = metric.split("_")[-1] if "_" in metric else metric return f"YSJ{unit}_{suffix}" if unit and suffix else metric return str(row.get("device_point") or "") def _row_dict(row: dict[str, Any]) -> dict[str, Any]: features: Any = None raw = row.get("features_json") if raw: try: features = json.loads(raw) except (TypeError, ValueError): features = raw return { "id": int(row["id"]), "sourceType": row.get("source_type") or "", "origin": row.get("origin") or "", "targetType": row.get("target_type") or "", "targetId": int(row["target_id"]) if row.get("target_id") is not None else None, "targetKey": row.get("target_key") or "", "channelKey": row.get("channel_key") or "", "unit": row.get("unit") or "", "deviceCode": row.get("device_code") or "", "devicePart": row.get("device_part") or "", "devicePoint": row.get("device_point") or "", "metricKey": row.get("metric_key") or "", "metricLabel": row.get("metric_label") or "", "startTime": _time_string(row.get("start_time")), "endTime": _time_string(row.get("end_time")), "labelType": row.get("label_type") or "", "labelSubtype": row.get("label_subtype"), "severity": int(row["severity"]) if row.get("severity") is not None else None, "confirmStatus": row.get("confirm_status") or "", "usableForTrain": bool(row.get("usable_for_train")), "sourceRow": row.get("source_row") or "", "features": features, "note": row.get("note"), "rowKey": _row_key(row), } def _fetch_label(cursor: Any, label_id: int) -> dict[str, Any]: cursor.execute(f"SELECT {_SELECT_FIELDS} FROM compressor_label WHERE id = %s", (label_id,)) row = cursor.fetchone() if not row: raise ValueError("标注不存在") return _row_dict(row) def query_labels( target_type: str | None = None, unit: str | None = None, device_code: str | None = None, metric_keys: list[str] | None = None, channel_keys: list[str] | None = None, wave_file_ids: list[int] | None = None, start: str | None = None, end: str | None = None, ) -> dict[str, Any]: clauses = ["deleted_at IS NULL"] params: list[Any] = [] if target_type: clauses.append("target_type = %s") params.append(target_type) if unit: clauses.append("unit = %s") params.append(unit) if device_code: clauses.append("device_code = %s") params.append(device_code) if metric_keys: placeholders = ", ".join(["%s"] * len(metric_keys)) clauses.append(f"metric_key IN ({placeholders})") params.extend(metric_keys) if channel_keys: placeholders = ", ".join(["%s"] * len(channel_keys)) clauses.append(f"channel_key IN ({placeholders})") params.extend(channel_keys) if wave_file_ids: placeholders = ", ".join(["%s"] * len(wave_file_ids)) clauses.append(f"target_id IN ({placeholders})") params.extend(wave_file_ids) if start: clauses.append("end_time > %s") params.append(start) if end: clauses.append("start_time < %s") params.append(end) sql = f""" SELECT {_SELECT_FIELDS} FROM compressor_label WHERE {' AND '.join(clauses)} ORDER BY channel_key ASC, start_time ASC, id ASC """ with get_connection() as connection: with connection.cursor() as cursor: cursor.execute(sql, params) rows = cursor.fetchall() return { "source": "database", "labels": [_row_dict(row) for row in rows], "notice": None, } def _insert_row(cursor: Any, row: dict[str, Any]) -> int: fields = ( "source_type", "origin", "target_type", "target_id", "target_key", "channel_key", "segment_key", "unit", "device_code", "device_part", "device_point", "metric_key", "metric_label", "start_time", "end_time", "algorithm_start", "algorithm_end", "label_type", "label_subtype", "severity", "confirm_status", "usable_for_train", "source_file", "source_row", "algorithm_version", "features_json", "note", ) placeholders = ", ".join(["%s"] * len(fields)) cursor.execute( f"INSERT INTO compressor_label ({', '.join(fields)}) VALUES ({placeholders})", tuple(row.get(field) for field in fields), ) return int(cursor.lastrowid) def _clone_row_with_range(cursor: Any, row: dict[str, Any], start: datetime, end: datetime) -> int: clone = dict(row) clone["start_time"] = _time_string(start) clone["end_time"] = _time_string(end) if row.get("origin") == "algorithm": base = str(row.get("segment_key") or row.get("id")) clone["segment_key"] = hashlib.md5( f"{base}|range|{_time_string(start)}|{_time_string(end)}".encode("utf-8") ).hexdigest() else: clone["segment_key"] = None return _insert_row(cursor, clone) def _trim_algorithm_row(cursor: Any, row: dict[str, Any], start: datetime, end: datetime) -> None: row_start = _parse_time(_time_string(row["start_time"])) row_end = _parse_time(_time_string(row["end_time"])) if row_end <= start or row_start >= end: return keep_before = row_start < start keep_after = row_end > end if keep_before and keep_after: cursor.execute( "UPDATE compressor_label SET end_time = %s WHERE id = %s", (_time_string(start), row["id"]), ) clone = dict(row) clone["start_time"] = _time_string(end) clone["end_time"] = _time_string(row_end) clone["algorithm_start"] = row.get("algorithm_start") or row["start_time"] clone["algorithm_end"] = row.get("algorithm_end") or row["end_time"] base = str(row.get("segment_key") or row["id"]) clone["segment_key"] = hashlib.md5(f"{base}|split|{_time_string(end)}".encode("utf-8")).hexdigest() clone["origin"] = "algorithm" _insert_row(cursor, clone) elif keep_before: cursor.execute( "UPDATE compressor_label SET end_time = %s WHERE id = %s", (_time_string(start), row["id"]), ) elif keep_after: cursor.execute( "UPDATE compressor_label SET start_time = %s WHERE id = %s", (_time_string(end), row["id"]), ) else: cursor.execute( "UPDATE compressor_label SET deleted_at = NOW() WHERE id = %s", (row["id"],), ) def _apply_manual( cursor: Any, payload: dict[str, Any], exclude_id: int | None = None, carry: dict[str, Any] | None = None, ) -> int: label_type = str(payload.get("label_type") or "").strip() if label_type not in LEVELS: raise ValueError("标签类型只能是 轻微 / 异常 / 严重异常") start = _snap_hour(_parse_time(payload["start_time"])) end = _snap_hour(_parse_time(payload["end_time"])) if end <= start: raise ValueError("结束时间必须大于开始时间") channel_key = str(payload.get("channel_key") or "").strip() if not channel_key: raise ValueError("channel_key 不能为空") cursor.execute( f""" SELECT {_SELECT_FIELDS} FROM compressor_label WHERE channel_key = %s AND deleted_at IS NULL AND start_time <= %s AND end_time >= %s AND (%s IS NULL OR id <> %s) """, (channel_key, _time_string(end), _time_string(start), exclude_id, exclude_id), ) rows = cursor.fetchall() for row in rows: if row["origin"] == "manual" and row["label_type"] != label_type: raise ValueError("同一子图内不同类型不能重叠或相邻,请先修改或删除相邻标注") same = [row for row in rows if row["label_type"] == label_type] merged_start = min([start] + [_parse_time(_time_string(row["start_time"])) for row in same]) merged_end = max([end] + [_parse_time(_time_string(row["end_time"])) for row in same]) for row in rows: if row["label_type"] != label_type and row["origin"] != "manual": _trim_algorithm_row(cursor, row, merged_start, merged_end) same_ids = [row["id"] for row in same] if same_ids: placeholders = ", ".join(["%s"] * len(same_ids)) cursor.execute( f"UPDATE compressor_label SET deleted_at = NOW() WHERE id IN ({placeholders})", same_ids, ) carry = carry or {} source_type = "pks_manual" if payload.get("target_type") == "pks" else "wave_manual" new_row = { "source_type": source_type, "origin": "manual", "target_type": payload.get("target_type") or ("pks" if payload.get("metric_key") else "wave_file"), "target_id": payload.get("target_id"), "target_key": channel_key, "channel_key": channel_key, "segment_key": None, "unit": str(payload.get("unit") or carry.get("unit") or ""), "device_code": payload.get("device_code") or carry.get("device_code"), "device_part": payload.get("device_part") or carry.get("device_part"), "device_point": payload.get("device_point") or carry.get("device_point"), "metric_key": payload.get("metric_key") or carry.get("metric_key"), "metric_label": payload.get("metric_label") or carry.get("metric_label"), "start_time": _time_string(merged_start), "end_time": _time_string(merged_end), "algorithm_start": carry.get("algorithm_start"), "algorithm_end": carry.get("algorithm_end"), "label_type": label_type, "label_subtype": payload.get("label_subtype") or carry.get("label_subtype"), "severity": SEVERITY_BY_LEVEL.get(label_type), "confirm_status": "人工确认", "usable_for_train": 1, "source_file": carry.get("source_file"), "source_row": carry.get("source_row"), "algorithm_version": carry.get("algorithm_version"), "features_json": carry.get("features_json"), "note": payload.get("note"), } return _insert_row(cursor, new_row) def create_label(payload: dict[str, Any]) -> dict[str, Any]: with get_connection() as connection: with connection.cursor() as cursor: connection.begin() try: label_id = _apply_manual(cursor, payload) result = _fetch_label(cursor, label_id) connection.commit() return result except Exception: connection.rollback() raise def update_label(label_id: int, payload: dict[str, Any]) -> dict[str, Any]: if label_id <= 0: raise ValueError("标注 id 必须为正整数") with get_connection() as connection: with connection.cursor() as cursor: connection.begin() try: cursor.execute( f"SELECT {_SELECT_FIELDS} FROM compressor_label WHERE id = %s AND deleted_at IS NULL", (label_id,), ) old = cursor.fetchone() if not old: raise ValueError("标注不存在") new_id = _apply_manual(cursor, payload, exclude_id=label_id, carry=old) result = _fetch_label(cursor, new_id) if old["origin"] == "manual": cursor.execute( "UPDATE compressor_label SET deleted_at = NOW() WHERE id = %s", (label_id,), ) else: _trim_algorithm_row( cursor, old, _parse_time(result["startTime"]), _parse_time(result["endTime"]), ) connection.commit() return result except Exception: connection.rollback() raise def delete_label(label_id: int) -> dict[str, Any]: if label_id <= 0: raise ValueError("标注 id 必须为正整数") with get_connection() as connection: with connection.cursor() as cursor: cursor.execute( "UPDATE compressor_label SET deleted_at = NOW() WHERE id = %s AND deleted_at IS NULL", (label_id,), ) if cursor.rowcount == 0: raise ValueError("标注不存在") return {"deleted": label_id} def delete_label_range(label_id: int, start: str, end: str) -> dict[str, Any]: if label_id <= 0: raise ValueError("标注 id 必须为正整数") cut_start = _snap_hour(_parse_time(start)) cut_end = _snap_hour(_parse_time(end)) if cut_end <= cut_start: raise ValueError("结束时间必须大于开始时间") with get_connection() as connection: with connection.cursor() as cursor: connection.begin() try: cursor.execute( f"SELECT {_SELECT_FIELDS} FROM compressor_label WHERE id = %s AND deleted_at IS NULL", (label_id,), ) row = cursor.fetchone() if not row: raise ValueError("标注不存在") row_start = _parse_time(_time_string(row["start_time"])) row_end = _parse_time(_time_string(row["end_time"])) lo = max(row_start, cut_start) hi = min(row_end, cut_end) if hi <= lo: raise ValueError("删除区间不在该标注内") cursor.execute( "UPDATE compressor_label SET deleted_at = NOW() WHERE id = %s", (label_id,), ) if lo <= row_start and hi >= row_end: connection.commit() return { "deleted": label_id, "deletedRange": [_time_string(lo), _time_string(hi)], "beforeId": None, "afterId": None, "whole": True, } before_id = None after_id = None if row_start < lo: before_id = _clone_row_with_range(cursor, row, row_start, lo) if hi < row_end: after_id = _clone_row_with_range(cursor, row, hi, row_end) connection.commit() return { "deleted": label_id, "deletedRange": [_time_string(lo), _time_string(hi)], "beforeId": before_id, "afterId": after_id, "whole": False, } except Exception: connection.rollback() raise