sing-ui/lib/atlas.ts

import { execFile } from "child_process";
import { Connection, Pipeline, StreamConfig, getConnection } from "./db";
import { appendParams, connectionUrl, discoverStreams, execSql } from "./sling";
import { resolveBinary } from "./binaries";

// Atlas needs an empty scratch database ("dev database") to plan diffs.
// The usual `docker://mysql/8` dev-url spins up a throwaway container, but
// its dial to host.docker.internal is broken on some Docker runtimes (e.g.
// OrbStack: connections establish, then the handshake EOFs). Use a dedicated
// scratch schema on the target server instead — no Docker needed.
const DEV_SCHEMA = "sling_ui_atlas_dev";

// The scratch schema lives on the target server, so it inherits the target
// connection's URL params (e.g. tidb_skip_isolation_level_check for TiDB).
function devDatabaseUrl(tgt: Connection): string {
  const user = encodeURIComponent(tgt.user);
  const pass = encodeURIComponent(tgt.password);
  const base = `mysql://${user}:${pass}@${tgt.host}:${tgt.port}/${DEV_SCHEMA}`;
  let params: Record<string, string> = {};
  try {
    params = JSON.parse(tgt.params || "{}") as Record<string, string>;
  } catch {
    // ignore malformed params
  }
  return appendParams(base, params);
}

// Create the scratch schema if missing and return its URL. Atlas keeps it
// clean after each run.
export function prepareDevDatabase(tgt: Connection): string {
  execSql(tgt, `CREATE DATABASE IF NOT EXISTS \`${DEV_SCHEMA}\``);
  return devDatabaseUrl(tgt);
}

function baseApplyArgs(
  srcUrl: string,
  tgtUrl: string,
  devUrl: string,
  excludes: string[]
): string[] {
  const args = [
    "schema",
    "apply",
    "--url",
    tgtUrl,
    "--to",
    srcUrl,
    "--dev-url",
    devUrl,
  ];
  // Community Atlas has no --include; --exclude globs match bare resource
  // names (no schema prefix). Excluded objects are dropped from the diff
  // entirely, so e.g. unsupported column types in them never surface.
  for (const p of excludes) args.push("--exclude", p);
  return args;
}

export function buildApplyArgs(
  srcUrl: string,
  tgtUrl: string,
  devUrl: string,
  excludes: string[] = []
): string[] {
  return [...baseApplyArgs(srcUrl, tgtUrl, devUrl, excludes), "--auto-approve"];
}

// atlas rejects --dry-run combined with --auto-approve.
export function buildDryRunArgs(
  srcUrl: string,
  tgtUrl: string,
  devUrl: string,
  excludes: string[] = []
): string[] {
  return [...baseApplyArgs(srcUrl, tgtUrl, devUrl, excludes), "--dry-run"];
}

// Compute the --exclude patterns for a pipeline's schema step. Atlas diffs
// whole databases, and tables that exist only on the target would be planned
// as DROPs — so for 'selected' scope we must exclude every non-selected table
// from BOTH sides. Bare table names only: exclude globs match resource names
// without schema prefixes.
export async function computeSchemaExcludes(
  pipeline: Pipeline
): Promise<string[]> {
  if ((pipeline.schema_scope || "selected") === "all") return [];

  const streams = JSON.parse(pipeline.streams || "[]") as StreamConfig[];
  const selected = new Set(
    streams.map((s) => {
      const i = s.name.lastIndexOf(".");
      return i >= 0 ? s.name.slice(i + 1) : s.name;
    })
  );

  const tableNames = (connId: number, side: string): Set<string> => {
    try {
      return new Set(
        discoverStreams(connId).map((s) => {
          const i = s.name.lastIndexOf(".");
          return i >= 0 ? s.name.slice(i + 1) : s.name;
        })
      );
    } catch (e) {
      // Silently falling back would risk DROPping target-only tables.
      throw new Error(
        `无法获取${side}库表列表，已中止 schema 同步：${(e as Error).message}`
      );
    }
  };

  const srcTables = tableNames(pipeline.source_conn_id, "源");
  const tgtTables = tableNames(pipeline.target_conn_id, "目标");

  const excludes: string[] = [];
  for (const t of new Set([...srcTables, ...tgtTables])) {
    if (!selected.has(t)) excludes.push(t);
  }
  // FKs referencing excluded tables vanish from the desired state anyway;
  // excluding all FKs additionally avoids touching constraints on the
  // selected tables themselves.
  if (!pipeline.include_fk) excludes.push("*[type=foreign-key]");
  return excludes;
}

function atlasBin(): string {
  return resolveBinary("atlas") ?? "atlas";
}

export async function schemaDiff(
  pipeline: Pipeline
): Promise<{ ok: boolean; sql: string }> {
  const src = getConnection(pipeline.source_conn_id);
  const tgt = getConnection(pipeline.target_conn_id);
  if (!src || !tgt) throw new Error("connection not found");
  let excludes: string[];
  try {
    excludes = await computeSchemaExcludes(pipeline);
  } catch (e) {
    return { ok: false, sql: (e as Error).message };
  }
  let devUrl: string;
  try {
    devUrl = prepareDevDatabase(tgt);
  } catch (e) {
    const err = e as { stdout?: string; stderr?: string; message?: string };
    return {
      ok: false,
      sql: [err.stdout, err.stderr, err.message].filter(Boolean).join("\n"),
    };
  }
  const args = buildDryRunArgs(
    connectionUrl(src),
    connectionUrl(tgt),
    devUrl,
    excludes
  );
  return new Promise((resolve) => {
    execFile(
      /* turbopackIgnore: true */ atlasBin(),
      args,
      { timeout: 120000, maxBuffer: 32 * 1024 * 1024 },
      (error, stdout, stderr) => {
        const sql = [stdout, stderr].filter(Boolean).join("\n");
        resolve({ ok: !error, sql });
      }
    );
  });
}