Compressor/backend/app/db.py

import pymysql

from .config import settings


def get_connection():
    return pymysql.connect(
        host=settings.db_host,
        port=settings.db_port,
        user=settings.db_user,
        password=settings.db_password,
        database=settings.db_name,
        charset="utf8mb4",
        cursorclass=pymysql.cursors.DictCursor,
        connect_timeout=settings.db_connect_timeout,
        read_timeout=30,
        write_timeout=30,
        autocommit=True,
    )


_ANNOTATION_TABLE_DDL = """
CREATE TABLE IF NOT EXISTS wave_annotation (
    id                 BIGINT UNSIGNED NOT NULL AUTO_INCREMENT COMMENT '标注ID',
    wave_file_id       BIGINT UNSIGNED NOT NULL COMMENT '所属波形文件ID（wave_file.id）',
    label              VARCHAR(8)      NOT NULL COMMENT '样本类型：正常 / 异常',
    period_start       INT             NOT NULL COMMENT '起始周期编号',
    period_end         INT             NOT NULL COMMENT '结束周期编号（含）',
    sample_index_start INT             NOT NULL COMMENT '起始采样点索引',
    sample_index_end   INT             NOT NULL COMMENT '结束采样点索引（含）',
    created_at         DATETIME        NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
    updated_at         DATETIME        NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
    PRIMARY KEY (id),
    KEY idx_wave_file (wave_file_id)
) ENGINE = InnoDB DEFAULT CHARSET = utf8mb4 COMMENT = '波形标注索引表';
"""


def ensure_annotation_table() -> None:
    with get_connection() as connection:
        with connection.cursor() as cursor:
            cursor.execute(_ANNOTATION_TABLE_DDL)


_LABEL_TABLE_DDL = """
CREATE TABLE IF NOT EXISTS compressor_label (
    id                BIGINT UNSIGNED NOT NULL AUTO_INCREMENT COMMENT '标签ID',
    source_type       VARCHAR(32)     NOT NULL COMMENT 'vibration/oil_pressure/temperature/pressure_wave/pks_manual',
    origin            VARCHAR(16)     NOT NULL DEFAULT 'manual' COMMENT 'algorithm/manual/imported',
    target_type       VARCHAR(16)     NOT NULL COMMENT 'wave_file/pks',
    target_id         BIGINT UNSIGNED NULL COMMENT 'wave_file.id',
    target_key        VARCHAR(200)    NOT NULL COMMENT '通道+目标标识',
    channel_key       VARCHAR(200)    NOT NULL COMMENT '子图通道键',
    segment_key       VARCHAR(64)     NULL COMMENT '算法原始段分组键',
    unit              VARCHAR(16)     NOT NULL COMMENT '机组号，如 7',
    device_code       VARCHAR(64)     NULL COMMENT '如 7#',
    device_part       VARCHAR(200)    NULL COMMENT '如 7号机组一缸 / 7号机',
    device_point      VARCHAR(200)    NULL COMMENT 'wave 点位或 PKS 指标名',
    metric_key        VARCHAR(64)     NULL COMMENT 'PKS 列名，如 YSJ_11',
    metric_label      VARCHAR(200)    NULL COMMENT 'PKS 指标中文名',
    start_time        DATETIME        NOT NULL COMMENT '标签段起点',
    end_time          DATETIME        NOT NULL COMMENT '标签段终点（不含）',
    algorithm_start   DATETIME        NULL COMMENT '算法原始段起点',
    algorithm_end     DATETIME        NULL COMMENT '算法原始段终点',
    label_type        VARCHAR(32)     NOT NULL COMMENT '轻微/异常/严重异常',
    label_subtype     VARCHAR(64)     NULL COMMENT '振动侧/温度部位/方向等',
    severity          TINYINT         NULL COMMENT '1轻微/2异常/3严重异常',
    confirm_status    VARCHAR(16)     NOT NULL DEFAULT '未确认' COMMENT '未确认/已确认/人工否决',
    usable_for_train  TINYINT(1)      NOT NULL DEFAULT 1,
    source_file       VARCHAR(500)    NULL,
    source_row        VARCHAR(64)     NULL COMMENT '周期编号等',
    algorithm_version VARCHAR(64)     NULL,
    features_json     MEDIUMTEXT      NULL COMMENT '算法原始属性',
    note              VARCHAR(1000)   NULL,
    created_by        VARCHAR(64)     NULL,
    updated_by        VARCHAR(64)     NULL,
    created_at        DATETIME        NOT NULL DEFAULT CURRENT_TIMESTAMP,
    updated_at        DATETIME        NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
    deleted_at        DATETIME        NULL,
    PRIMARY KEY (id),
    UNIQUE KEY uk_segment (segment_key),
    KEY idx_target_time (target_key, start_time, end_time),
    KEY idx_channel_time (channel_key, start_time),
    KEY idx_unit_time (unit, start_time),
    KEY idx_source (source_type, origin)
) ENGINE = InnoDB DEFAULT CHARSET = utf8mb4 COMMENT = '统一训练标签表';
"""


def ensure_label_table() -> None:
    with get_connection() as connection:
        with connection.cursor() as cursor:
            cursor.execute(_LABEL_TABLE_DDL)