massive-cloud/scripts/backfill-flatness-suite-measurements.ts

#!/usr/bin/env node

import { execFile as execFileCallback } from "node:child_process";
import { promises as fs } from "node:fs";
import path from "node:path";
import { promisify } from "node:util";
import { pathToFileURL } from "node:url";
import mysql, { type Pool, type ResultSetHeader, type RowDataPacket } from "mysql2/promise";

const execFile = promisify(execFileCallback);

export type Measurement = Record<string, unknown>;

export interface CandidateRow extends RowDataPacket {
  id: number;
  taskId: number;
  measurements: unknown;
  resultData: Record<string, unknown> | null;
}

export interface BackfillPlan {
  row: CandidateRow;
  resultJsonPath: string;
  originalResultJson: Record<string, unknown>;
  measurements: Measurement[];
}

export interface ExecutionSummary {
  planned: number;
  applied: number;
  skipped: number;
  failed: number;
}

export function mergeMeasurements<T extends Record<string, unknown>>(
  value: T | null,
  measurements: Measurement[],
): T & { measurements: Measurement[] } {
  return { ...(value ?? ({} as T)), measurements };
}

export async function executePlans(
  plans: BackfillPlan[],
  apply: boolean,
  persist: (plan: BackfillPlan) => Promise<boolean>,
): Promise<ExecutionSummary> {
  const summary: ExecutionSummary = {
    planned: plans.length,
    applied: 0,
    skipped: apply ? 0 : plans.length,
    failed: 0,
  };
  if (!apply) return summary;

  for (const plan of plans) {
    try {
      if (await persist(plan)) summary.applied += 1;
      else summary.skipped += 1;
    } catch (error) {
      summary.failed += 1;
      console.error(`[backfill] 结果 ${plan.row.id} 写入失败:`, error);
    }
  }
  return summary;
}

function argumentValue(name: string): string | undefined {
  const index = process.argv.indexOf(name);
  return index >= 0 ? process.argv[index + 1] : undefined;
}

function timestamp(): string {
  return new Date().toISOString().replace(/[:.]/g, "-");
}

async function parseMeasurements(
  excelPath: string,
  templatePath: string,
  helperPath: string,
): Promise<Measurement[]> {
  const python = process.env.PYTHON3_PATH || "/usr/bin/python3";
  const { stdout } = await execFile(
    python,
    [helperPath, "--excel", excelPath, "--template", templatePath],
    { maxBuffer: 10 * 1024 * 1024 },
  );
  const parsed = JSON.parse(stdout) as unknown;
  if (!Array.isArray(parsed) || parsed.length === 0) {
    throw new Error("解析结果为空")
  }
  return parsed.filter((item): item is Measurement => !!item && typeof item === "object");
}

async function listCandidates(pool: Pool, resultId?: number): Promise<CandidateRow[]> {
  const filters = [
    "pluginId = 'flatness-suite'",
    "status = 'completed'",
    "(measurements IS NULL OR JSON_LENGTH(measurements) = 0)",
  ];
  const params: number[] = [];
  if (resultId !== undefined) {
    filters.push("id = ?");
    params.push(resultId);
  }
  const [rows] = await pool.query<CandidateRow[]>(
    `SELECT id, taskId, measurements, resultData
       FROM task_analysis_results
      WHERE ${filters.join(" AND ")}
      ORDER BY id`,
    params,
  );
  return rows;
}

async function buildPlans(
  rows: CandidateRow[],
  dataDir: string,
  templatePath: string,
  helperPath: string,
): Promise<{ plans: BackfillPlan[]; skipped: number }> {
  const plans: BackfillPlan[] = [];
  let skipped = 0;
  for (const row of rows) {
    const resultDir = path.join(dataDir, "tasks", String(row.taskId), "analysis", String(row.id));
    const resultJsonPath = path.join(resultDir, "result.json");
    const excelPath = path.join(resultDir, "contour.xlsx");
    try {
      const [resultText] = await Promise.all([
        fs.readFile(resultJsonPath, "utf8"),
        fs.access(excelPath),
      ]);
      const resultJson = JSON.parse(resultText) as Record<string, unknown>;
      const extra = resultJson.extra as Record<string, unknown> | undefined;
      const resultDataExtra = row.resultData?.extra as Record<string, unknown> | undefined;
      if (extra?.module !== "contour_batch" && resultDataExtra?.module !== "contour_batch") {
        throw new Error("不是 contour_batch 批量结果")
      }
      const measurements = await parseMeasurements(excelPath, templatePath, helperPath);
      plans.push({ row, resultJsonPath, originalResultJson: resultJson, measurements });
      console.log(`[backfill] 结果 ${row.id}: 可回填 ${measurements.length} 条测量值`);
    } catch (error) {
      skipped += 1;
      console.warn(`[backfill] 跳过结果 ${row.id}: ${error instanceof Error ? error.message : String(error)}`);
    }
  }
  return { plans, skipped };
}

async function writeJsonAtomically(filePath: string, value: unknown): Promise<void> {
  const temporaryPath = `${filePath}.backfill-${process.pid}.tmp`;
  await fs.writeFile(temporaryPath, `${JSON.stringify(value, null, 2)}\n`, "utf8");
  await fs.rename(temporaryPath, filePath);
}

async function writeBackup(plans: BackfillPlan[], dataDir: string): Promise<string> {
  const backupDir = path.join(dataDir, "maintenance-backups");
  await fs.mkdir(backupDir, { recursive: true });
  const backupPath = path.join(backupDir, `flatness-suite-measurements-${timestamp()}.json`);
  await writeJsonAtomically(backupPath, {
    version: 1,
    createdAt: new Date().toISOString(),
    entries: plans.map((plan) => ({
      id: plan.row.id,
      taskId: plan.row.taskId,
      resultJsonPath: plan.resultJsonPath,
      measurements: plan.row.measurements,
      resultData: plan.row.resultData,
      resultJson: plan.originalResultJson,
    })),
  });
  return backupPath;
}

async function persistPlan(pool: Pool, plan: BackfillPlan): Promise<boolean> {
  const connection = await pool.getConnection();
  const updatedResultJson = mergeMeasurements(plan.originalResultJson, plan.measurements);
  const updatedResultData = mergeMeasurements(plan.row.resultData, plan.measurements);
  try {
    await connection.beginTransaction();
    await writeJsonAtomically(plan.resultJsonPath, updatedResultJson);
    const [result] = await connection.execute<ResultSetHeader>(
      `UPDATE task_analysis_results
          SET measurements = ?, resultData = ?
        WHERE id = ?
          AND pluginId = 'flatness-suite'
          AND status = 'completed'
          AND (measurements IS NULL OR JSON_LENGTH(measurements) = 0)`,
      [JSON.stringify(plan.measurements), JSON.stringify(updatedResultData), plan.row.id],
    );
    if (result.affectedRows !== 1) {
      await writeJsonAtomically(plan.resultJsonPath, plan.originalResultJson);
      await connection.rollback();
      return false;
    }
    await connection.commit();
    return true;
  } catch (error) {
    await connection.rollback();
    await writeJsonAtomically(plan.resultJsonPath, plan.originalResultJson).catch(() => undefined);
    throw error;
  } finally {
    connection.release();
  }
}

interface BackupManifest {
  entries: Array<{
    id: number;
    resultJsonPath: string;
    measurements: unknown;
    resultData: unknown;
    resultJson: unknown;
  }>;
}

async function restoreBackup(pool: Pool, backupPath: string): Promise<void> {
  const manifest = JSON.parse(await fs.readFile(backupPath, "utf8")) as BackupManifest;
  if (!Array.isArray(manifest.entries)) throw new Error("备份清单格式无效")
  for (const entry of manifest.entries) {
    await writeJsonAtomically(entry.resultJsonPath, entry.resultJson);
    await pool.execute(
      "UPDATE task_analysis_results SET measurements = ?, resultData = ? WHERE id = ? AND pluginId = 'flatness-suite'",
      [
        entry.measurements == null ? null : JSON.stringify(entry.measurements),
        entry.resultData == null ? null : JSON.stringify(entry.resultData),
        entry.id,
      ],
    );
    console.log(`[backfill] 已恢复结果 ${entry.id}`);
  }
}

async function main(): Promise<void> {
  const databaseUrl = process.env.DATABASE_URL;
  if (!databaseUrl) throw new Error("缺少 DATABASE_URL")
  const dataDir = process.env.DATA_DIR || "/data";
  const algRoot = process.env.ALG_ROOT || "/alg";
  const resultIdRaw = argumentValue("--result-id");
  const resultId = resultIdRaw === undefined ? undefined : Number(resultIdRaw);
  if (resultIdRaw !== undefined && !Number.isInteger(resultId)) throw new Error("--result-id 必须是整数")
  const apply = process.argv.includes("--apply");
  const restorePath = argumentValue("--restore");
  const templatePath = path.join(algRoot, "平整度、平面度、轮廓测量v1.5", "excel", "测量结果.xlsx");
  const helperPath = path.resolve("server/plugins/adapters/flatness_suite_measurements.py");
  const pool = mysql.createPool(databaseUrl);

  try {
    if (restorePath) {
      await restoreBackup(pool, restorePath);
      return;
    }
    const rows = await listCandidates(pool, resultId);
    const { plans, skipped } = await buildPlans(rows, dataDir, templatePath, helperPath);
    let backupPath: string | undefined;
    if (apply && plans.length > 0) {
      backupPath = await writeBackup(plans, dataDir);
      console.log(`[backfill] 写入前备份: ${backupPath}`);
    }
    const summary = await executePlans(plans, apply, (plan) => persistPlan(pool, plan));
    summary.skipped += skipped;
    console.log(`[backfill] ${apply ? "apply" : "dry-run"} 完成: ${JSON.stringify(summary)}`);
    if (!apply && plans.length > 0) console.log("[backfill] 确认无误后追加 --apply 执行写入");
    if (backupPath) console.log(`[backfill] 恢复命令: pnpm tsx scripts/backfill-flatness-suite-measurements.ts --restore ${backupPath}`);
  } finally {
    await pool.end();
  }
}

const isMain = process.argv[1] && import.meta.url === pathToFileURL(path.resolve(process.argv[1])).href;
if (isMain) {
  main().catch((error) => {
    console.error("[backfill] 失败:", error);
    process.exitCode = 1;
  });
}