feat(data-compare): support per-table match column configuration

apps/desktop/src/composables/__tests__/useDataCompareSession.spec.ts

@@ -1,5 +1,7 @@
 import { strict as assert } from "node:assert";
 import { beforeEach, test, vi } from "vitest";
+import type { ColumnInfo } from "@/types/database";
+import type { DataCompareSession, DataCompareSessionConfig } from "../useDataCompareSession.ts";
 
 const apiMock = vi.hoisted(() => ({
   buildDataCompareSyncPlan: vi.fn(),
@@ -18,27 +20,76 @@ vi.mock("@/lib/backend/api", () => apiMock);
 vi.mock("@/composables/useDialogSources", () => ({ openDataCompareSession: openMock }));
 vi.mock("@/composables/useExportTracker", () => ({ useExportTracker: () => trackerMock }));
 
-const { startDataCompareSession } = await import("../useDataCompareSession.ts");
+const { normalizeKeyColumnOverrides, startDataCompareSession } = await import("../useDataCompareSession.ts");
 
-test("runs a data compare session independently of the dialog and retains its result", async () => {
-  apiMock.getColumns.mockResolvedValue([
-    { name: "id", is_primary_key: true },
-    { name: "name", is_primary_key: false },
-  ]);
-  apiMock.prepareDataCompareFromTables.mockResolvedValue({
-    result: {
-      added: [{ key: "1", keyValues: { id: 1 }, values: { id: 1, name: "new" } }],
-      removed: [],
-      modified: [],
-    },
-    syncStatements: ["INSERT INTO users (id, name) VALUES (1, 'new')"],
-    syncSql: "INSERT INTO users (id, name) VALUES (1, 'new')",
+function column(name: string, isPrimaryKey = false): ColumnInfo {
+  return { name, data_type: "TEXT", is_nullable: true, column_default: null, is_primary_key: isPrimaryKey, extra: null };
+}
+
+const TABLES: Record<string, ColumnInfo[]> = {
+  users: [column("user_id", true), column("email"), column("name")],
+  orders: [column("tenant_id", true), column("order_id", true), column("amount")],
+  payments: [column("payment_no", true), column("tenant_id"), column("amount")],
+  logs: [column("category"), column("created_at")],
+  legacy: [column("snid", true), column("username")],
+};
+
+function preparationFor(table: string) {
+  return {
+    result: { added: [{ key: "1", keyValues: { id: 1 }, values: { id: 1 } }], removed: [], modified: [] },
+    syncStatements: [`INSERT INTO ${table} (id) VALUES (1)`],
+    syncSql: `INSERT INTO ${table} (id) VALUES (1)`,
     preSyncStatements: [],
     sourceRowCount: 1,
     targetRowCount: 0,
     sourceTruncated: false,
     targetTruncated: false,
-  });
+  };
+}
+
+function baseConfig(overrides: Partial<DataCompareSessionConfig> = {}): DataCompareSessionConfig {
+  return {
+    sourceConnectionId: "source",
+    sourceDatabase: "app",
+    sourceSchema: "public",
+    sourceDatabases: ["app"],
+    sourceSchemas: ["public"],
+    sourceTables: Object.keys(TABLES),
+    selectedSourceTables: [],
+    targetConnectionId: "target",
+    targetDatabase: "warehouse",
+    targetSchema: "public",
+    targetDatabases: ["warehouse"],
+    targetSchemas: ["public"],
+    targetTables: Object.keys(TABLES),
+    targetTable: "",
+    keyColumnsByTable: {},
+    label: "app → warehouse",
+    ...overrides,
+  };
+}
+
+const dependencies = {
+  ensureConnected: () => Promise.resolve(undefined),
+  getConfig: () => ({ db_type: "mysql" as const }),
+};
+
+async function runSession(input: DataCompareSessionConfig, tasks: { sourceTable: string; targetTable: string }[]): Promise<DataCompareSession> {
+  const session = startDataCompareSession(input, tasks, dependencies);
+  for (let attempt = 0; attempt < 50 && session.status === "running"; attempt += 1) {
+    await new Promise((resolve) => setTimeout(resolve, 0));
+  }
+  return session;
+}
+
+function keyColumnsOf(call: number): string[] {
+  return apiMock.prepareDataCompareFromTables.mock.calls[call]?.[0]?.keyColumns as string[];
+}
+
+beforeEach(() => {
+  vi.clearAllMocks();
+  apiMock.getColumns.mockImplementation((_connectionId: string, _database: string, _schema: string, table: string) => Promise.resolve(TABLES[table] ?? []));
+  apiMock.prepareDataCompareFromTables.mockImplementation((options: { targetTable: string }) => Promise.resolve(preparationFor(options.targetTable)));
   apiMock.buildDataCompareSyncPlan.mockResolvedValue({
     insertCount: 1,
     updateCount: 0,
@@ -47,36 +98,10 @@ test("runs a data compare session independently of the dialog and retains its re
     syncStatements: ["INSERT"],
     syncSql: "INSERT",
   });
+});
 
-  const session = startDataCompareSession(
-    {
-      sourceConnectionId: "source",
-      sourceDatabase: "app",
-      sourceSchema: "public",
-      sourceDatabases: ["app"],
-      sourceSchemas: ["public"],
-      sourceTables: ["users"],
-      selectedSourceTables: ["users"],
-      targetConnectionId: "target",
-      targetDatabase: "warehouse",
-      targetSchema: "public",
-      targetDatabases: ["warehouse"],
-      targetSchemas: ["public"],
-      targetTables: ["users"],
-      targetTable: "users",
-      keyColumns: ["id"],
-      label: "app → warehouse",
-    },
-    [{ sourceTable: "users", targetTable: "users" }],
-    {
-      ensureConnected: vi.fn().mockResolvedValue(undefined),
-      getConfig: () => ({ db_type: "mysql" }),
-    },
-  );
-
-  for (let attempt = 0; attempt < 20 && session.status === "running"; attempt += 1) {
-    await new Promise((resolve) => setTimeout(resolve, 0));
-  }
+test("runs a data compare session independently of the dialog and retains its result", async () => {
+  const session = await runSession(baseConfig({ selectedSourceTables: ["users"], targetTable: "users" }), [{ sourceTable: "users", targetTable: "users" }]);
 
   assert.equal(session.status, "completed");
   assert.equal(session.batchResults.length, 1);
@@ -91,6 +116,117 @@ test("runs a data compare session independently of the dialog and retains its re
   assert.equal(openMock.mock.calls.at(-1)?.[0], session.id);
 });
 
-beforeEach(() => {
-  vi.clearAllMocks();
+test("infers the primary key of every table when no match columns are configured", async () => {
+  const session = await runSession(baseConfig({ selectedSourceTables: ["users", "orders", "payments"] }), [
+    { sourceTable: "users", targetTable: "users" },
+    { sourceTable: "orders", targetTable: "orders" },
+    { sourceTable: "payments", targetTable: "payments" },
+  ]);
+
+  assert.equal(session.status, "completed");
+  assert.deepEqual(
+    session.batchResults.map((result) => result.keyColumns),
+    [["user_id"], ["tenant_id", "order_id"], ["payment_no"]],
+  );
+  assert.deepEqual(
+    apiMock.prepareDataCompareFromTables.mock.calls.map((call) => call[0].keyColumns),
+    [["user_id"], ["tenant_id", "order_id"], ["payment_no"]],
+  );
+});
+
+test("gives one table its own match columns without changing the other tables", async () => {
+  const session = await runSession(baseConfig({ selectedSourceTables: ["users", "orders", "payments"], keyColumnsByTable: { payments: ["tenant_id", "payment_no"] } }), [
+    { sourceTable: "users", targetTable: "users" },
+    { sourceTable: "orders", targetTable: "orders" },
+    { sourceTable: "payments", targetTable: "payments" },
+  ]);
+
+  assert.equal(session.status, "completed");
+  assert.deepEqual(
+    session.batchResults.map((result) => ({ table: result.sourceTable, keys: result.keyColumns })),
+    [
+      { table: "users", keys: ["user_id"] },
+      { table: "orders", keys: ["tenant_id", "order_id"] },
+      { table: "payments", keys: ["tenant_id", "payment_no"] },
+    ],
+  );
+  assert.deepEqual(keyColumnsOf(2), ["tenant_id", "payment_no"]);
+});
+
+test("keeps a manual match-column selection on one table even when another table is overridden too", async () => {
+  const session = await runSession(baseConfig({ selectedSourceTables: ["users", "orders"], keyColumnsByTable: { users: ["email"], orders: ["amount"] } }), [
+    { sourceTable: "users", targetTable: "users" },
+    { sourceTable: "orders", targetTable: "orders" },
+  ]);
+
+  assert.equal(session.status, "completed");
+  assert.deepEqual(keyColumnsOf(0), ["email"]);
+  assert.deepEqual(keyColumnsOf(1), ["amount"]);
+});
+
+test("reports a missing key column only for the table that configures it", async () => {
+  const session = await runSession(baseConfig({ selectedSourceTables: ["users", "orders"], keyColumnsByTable: { users: ["does_not_exist"] } }), [
+    { sourceTable: "users", targetTable: "users" },
+    { sourceTable: "orders", targetTable: "orders" },
+  ]);
+
+  assert.equal(session.status, "completed");
+  assert.equal(session.batchResults[0]?.status, "error");
+  assert.match(session.batchResults[0]?.error ?? "", /does_not_exist/);
+  assert.equal(session.batchResults[1]?.status, "different");
+});
+
+test("never falls back to the first column when a table has no primary key", async () => {
+  const session = await runSession(baseConfig({ selectedSourceTables: ["logs"], targetTable: "logs" }), [{ sourceTable: "logs", targetTable: "logs" }]);
+
+  assert.equal(session.batchResults.length, 1);
+  assert.equal(session.batchResults[0]?.status, "error");
+  assert.equal(apiMock.prepareDataCompareFromTables.mock.calls.length, 0);
+});
+
+test("treats an explicitly cleared selection as missing match columns instead of restoring the primary key", async () => {
+  const session = await runSession(baseConfig({ selectedSourceTables: ["users"], targetTable: "users", keyColumnsByTable: { users: [] } }), [{ sourceTable: "users", targetTable: "users" }]);
+
+  assert.equal(session.batchResults[0]?.status, "error");
+  assert.equal(apiMock.prepareDataCompareFromTables.mock.calls.length, 0);
+});
+
+test("canonicalizes a manual match column against an upper-cased target column", async () => {
+  apiMock.getColumns.mockImplementation((connectionId: string, _database: string, _schema: string, table: string) => {
+    if (table !== "legacy") return Promise.resolve(TABLES[table] ?? []);
+    return Promise.resolve(connectionId === "source" ? [column("snid", true), column("username")] : [column("SNID", true), column("USERNAME")]);
+  });
+
+  const session = await runSession(baseConfig({ selectedSourceTables: ["legacy"], targetTable: "legacy", keyColumnsByTable: { legacy: ["snid", "Username"] } }), [{ sourceTable: "legacy", targetTable: "legacy" }]);
+
+  assert.equal(session.batchResults[0]?.status, "different");
+  assert.deepEqual(keyColumnsOf(0), ["SNID", "USERNAME"]);
+  assert.deepEqual(apiMock.prepareDataCompareFromTables.mock.calls[0]?.[0].sourceColumns, ["snid", "username"]);
+  assert.deepEqual(apiMock.prepareDataCompareFromTables.mock.calls[0]?.[0].columns, ["SNID", "USERNAME"]);
+});
+
+test("uses a per-table override for a source table whose target table is missing", async () => {
+  apiMock.prepareDataCompareMissingTarget.mockResolvedValue(preparationFor("users"));
+
+  const session = await runSession(baseConfig({ selectedSourceTables: ["users"], targetTable: "users", keyColumnsByTable: { users: ["email"] }, targetTables: ["orders"] }), [{ sourceTable: "users", targetTable: "users" }]);
+
+  assert.equal(session.batchResults[0]?.status, "different");
+  assert.deepEqual(session.batchResults[0]?.keyColumns, ["email"]);
+  assert.deepEqual(apiMock.prepareDataCompareMissingTarget.mock.calls[0]?.[0].keyColumns, ["email"]);
+});
+
+test("keeps the configured per-table match columns after cloning the session config", async () => {
+  const input = baseConfig({ selectedSourceTables: ["users"], targetTable: "users", keyColumnsByTable: { users: [" user_id ", "USER_ID", "email"] } });
+  const session = await runSession(input, [{ sourceTable: "users", targetTable: "users" }]);
+
+  assert.deepEqual(session.config.keyColumnsByTable, { users: ["user_id", "email"] });
+  assert.deepEqual(keyColumnsOf(0), ["user_id", "email"]);
+  // The session must not share the caller's arrays.
+  input.keyColumnsByTable.users?.push("name");
+  assert.deepEqual(session.config.keyColumnsByTable, { users: ["user_id", "email"] });
+});
+
+test("normalizes a per-table match-column map", () => {
+  assert.deepEqual(normalizeKeyColumnOverrides({ users: [" id ", ""], orders: [], "": ["x"] }), { users: ["id"], orders: [] });
+  assert.deepEqual(normalizeKeyColumnOverrides(undefined), {});
 });

apps/desktop/src/composables/useDataCompareSession.ts

@@ -3,11 +3,18 @@ import { uuid } from "@/lib/common/utils";
 import * as api from "@/lib/backend/api";
 import { openDataCompareSession } from "@/composables/useDialogSources";
 import { useExportTracker } from "@/composables/useExportTracker";
-import { inferCompareKeyColumns, intersectCompareColumns, matchColumnNameIgnoreCase, type DataCompareCellValue, type DataCompareFromTablesPreparation, type DataCompareResult, type DataCompareSyncPlan, type DataCompareSyncPlanTableOptions } from "@/lib/dataGrid/dataCompare";
+import { inferCompareKeyColumns, intersectCompareColumns, matchColumnNameIgnoreCase, normalizeKeyColumns, type DataCompareCellValue, type DataCompareFromTablesPreparation, type DataCompareResult, type DataCompareSyncPlan, type DataCompareSyncPlanTableOptions } from "@/lib/dataGrid/dataCompare";
 import type { ColumnInfo, DatabaseType } from "@/types/database";
 
 export type CompareColumn = ColumnInfo;
 
+/**
+ * One source-table -> target-table compare task.
+ *
+ * Match columns are not part of the task itself: they are stored per source
+ * table in `DataCompareSessionConfig.keyColumnsByTable` so that every task of a
+ * batch compare can carry its own selection instead of sharing one global list.
+ */
 export interface DataCompareTableTask {
   sourceTable: string;
   targetTable: string;
@@ -75,7 +82,15 @@ export interface DataCompareSessionConfig {
   targetSchemas: string[];
   targetTables: string[];
   targetTable: string;
-  keyColumns: string[];
+  /**
+   * Explicit match columns per source table (keyed by the source table name).
+   *
+   * An absent entry means "no override": the primary key of that source table is
+   * inferred automatically. A present entry is an explicit user choice, so an
+   * empty array means the user cleared the selection and must pick columns
+   * again — it must not silently fall back to the primary key.
+   */
+  keyColumnsByTable: Record<string, string[]>;
   label: string;
 }
 
@@ -130,16 +145,36 @@ function cloneConfig(config: DataCompareSessionConfig): DataCompareSessionConfig
     targetDatabases: [...config.targetDatabases],
     targetSchemas: [...config.targetSchemas],
     targetTables: [...config.targetTables],
-    keyColumns: [...config.keyColumns],
+    keyColumnsByTable: normalizeKeyColumnOverrides(config.keyColumnsByTable),
   };
 }
 
+/**
+ * Normalizes a per-table match-column map so a cloned session never shares
+ * arrays with the caller: every entry is trimmed and de-duplicated while an
+ * explicit empty selection stays an empty array (it must not become "auto").
+ */
+export function normalizeKeyColumnOverrides(overrides: Record<string, string[] | undefined> | null | undefined): Record<string, string[]> {
+  const normalized: Record<string, string[]> = {};
+  for (const [table, columns] of Object.entries(overrides ?? {})) {
+    if (!table) continue;
+    normalized[table] = normalizeKeyColumns(columns ?? []);
+  }
+  return normalized;
+}
+
 function sessionError(error: unknown): string {
   return error instanceof Error ? error.message : (error as { message?: string } | null)?.message || String(error);
 }
 
+function fallbackCompareErrorMessage(kind: DataCompareSessionErrorKind, columns?: string): string {
+  if (kind === "missingKeyColumns") return `Missing key columns: ${columns ?? ""}`;
+  if (kind === "noCommonColumns") return "No common columns";
+  return "No key columns available";
+}
+
 function compareErrorMessage(dependencies: DataCompareSessionDependencies, kind: DataCompareSessionErrorKind, columns?: string): string {
-  return dependencies.formatError?.(kind, columns) ?? (kind === "missingKeyColumns" ? `Missing key columns: ${columns ?? ""}` : kind === "noCommonColumns" ? "No common columns" : "No key columns available");
+  return dependencies.formatError?.(kind, columns) ?? fallbackCompareErrorMessage(kind, columns);
 }
 
 function toSelectableDiff(diff: DataCompareResult): SelectableDataCompareResult {
@@ -284,6 +319,18 @@ function buildErrorResult(task: DataCompareTableTask, keyColumns: string[], data
   };
 }
 
+/**
+ * Resolves the configured match columns of one compare task.
+ *
+ * Returns the explicit per-source-table override when the user configured one
+ * (an empty array counts as "explicitly nothing selected") and `undefined`
+ * otherwise, which means the caller must infer the primary key for that table.
+ */
+function configuredTaskKeyColumns(input: DataCompareSessionConfig, task: DataCompareTableTask): string[] | undefined {
+  const override = input.keyColumnsByTable[task.sourceTable];
+  return override === undefined ? undefined : normalizeKeyColumns(override);
+}
+
 async function runDataCompareSession(session: DataCompareSession, tasks: DataCompareTableTask[], dependencies: DataCompareSessionDependencies): Promise<void> {
   const input = session.config;
   const sourceColumnCache = new Map<string, CompareColumn[]>();
@@ -297,10 +344,11 @@ async function runDataCompareSession(session: DataCompareSession, tasks: DataCom
 
     for (const [index, task] of tasks.entries()) {
       publishProgress(session, { current: index + 1, total: tasks.length, table: task.sourceTable });
-      let keyColumns = input.keyColumns;
+      const configuredKeys = configuredTaskKeyColumns(input, task);
+      let keyColumns = configuredKeys ?? [];
       try {
         if (input.targetTables.includes(task.targetTable)) {
-          const resolvedKeys = input.keyColumns.length > 0 ? input.keyColumns : await inferKeyColumnsForTable(input, task.sourceTable, sourceColumnCache);
+          const resolvedKeys = configuredKeys ?? (await inferKeyColumnsForTable(input, task.sourceTable, sourceColumnCache));
           if (resolvedKeys.length === 0) throw new Error(compareErrorMessage(dependencies, "noKeyColumns"));
 
           const sourceColumns = await loadColumnsWithCache(sourceColumnCache, input.sourceConnectionId, input.sourceDatabase, input.sourceSchema, task.sourceTable);
@@ -335,7 +383,7 @@ async function runDataCompareSession(session: DataCompareSession, tasks: DataCom
           results.push(buildTableResult(task, preparation, canonicalKeyColumns, columns, columnInfo, databaseType));
         } else {
           const sourceColumns = await loadColumnsWithCache(sourceColumnCache, input.sourceConnectionId, input.sourceDatabase, input.sourceSchema, task.sourceTable);
-          const resolvedKeys = input.keyColumns.length > 0 ? input.keyColumns : [];
+          const resolvedKeys = configuredTaskKeyColumns(input, task) ?? [];
           const sourceColumnNames = sourceColumns.map((column) => column.name);
           keyColumns = resolvedKeys.map((key) => matchColumnNameIgnoreCase(key, sourceColumnNames) ?? key);
           const preparation = await api.prepareDataCompareMissingTarget({