sing-ui/lib/sling.ts

import { execFileSync } from "child_process";
import fs from "fs";
import path from "path";
import { dump } from "js-yaml";
import {
  Connection,
  DATA_DIR,
  Pipeline,
  StreamConfig,
  getConnection,
  listConnections,
} from "./db";
import { resolveBinary } from "./binaries";

export const SLING_HOME = path.join(DATA_DIR, "sling-home");

export interface DiscoveredStream {
  name: string;
  columns?: string[];
}

// Parse a connection/pipeline JSON-object column; tolerate bad data.
function parseJsonObject(json: string | undefined | null): Record<string, string> {
  if (!json) return {};
  try {
    const v = JSON.parse(json) as unknown;
    if (v && typeof v === "object" && !Array.isArray(v)) {
      return Object.fromEntries(
        Object.entries(v as Record<string, unknown>).map(([k, val]) => [
          k,
          String(val),
        ])
      );
    }
  } catch {
    // fall through
  }
  return {};
}

// Append query params to a URL, percent-encoding keys and values.
export function appendParams(url: string, params: Record<string, string>): string {
  const entries = Object.entries(params);
  if (entries.length === 0) return url;
  const qs = entries
    .map(([k, v]) => `${encodeURIComponent(k)}=${encodeURIComponent(v)}`)
    .join("&");
  return `${url}${url.includes("?") ? "&" : "?"}${qs}`;
}

// Build a mysql:// URL from a connection record, escaping credentials and
// merging the connection's extra params into the query string. These params
// reach both sling (env.yaml) and atlas (--url/--to/--dev-url); the
// go-sql-driver SETs unknown DSN params as session variables, which is how
// e.g. tidb_skip_isolation_level_check=1 takes effect.
export function connectionUrl(conn: Connection): string {
  const user = encodeURIComponent(conn.user);
  const pass = encodeURIComponent(conn.password);
  const base = `mysql://${user}:${pass}@${conn.host}:${conn.port}/${conn.database}`;
  return appendParams(base, parseJsonObject(conn.params));
}

// Regenerate env.yaml from the connections table. Names are uppercased
// because sling treats connection names case-insensitively. Session-level
// foreign_key_checks=0 lets sling truncate/drop FK-referenced parent tables
// in any order (same approach as mysqldump); constraints stay defined. It is
// appended only here (atlas doesn't need it) and never overrides a value the
// user set explicitly in connection params.
export function syncEnvYaml(): string {
  fs.mkdirSync(SLING_HOME, { recursive: true });
  const connections: Record<string, { url: string }> = {};
  for (const c of listConnections()) {
    let url = connectionUrl(c);
    if (!parseJsonObject(c.params)["foreign_key_checks"]) {
      url = appendParams(url, { foreign_key_checks: "0" });
    }
    connections[c.name.toUpperCase()] = { url };
  }
  const file = path.join(SLING_HOME, "env.yaml");
  fs.writeFileSync(file, dump({ connections }));
  return file;
}

function slingBin(): string {
  return resolveBinary("sling") ?? "sling";
}

export function testConnection(id: number): { ok: boolean; output: string } {
  const conn = getConnection(id);
  if (!conn) throw new Error(`connection ${id} not found`);
  syncEnvYaml();
  try {
    const output = execFileSync(
      /* turbopackIgnore: true */ slingBin(),
      ["conns", "test", conn.name, "--home-dir", SLING_HOME],
      { encoding: "utf8", timeout: 30000 }
    );
    return { ok: true, output };
  } catch (e) {
    const err = e as { stdout?: string; stderr?: string; message?: string };
    return {
      ok: false,
      output: [err.stdout, err.stderr, err.message]
        .filter(Boolean)
        .join("\n"),
    };
  }
}

// `conns discover --output json --columns` emits a tabular payload:
//   { fields: ["Database","Schema","Table","ID","Column","Native Type","General Type"],
//     rows: [[db, schema, table, id, column, native, general], ...] }
// covering every schema on the server. Group rows by table, keep only the
// connection's own database, and collect column names in row order.
function fromFieldsRows(
  d: Record<string, unknown>,
  database: string
): DiscoveredStream[] | null {
  if (!Array.isArray(d.fields) || !Array.isArray(d.rows)) return null;
  const fields = d.fields.map((f) => String(f).toLowerCase());
  // The "Database" column holds the connection's database for every row, so
  // prefer "Schema" (the table's actual schema) for filtering.
  const schemaIdx =
    fields.findIndex((f) => f === "schema") >= 0
      ? fields.findIndex((f) => f === "schema")
      : fields.findIndex((f) => f === "database");
  const tableIdx = fields.findIndex((f) => f === "table");
  const colIdx = fields.findIndex((f) => f === "column");
  if (tableIdx < 0) return null;
  const byTable = new Map<string, string[]>();
  for (const row of d.rows) {
    if (!Array.isArray(row)) continue;
    const schema = schemaIdx >= 0 ? String(row[schemaIdx] ?? "") : "";
    if (database && schema && schema !== database) continue;
    const table = String(row[tableIdx] ?? "");
    if (!table) continue;
    let cols = byTable.get(table);
    if (!cols) byTable.set(table, (cols = []));
    if (colIdx >= 0 && row[colIdx] != null) cols.push(String(row[colIdx]));
  }
  return [...byTable.entries()].map(([name, columns]) => ({ name, columns }));
}

// Defensive fallbacks for other plausible shapes: string array, object array,
// { streams: [...] }, or a map of stream name -> { columns }.
function normalizeDiscover(
  data: unknown,
  database: string
): DiscoveredStream[] {
  if (data && typeof data === "object" && !Array.isArray(data)) {
    const tabular = fromFieldsRows(data as Record<string, unknown>, database);
    if (tabular) return tabular;
  }
  const out: DiscoveredStream[] = [];
  const toColumns = (raw: unknown): string[] | undefined => {
    if (!Array.isArray(raw)) return undefined;
    return raw.map((c) =>
      typeof c === "string"
        ? c
        : String((c as Record<string, unknown>)?.name ?? c)
    );
  };
  const push = (item: unknown) => {
    if (typeof item === "string") {
      out.push({ name: item });
      return;
    }
    if (item && typeof item === "object") {
      const o = item as Record<string, unknown>;
      const name = String(o.name ?? o.stream ?? o.table ?? "");
      if (!name) return;
      const columns = toColumns(o.columns);
      out.push(columns ? { name, columns } : { name });
    }
  };
  if (Array.isArray(data)) {
    data.forEach(push);
  } else if (data && typeof data === "object") {
    const d = data as Record<string, unknown>;
    if (Array.isArray(d.streams)) {
      d.streams.forEach(push);
    } else {
      for (const [key, value] of Object.entries(d)) {
        const columns = toColumns((value as Record<string, unknown>)?.columns);
        out.push(columns ? { name: key, columns } : { name: key });
      }
    }
  }
  return out;
}

function parseJsonLoose(text: string): unknown {
  try {
    return JSON.parse(text);
  } catch {
    // stdout may carry log noise around the payload; slice from the first brace
    const start = text.search(/[[{]/);
    if (start >= 0) return JSON.parse(text.slice(start));
    throw new Error("sling discover did not output JSON");
  }
}

// Execute a SQL statement on a connection via `sling conns exec`.
export function execSql(conn: Connection, sql: string): void {
  syncEnvYaml();
  execFileSync(
    /* turbopackIgnore: true */ slingBin(),
    ["conns", "exec", conn.name.toUpperCase(), sql, "--home-dir", SLING_HOME],
    { encoding: "utf8", timeout: 30000 }
  );
}

export function discoverStreams(id: number): DiscoveredStream[] {
  const conn = getConnection(id);
  if (!conn) throw new Error(`connection ${id} not found`);
  syncEnvYaml();
  const output = execFileSync(
    /* turbopackIgnore: true */ slingBin(),
    [
      "conns",
      "discover",
      conn.name,
      "--output",
      "json",
      "--columns",
      "--home-dir",
      SLING_HOME,
    ],
    { encoding: "utf8", timeout: 60000, maxBuffer: 32 * 1024 * 1024 }
  );
  return normalizeDiscover(parseJsonLoose(output), conn.database);
}

// Render the pipeline into a sling replication.yaml for a given run.
export function buildReplicationYaml(
  pipeline: Pipeline,
  runId: number
): string {
  const src = getConnection(pipeline.source_conn_id);
  const tgt = getConnection(pipeline.target_conn_id);
  if (!src || !tgt) throw new Error("pipeline connections not found");
  syncEnvYaml();

  const streams = JSON.parse(pipeline.streams || "[]") as StreamConfig[];
  const streamsYaml: Record<string, Record<string, unknown> | null> = {};
  for (const s of streams) {
    const cfg: Record<string, unknown> = {};
    if (s.object) cfg.object = s.object;
    if (s.primary_key) cfg.primary_key = s.primary_key;
    if (s.update_key) cfg.update_key = s.update_key;
    if (s.select) cfg.select = s.select;
    // sling requires database-qualified stream names (e.g. "mydb.users") to
    // resolve source columns; bare table names fail with "did not find any
    // columns".
    const key = s.name.includes(".") ? s.name : `${src.database}.${s.name}`;
    streamsYaml[key] = Object.keys(cfg).length > 0 ? cfg : null;
  }

  // With schema_sync on, Atlas owns the target schema (indexes, FKs, exact
  // column types). sling's full-refresh would drop & recreate tables with
  // inferred types and no constraints, so translate it to truncate, which
  // keeps the Atlas-managed schema and reloads all rows.
  const mode =
    pipeline.schema_sync && (pipeline.mode || "full-refresh") === "full-refresh"
      ? "truncate"
      : pipeline.mode || "full-refresh";

  const doc = {
    source: src.name.toUpperCase(),
    target: tgt.name.toUpperCase(),
    defaults: {
      mode,
      // sling also requires a schema-qualified target object name.
      object: `${tgt.database}.{stream_table}`,
    },
    streams: streamsYaml,
    // Pipeline-level env is exported as process env vars by sling. User
    // values override the built-in defaults (e.g. SLING_THREADS).
    env: { SLING_THREADS: "4", ...parseJsonObject(pipeline.env) },
  };

  const dir = path.join(DATA_DIR, "replications");
  fs.mkdirSync(dir, { recursive: true });
  const file = path.join(dir, `${runId}.yaml`);
  fs.writeFileSync(file, dump(doc));
  return file;
}