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

apps/desktop/src/components/diff/CompareKeyColumnsSelect.vue

@@ -0,0 +1,186 @@
+<script setup lang="ts">
+import { computed, nextTick, ref, watch } from "vue";
+import type { HTMLAttributes } from "vue";
+import { CheckSquare, ChevronDown, Search, Square } from "@lucide/vue";
+import { useI18n } from "vue-i18n";
+import { Button } from "@/components/ui/button";
+import { Input } from "@/components/ui/input";
+import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
+import { filterDatabaseOptions } from "@/lib/database/databaseOptionSearch";
+import { cn } from "@/lib/common/utils";
+import type { CompareKeyColumnOption } from "@/lib/dataGrid/dataCompare";
+
+/**
+ * Searchable multi-select for database table columns, used to configure the
+ * match columns (key columns) of one Data Compare task.
+ *
+ * It is deliberately read-only over the column metadata it receives: options
+ * come from `getColumns`, so a user can only pick columns that really exist on
+ * the table. Primary keys are labelled instead of pre-filtered, because a
+ * composite unique key is a legitimate explicit choice.
+ */
+const props = withDefaults(
+  defineProps<{
+    /** Currently selected column names. */
+    modelValue: string[];
+    /** All columns of the table, in table order. */
+    columns: CompareKeyColumnOption[];
+    disabled?: boolean;
+    /** Column metadata is still being fetched for this table. */
+    loading?: boolean;
+    /** Shown on the trigger when nothing is selected. */
+    placeholder?: string;
+    searchPlaceholder?: string;
+    emptyText?: string;
+    /** Extra classes for the popover content, mirroring SearchableSelect. */
+    contentClass?: HTMLAttributes["class"];
+  }>(),
+  {
+    disabled: false,
+    loading: false,
+    placeholder: "",
+    searchPlaceholder: "",
+    emptyText: "",
+  },
+);
+
+const emit = defineEmits<{ "update:modelValue": [value: string[]] }>();
+
+const { t } = useI18n();
+const open = ref(false);
+const searchText = ref("");
+const searchInput = ref<InstanceType<typeof Input>>();
+
+const selectedSet = computed(() => new Set(props.modelValue.map((column) => column.toLowerCase())));
+const columnNames = computed(() => props.columns.map((column) => column.name));
+const filteredColumns = computed(() => {
+  const allowed = new Set(filterDatabaseOptions(columnNames.value, searchText.value));
+  return props.columns.filter((column) => allowed.has(column.name));
+});
+const selectedLabel = computed(() => props.modelValue.join(", "));
+const selectedCountLabel = computed(() => t("dataCompare.keyColumnsSelectedCount", { selected: props.modelValue.length, total: props.columns.length }));
+
+function isSelected(column: CompareKeyColumnOption): boolean {
+  return selectedSet.value.has(column.name.toLowerCase());
+}
+
+function toggleColumn(column: CompareKeyColumnOption): void {
+  if (props.disabled) return;
+  const next = isSelected(column)
+    ? props.modelValue.filter((name) => name.toLowerCase() !== column.name.toLowerCase())
+    : // Append in click order so a composite key keeps the order the user picked.
+      [...props.modelValue, column.name];
+  emit("update:modelValue", next);
+}
+
+watch(open, async (value) => {
+  if (!value) {
+    searchText.value = "";
+    return;
+  }
+  await nextTick();
+  const input = searchInput.value?.$el as HTMLInputElement | undefined;
+  input?.focus();
+});
+</script>
+
+<template>
+  <Popover v-model:open="open">
+    <PopoverTrigger as-child>
+      <Button
+        type="button"
+        variant="outline"
+        :disabled="disabled"
+        :title="selectedLabel || placeholder"
+        :class="
+          cn(
+            'dbx-compare-key-columns-trigger dbx-control-chrome h-8 w-full min-w-0 justify-between gap-1.5 rounded-md border border-input bg-transparent px-2.5 text-xs font-normal shadow-none hover:bg-muted/40 focus-visible:border-ring focus-visible:ring-3 focus-visible:ring-ring/50 dark:bg-input/30 dark:hover:bg-input/50',
+          )
+        "
+      >
+        <span v-if="loading" class="truncate text-muted-foreground">{{ t("dataCompare.keyColumnsLoading") }}</span>
+        <span v-else-if="selectedLabel" class="truncate font-mono">{{ selectedLabel }}</span>
+        <span v-else class="truncate text-muted-foreground">{{ placeholder || t("dataCompare.keyColumnsSelectPlaceholder") }}</span>
+        <ChevronDown class="h-3.5 w-3.5 shrink-0 opacity-60" />
+      </Button>
+    </PopoverTrigger>
+    <PopoverContent align="start" :class="cn('w-[var(--reka-popover-trigger-width)] min-w-56 max-w-[min(30rem,calc(100vw-1rem))] border-0 bg-transparent p-0 shadow-none ring-0', contentClass)">
+      <div class="rounded-md border bg-popover p-1.5 shadow-md">
+        <div class="relative rounded-md border bg-background">
+          <Search class="pointer-events-none absolute left-2 top-1/2 h-3 w-3 -translate-y-1/2 text-muted-foreground" />
+          <span v-if="!searchText" class="pointer-events-none absolute left-[25px] top-1/2 -translate-y-1/2 text-xs text-muted-foreground">
+            {{ searchPlaceholder || t("dataCompare.keyColumnsSearch") }}
+          </span>
+          <Input ref="searchInput" :model-value="searchText" class="h-6 border-0 pl-6 pr-2 text-xs caret-foreground shadow-none focus-visible:ring-0" @update:model-value="(value) => (searchText = String(value))" @keydown.escape="open = false" />
+        </div>
+        <div class="dbx-compare-key-columns-list max-h-64 overflow-y-auto py-1">
+          <div v-if="loading" class="px-2 py-2 text-xs text-muted-foreground">{{ t("dataCompare.keyColumnsLoading") }}</div>
+          <template v-else-if="filteredColumns.length">
+            <button
+              v-for="column in filteredColumns"
+              :key="column.name"
+              type="button"
+              :data-column-name="column.name"
+              :data-primary-key="column.is_primary_key ? 'true' : 'false'"
+              :disabled="disabled"
+              class="flex h-7 w-full min-w-0 items-center gap-2 rounded-md px-2 text-left text-xs hover:bg-accent hover:text-accent-foreground disabled:cursor-not-allowed disabled:opacity-50"
+              @click="toggleColumn(column)"
+            >
+              <CheckSquare v-if="isSelected(column)" class="h-3.5 w-3.5 shrink-0 text-primary" />
+              <Square v-else class="h-3.5 w-3.5 shrink-0 text-muted-foreground/40" />
+              <span class="min-w-0 flex-1 truncate font-mono">{{ column.name }}</span>
+              <span v-if="column.is_primary_key" class="shrink-0 rounded bg-amber-500/15 px-1 text-[10px] font-medium text-amber-700 dark:text-amber-400">
+                {{ t("dataCompare.keyColumnsPrimaryKeyBadge") }}
+              </span>
+            </button>
+          </template>
+          <div v-else class="px-2 py-2 text-xs text-muted-foreground">
+            {{ emptyText || t("dataCompare.keyColumnsEmpty") }}
+          </div>
+        </div>
+        <div class="border-t px-2 pt-1.5 text-[11px] text-muted-foreground">{{ selectedCountLabel }}</div>
+      </div>
+    </PopoverContent>
+  </Popover>
+</template>
+
+<style>
+.dbx-compare-key-columns-list {
+  scrollbar-width: thin;
+  scrollbar-color: color-mix(in oklch, var(--foreground) 30%, transparent) transparent;
+}
+
+.dbx-compare-key-columns-list::-webkit-scrollbar {
+  width: 6px;
+  height: 6px;
+}
+
+.dbx-compare-key-columns-list::-webkit-scrollbar-track {
+  background: transparent;
+}
+
+.dbx-compare-key-columns-list::-webkit-scrollbar-thumb {
+  border: 1px solid transparent;
+  border-radius: 999px;
+  background: color-mix(in oklch, var(--foreground) 30%, transparent);
+  background-clip: padding-box;
+}
+
+.dbx-compare-key-columns-list:hover::-webkit-scrollbar-thumb {
+  border: 0;
+  background: color-mix(in oklch, var(--foreground) 48%, transparent);
+}
+
+.dark .dbx-compare-key-columns-list {
+  scrollbar-color: rgb(82, 82, 91) transparent;
+}
+
+.dark .dbx-compare-key-columns-list::-webkit-scrollbar-thumb {
+  background: rgb(82, 82, 91);
+  background-clip: padding-box;
+}
+
+.dark .dbx-compare-key-columns-list:hover::-webkit-scrollbar-thumb {
+  background: rgb(113, 113, 122);
+}
+</style>

apps/desktop/src/components/diff/DataCompareDialog.vue

@@ -3,7 +3,6 @@ import { computed, nextTick, onBeforeUnmount, ref, watch } from "vue";
 import { useI18n } from "vue-i18n";
 import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
 import { Button } from "@/components/ui/button";
-import { Input } from "@/components/ui/input";
 import { Label } from "@/components/ui/label";
 import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
 import SearchableSelect from "@/components/ui/searchable-select/SearchableSelect.vue";
@@ -13,14 +12,13 @@ import { useToast } from "@/composables/useToast";
 import { databaseOptionsForConnection, fetchNamespaceOptionsForConnection } from "@/composables/useDatabaseOptions";
 import { isSchemaAware } from "@/lib/database/databaseCapabilities";
 import { copyToClipboard } from "@/lib/common/clipboard";
-import type { DataCompareCellValue, DataCompareSyncPlan } from "@/lib/dataGrid/dataCompare";
-import { inferCompareKeyColumns } from "@/lib/dataGrid/dataCompare";
+import { inferCompareKeyColumns, normalizeKeyColumns, sameKeyColumns, type CompareKeyColumnOption, type DataCompareCellValue, type DataCompareSyncPlan } from "@/lib/dataGrid/dataCompare";
 import {
   buildDataCompareSyncPlanTables,
   emptyDataCompareSyncPlan,
   getDataCompareSession,
+  normalizeKeyColumnOverrides,
   startDataCompareSession,
-  type CompareColumn,
   type DataCompareSession,
   type DataCompareTableResult,
   type DataCompareTableStatus,
@@ -29,6 +27,7 @@ import {
   type SelectableDataCompareModifiedRow,
   type SelectableDataCompareRow,
 } from "@/composables/useDataCompareSession";
+import CompareKeyColumnsSelect from "@/components/diff/CompareKeyColumnsSelect.vue";
 import * as api from "@/lib/backend/api";
 import { executeWithProductionSqlGuard } from "@/lib/database/productionExecutionGuard";
 import TableMultiSelect from "@/components/diff/TableMultiSelect.vue";
@@ -67,7 +66,6 @@ const targetDatabases = ref<string[]>([]);
 const targetSchemas = ref<string[]>([]);
 const targetTables = ref<string[]>([]);
 
-const keyColumnsText = ref("");
 const detailPreviewLimit = ref(String(PREVIEW_LIMIT_OPTIONS[1]));
 const batchResults = ref<DataCompareTableResult[]>([]);
 const syncPlan = ref<DataCompareSyncPlan>(emptyDataCompareSyncPlan());
@@ -84,6 +82,21 @@ const showAdded = ref(true);
 const showRemoved = ref(true);
 const showModified = ref(true);
 
+/**
+ * Explicit match columns per source table. An absent entry means "auto": the
+ * primary key of that table is inferred, so a batch compare never shares one
+ * global key column list across different tables.
+ */
+const keyColumnOverrides = ref<Record<string, string[]>>({});
+/** Column metadata per source table, keyed by table name; cleared with the source endpoint. */
+const tableColumns = ref<Record<string, CompareKeyColumnOption[]>>({});
+const tableColumnErrors = ref<Record<string, string>>({});
+const loadingColumnTables = ref<string[]>([]);
+const expandedKeyColumnTable = ref("");
+let columnRequests = new Map<string, Promise<CompareKeyColumnOption[]>>();
+/** Invalidates in-flight metadata responses after a source endpoint change. */
+let columnMetadataGeneration = 0;
+
 const activeSessionId = ref<string | null>(props.sessionId ?? null);
 let syncPlanRequestId = 0;
 let initializingPrefill = false;
@@ -112,13 +125,52 @@ const compareTasksPreview = computed(() =>
 );
 const matchedTaskCount = computed(() => compareTasksPreview.value.filter((task) => task.matched).length);
 const missingTargetTables = computed(() => compareTasksPreview.value.filter((task) => !task.matched).map((task) => task.targetTable || task.sourceTable));
-const canCompare = computed(() => sourceConnectionId.value && sourceDatabase.value && sourceSchema.value && selectedSourceTableNames.value.length > 0 && targetConnectionId.value && targetDatabase.value && targetSchema.value);
-const keyColumns = computed(() =>
-  keyColumnsText.value
-    .split(",")
-    .map((value) => value.trim())
-    .filter(Boolean),
+
+interface KeyColumnRow {
+  table: string;
+  matched: boolean;
+  manual: boolean;
+  columns: string[];
+  loading: boolean;
+  hint: string;
+  /** Set when the column metadata of this table could not be loaded. */
+  error?: string;
+  /** `missing` means neither a primary key nor a manual selection is available yet. */
+  status: "auto" | "manual" | "missing";
+}
+
+const keyColumnRows = computed<KeyColumnRow[]>(() =>
+  selectedSourceTableNames.value.map((table) => {
+    const matched = !!compareTasksPreview.value.find((task) => task.sourceTable === table)?.matched;
+    const manual = hasKeyColumnOverride(table);
+    const columns = effectiveKeyColumns(table);
+    const error = tableColumnErrors.value[table];
+    return {
+      table,
+      matched,
+      manual,
+      columns,
+      loading: isTableColumnsLoading(table),
+      hint: error || keyColumnHint(table, { matched, manual, columns }),
+      error,
+      status: manual ? "manual" : columns.length > 0 ? "auto" : "missing",
+    };
+  }),
 );
+const singleKeyColumnRow = computed(() => (selectedSourceTableNames.value.length === 1 ? keyColumnRows.value[0] : undefined));
+/**
+ * A single-table compare whose target table exists but whose match columns are
+ * unknown: primary key inference found nothing and the user has not picked
+ * columns yet. Blocking here is safe because no other table depends on it.
+ * Unloaded metadata is deliberately not treated as "no primary key".
+ */
+const singleTableNeedsKeyColumns = computed(() => {
+  const row = singleKeyColumnRow.value;
+  if (!row || !row.matched) return false;
+  if (!isTableColumnsKnown(row.table)) return false;
+  return row.columns.length === 0;
+});
+const canCompare = computed(() => sourceConnectionId.value && sourceDatabase.value && sourceSchema.value && selectedSourceTableNames.value.length > 0 && targetConnectionId.value && targetDatabase.value && targetSchema.value && !singleTableNeedsKeyColumns.value);
 const detailPreviewLimitNumber = computed(() => Number(detailPreviewLimit.value) || PREVIEW_LIMIT_OPTIONS[1]);
 
 const sameTableCount = computed(() => batchResults.value.filter((item) => item.status === "same").length);
@@ -172,6 +224,157 @@ function resetSelectedSourceTables(nextTables: Iterable<string>) {
   selectedSourceTables.value = new Set(nextTables);
 }
 
+// --- Match columns (key columns) -------------------------------------------------
+//
+// Every selected source table owns its match columns: the primary key is the
+// default, and a manual picker selection is stored as a per-table override in
+// `keyColumnOverrides`. The session receives the same per-table shape, so a
+// batch compare is never forced to share one global key column list.
+
+function resetTableColumnMetadata() {
+  columnMetadataGeneration += 1;
+  columnRequests = new Map();
+  tableColumns.value = {};
+  tableColumnErrors.value = {};
+  loadingColumnTables.value = [];
+  expandedKeyColumnTable.value = "";
+}
+
+function columnsForTable(table: string): CompareKeyColumnOption[] {
+  return tableColumns.value[table] ?? [];
+}
+
+function isTableColumnsKnown(table: string): boolean {
+  return tableColumns.value[table] !== undefined;
+}
+
+function isTableColumnsLoading(table: string): boolean {
+  return loadingColumnTables.value.includes(table);
+}
+
+function hasKeyColumnOverride(table: string): boolean {
+  return keyColumnOverrides.value[table] !== undefined;
+}
+
+function inferredKeyColumns(table: string): string[] {
+  return inferCompareKeyColumns(columnsForTable(table));
+}
+
+/** Effective match columns of a table: the manual override, otherwise the inferred primary key. */
+function effectiveKeyColumns(table: string): string[] {
+  const override = keyColumnOverrides.value[table];
+  return override === undefined ? inferredKeyColumns(table) : override;
+}
+
+async function fetchTableColumns(table: string): Promise<CompareKeyColumnOption[]> {
+  const connectionId = sourceConnectionId.value;
+  const database = sourceDatabase.value;
+  const schema = sourceSchema.value;
+  if (!connectionId || !database || !schema || !table) return [];
+  const generation = columnMetadataGeneration;
+  loadingColumnTables.value = [...loadingColumnTables.value, table];
+  try {
+    const columns = (await api.getColumns(connectionId, database, schema, table)) as CompareKeyColumnOption[];
+    if (generation !== columnMetadataGeneration) return [];
+    tableColumns.value = { ...tableColumns.value, [table]: columns };
+    delete tableColumnErrors.value[table];
+    return columns;
+  } catch (error: any) {
+    if (generation === columnMetadataGeneration) tableColumnErrors.value = { ...tableColumnErrors.value, [table]: error?.message || String(error) };
+    return [];
+  } finally {
+    if (generation === columnMetadataGeneration) loadingColumnTables.value = loadingColumnTables.value.filter((item) => item !== table);
+  }
+}
+
+function loadTableColumns(table: string): Promise<CompareKeyColumnOption[]> {
+  const cached = tableColumns.value[table];
+  if (cached) return Promise.resolve(cached);
+  const pending = columnRequests.get(table);
+  if (pending) return pending;
+  const request = fetchTableColumns(table).finally(() => columnRequests.delete(table));
+  columnRequests.set(table, request);
+  return request;
+}
+
+/** Loads metadata for the selected tables with bounded concurrency, so a wide selection cannot flood the backend. */
+async function prefetchSelectedTableColumns(tables: string[]): Promise<void> {
+  const pending = tables.filter((table) => !tableColumns.value[table]);
+  if (pending.length === 0) return;
+  let cursor = 0;
+  const workers = Array.from({ length: Math.min(4, pending.length) }, async () => {
+    while (cursor < pending.length) await loadTableColumns(pending[cursor++]);
+  });
+  await Promise.all(workers);
+}
+
+function setTableKeyColumns(table: string, columns: string[]) {
+  const normalized = normalizeKeyColumns(columns);
+  const next = { ...keyColumnOverrides.value };
+  if (sameKeyColumns(normalized, inferredKeyColumns(table))) delete next[table];
+  else next[table] = normalized;
+  keyColumnOverrides.value = next;
+}
+
+function clearTableKeyColumnsOverride(table: string) {
+  if (keyColumnOverrides.value[table] === undefined) return;
+  const next = { ...keyColumnOverrides.value };
+  delete next[table];
+  keyColumnOverrides.value = next;
+}
+
+/** Drops overrides and metadata of tables that are no longer selected. */
+function pruneKeyColumnState(tables: string[]) {
+  const keep = new Set(tables);
+  const nextOverrides: Record<string, string[]> = {};
+  let overridesChanged = false;
+  for (const [table, columns] of Object.entries(keyColumnOverrides.value)) {
+    if (keep.has(table)) nextOverrides[table] = columns;
+    else overridesChanged = true;
+  }
+  if (overridesChanged) keyColumnOverrides.value = nextOverrides;
+
+  const nextColumns: Record<string, CompareKeyColumnOption[]> = {};
+  const nextErrors: Record<string, string> = {};
+  let columnsChanged = false;
+  for (const [table, columns] of Object.entries(tableColumns.value)) {
+    if (keep.has(table)) {
+      nextColumns[table] = columns;
+      if (tableColumnErrors.value[table]) nextErrors[table] = tableColumnErrors.value[table];
+    } else {
+      columnsChanged = true;
+    }
+  }
+  if (columnsChanged) {
+    tableColumns.value = nextColumns;
+    tableColumnErrors.value = nextErrors;
+  }
+  if (expandedKeyColumnTable.value && !keep.has(expandedKeyColumnTable.value)) expandedKeyColumnTable.value = "";
+}
+
+function toggleKeyColumnTable(table: string) {
+  expandedKeyColumnTable.value = expandedKeyColumnTable.value === table ? "" : table;
+  if (expandedKeyColumnTable.value) void loadTableColumns(table);
+}
+
+function keyColumnHint(table: string, state: { matched: boolean; manual: boolean; columns: string[] }): string {
+  if (!state.matched) return t("dataCompare.keyColumnsTargetMissingHint");
+  if (state.manual) return t("dataCompare.keyColumnsManualHint", { columns: state.columns.length > 0 ? state.columns.join(", ") : t("dataCompare.keyColumnsNoneSelected") });
+  if (state.columns.length > 0) return t("dataCompare.keyColumnsAutoSummary", { columns: state.columns.join(", ") });
+  if (!isTableColumnsKnown(table)) return t("dataCompare.keyColumnsStatusAuto");
+  return t("dataCompare.keyColumnsNoPrimaryKey");
+}
+
+/** Only the selected tables are sent, so a deselected table can never leak a stale override. */
+function buildSessionKeyColumnsByTable(): Record<string, string[]> {
+  const overrides: Record<string, string[]> = {};
+  for (const table of selectedSourceTableNames.value) {
+    const override = keyColumnOverrides.value[table];
+    if (override !== undefined) overrides[table] = [...override];
+  }
+  return overrides;
+}
+
 function buildCompareTasks(): DataCompareTableTask[] {
   if (!selectedSourceTableNames.value.length) return [];
   if (!isBatchCompare.value) {
@@ -222,7 +425,10 @@ async function restoreDataCompareSession(session: DataCompareSession): Promise<v
     targetSchemas.value = [...config.targetSchemas];
     targetTables.value = [...config.targetTables];
     targetTable.value = config.targetTable;
-    keyColumnsText.value = config.keyColumns.join(", ");
+    // Restore the per-table match columns so a reopened session keeps every
+    // table's own selection instead of collapsing back to one global list.
+    keyColumnOverrides.value = normalizeKeyColumnOverrides(config.keyColumnsByTable);
+    resetTableColumnMetadata();
     batchResults.value = session.batchResults;
     syncPlan.value = session.syncPlan;
     syncErrors.value = [];
@@ -230,6 +436,7 @@ async function restoreDataCompareSession(session: DataCompareSession): Promise<v
     compareProgressTotal.value = session.progress?.total ?? 0;
     compareProgressTable.value = session.progress?.table ?? "";
     comparing.value = session.status === "running";
+    void prefetchSelectedTableColumns(selectedSourceTableNames.value);
   } finally {
     await nextTick();
     if (generation === initializingPrefillGeneration) initializingPrefill = false;
@@ -377,28 +584,6 @@ async function loadTables(side: "source" | "target") {
   }
 }
 
-async function loadColumnsWithCache(cache: Map<string, CompareColumn[]>, connectionId: string, database: string, schema: string, table: string): Promise<CompareColumn[]> {
-  const key = `${connectionId}:${database}:${schema}:${table}`;
-  const cached = cache.get(key);
-  if (cached) return cached;
-  const columns = (await api.getColumns(connectionId, database, schema, table)) as CompareColumn[];
-  cache.set(key, columns);
-  return columns;
-}
-
-async function inferKeyColumnsForTable(table: string, sourceColumnCache?: Map<string, CompareColumn[]>): Promise<string[]> {
-  if (!sourceConnectionId.value || !sourceDatabase.value || !sourceSchema.value || !table) return [];
-  const columns = sourceColumnCache ? await loadColumnsWithCache(sourceColumnCache, sourceConnectionId.value, sourceDatabase.value, sourceSchema.value, table) : (((await api.getColumns(sourceConnectionId.value, sourceDatabase.value, sourceSchema.value, table)) as CompareColumn[]) ?? []);
-  return inferCompareKeyColumns(columns);
-}
-
-async function inferKeyColumns() {
-  const table = selectedSourceTableNames.value.length === 1 ? selectedSourceTableNames.value[0] : "";
-  if (!table) return;
-  const inferred = await inferKeyColumnsForTable(table);
-  keyColumnsText.value = inferred.join(", ");
-}
-
 function resultStatusLabel(status: DataCompareTableStatus): string {
   if (status === "different") return t("dataCompare.statusDifferent");
   if (status === "same") return t("dataCompare.statusSame");
@@ -535,7 +720,7 @@ function startCompare(): void {
       targetSchemas: [...targetSchemas.value],
       targetTables: [...targetTables.value],
       targetTable: targetTable.value,
-      keyColumns: [...keyColumns.value],
+      keyColumnsByTable: buildSessionKeyColumnsByTable(),
       label: `${comparisonEndpointLabel(sourceConnectionId.value, sourceDatabase.value, sourceSchema.value)} → ${comparisonEndpointLabel(targetConnectionId.value, targetDatabase.value, targetSchema.value)}`,
     },
     tasks,
@@ -651,6 +836,8 @@ watch(sourceConnectionId, (id) => {
   sourceTables.value = [];
   sourceTable.value = "";
   resetSelectedSourceTables([]);
+  // Column metadata and match columns belong to the previous source endpoint.
+  resetTableColumnMetadata();
   loadDatabases(id, "source").catch((e) => toast(String(e), 5000));
 });
 watch(targetConnectionId, (id) => {
@@ -671,6 +858,7 @@ watch(sourceDatabase, () => {
   sourceTables.value = [];
   sourceTable.value = "";
   resetSelectedSourceTables([]);
+  resetTableColumnMetadata();
   loadSchemas("source", props.prefillSchema).catch((e) => toast(String(e), 5000));
 });
 watch(targetDatabase, () => {
@@ -688,6 +876,7 @@ watch(sourceSchema, () => {
   sourceTables.value = [];
   sourceTable.value = "";
   resetSelectedSourceTables([]);
+  resetTableColumnMetadata();
   if (sourceSchema.value) loadTables("source").catch((e) => toast(String(e), 5000));
 });
 watch(targetSchema, () => {
@@ -700,9 +889,12 @@ watch(targetSchema, () => {
 watch(selectedSourceTableNames, (tables, previous) => {
   if (initializingPrefill) return;
   clearResult();
+  // Drop the match columns and metadata of tables that are no longer compared,
+  // so a table that is selected again starts from its own primary key.
+  pruneKeyColumnState(tables);
   sourceTable.value = tables.length === 1 ? tables[0] : "";
   if (tables.length !== 1) {
-    keyColumnsText.value = "";
+    void prefetchSelectedTableColumns(tables);
     return;
   }
   const table = tables[0];
@@ -711,10 +903,7 @@ watch(selectedSourceTableNames, (tables, previous) => {
   } else if (previous?.length === 1 && targetTable.value === previous[0]) {
     targetTable.value = "";
   }
-  if (table !== previous?.[0]) {
-    keyColumnsText.value = "";
-  }
-  inferKeyColumns().catch(() => {});
+  void prefetchSelectedTableColumns(tables);
 });
 watch(targetTable, () => {
   if (initializingPrefill) return;
@@ -748,6 +937,7 @@ watch(
           if (sourceTables.value.includes(props.prefillTable)) {
             resetSelectedSourceTables([props.prefillTable]);
             sourceTable.value = props.prefillTable;
+            await loadTableColumns(props.prefillTable);
           }
         }
       } finally {
@@ -917,10 +1107,59 @@ onBeforeUnmount(() => {
         </div>
 
         <div class="space-y-1">
-          <Label class="text-xs font-medium">{{ t("dataCompare.keyColumns") }}</Label>
-          <Input v-model="keyColumnsText" class="h-8 text-xs" :placeholder="t('dataCompare.keyColumnsPlaceholder')" :disabled="comparing" />
-          <div class="text-[11px] text-muted-foreground">
-            {{ t("dataCompare.keyColumnsAutoHint") }}
+          <div class="flex flex-wrap items-center gap-2">
+            <Label class="text-xs font-medium">{{ t("dataCompare.keyColumns") }}</Label>
+            <span v-if="singleKeyColumnRow" class="inline-flex items-center rounded px-1.5 py-0.5 text-[10px] font-medium" :class="singleKeyColumnRow.manual ? 'bg-primary/15 text-primary' : 'bg-muted text-muted-foreground'">
+              {{ singleKeyColumnRow.manual ? t("dataCompare.keyColumnsStatusManual") : t("dataCompare.keyColumnsStatusAuto") }}
+            </span>
+            <Button v-if="singleKeyColumnRow?.manual" size="sm" variant="ghost" class="h-6 px-1.5 text-[11px]" :disabled="comparing" @click="clearTableKeyColumnsOverride(singleKeyColumnRow.table)">
+              {{ t("dataCompare.keyColumnsResetAuto") }}
+            </Button>
+          </div>
+
+          <div v-if="!selectedSourceTableNames.length" class="text-[11px] text-muted-foreground">
+            {{ t("dataCompare.keyColumnsSelectTableHint") }}
+          </div>
+
+          <!-- Single table: pick match columns straight from the real database columns. -->
+          <div v-else-if="singleKeyColumnRow" class="space-y-1">
+            <CompareKeyColumnsSelect :model-value="singleKeyColumnRow.columns" :columns="columnsForTable(singleKeyColumnRow.table)" :loading="singleKeyColumnRow.loading" :disabled="comparing" @update:model-value="(columns) => setTableKeyColumns(singleKeyColumnRow!.table, columns)" />
+            <div class="text-[11px]" data-key-column-hint :class="singleKeyColumnRow.error || singleKeyColumnRow.status === 'missing' ? 'text-destructive' : 'text-muted-foreground'">
+              {{ singleKeyColumnRow.hint }}
+            </div>
+          </div>
+
+          <!-- Batch compare: every table keeps its own match columns. -->
+          <div v-else class="space-y-1">
+            <div class="rounded-lg border">
+              <div class="border-b bg-muted/20 px-2 py-1 text-[11px] text-muted-foreground">
+                {{ t("dataCompare.keyColumnsPerTableHint") }}
+              </div>
+              <div class="max-h-56 divide-y overflow-auto">
+                <div v-for="row in keyColumnRows" :key="row.table">
+                  <button type="button" class="dbx-compare-key-table-row flex w-full items-center gap-2 px-2 py-1.5 text-left text-xs hover:bg-muted/40" :data-key-column-table="row.table" @click="toggleKeyColumnTable(row.table)">
+                    <ChevronDown v-if="expandedKeyColumnTable === row.table" class="h-3.5 w-3.5 shrink-0" />
+                    <ChevronRight v-else class="h-3.5 w-3.5 shrink-0" />
+                    <span class="min-w-0 flex-1 truncate font-mono">{{ row.table }}</span>
+                    <span class="min-w-0 max-w-[45%] shrink-0 truncate font-mono" :class="row.columns.length > 0 ? (row.manual ? 'text-foreground' : 'text-muted-foreground') : 'text-destructive'">
+                      {{ row.loading ? t("dataCompare.keyColumnsLoading") : row.columns.length > 0 ? row.columns.join(", ") : t("dataCompare.keyColumnsNeedsSelection") }}
+                    </span>
+                    <span class="shrink-0 rounded px-1.5 py-0.5 text-[10px] font-medium" :class="row.status === 'manual' ? 'bg-primary/15 text-primary' : row.status === 'auto' ? 'bg-muted text-muted-foreground' : 'bg-destructive/15 text-destructive'">
+                      {{ row.status === "manual" ? t("dataCompare.keyColumnsStatusManual") : row.status === "auto" ? t("dataCompare.keyColumnsStatusAuto") : t("dataCompare.keyColumnsNeedsSelection") }}
+                    </span>
+                  </button>
+                  <div v-if="expandedKeyColumnTable === row.table" class="space-y-1 border-t bg-muted/10 px-2 py-2">
+                    <CompareKeyColumnsSelect :model-value="row.columns" :columns="columnsForTable(row.table)" :loading="row.loading" :disabled="comparing" @update:model-value="(columns) => setTableKeyColumns(row.table, columns)" />
+                    <div class="flex flex-wrap items-center gap-2 text-[11px]">
+                      <span data-key-column-hint :class="row.error || (row.status === 'missing' && row.matched) ? 'text-destructive' : 'text-muted-foreground'">{{ row.hint }}</span>
+                      <Button v-if="row.manual" size="sm" variant="ghost" class="h-6 px-1.5 text-[11px]" :disabled="comparing" @click="clearTableKeyColumnsOverride(row.table)">
+                        {{ t("dataCompare.keyColumnsResetAuto") }}
+                      </Button>
+                    </div>
+                  </div>
+                </div>
+              </div>
+            </div>
           </div>
         </div>
 

apps/desktop/src/components/diff/TableMultiSelect.vue

@@ -79,7 +79,7 @@ function toggleSelectAll() {
       {{ emptyText || t("tableMultiSelect.noTables") }}
     </div>
     <div v-else class="max-h-40 overflow-auto rounded border">
-      <button v-for="table in filteredTables" :key="table" type="button" :disabled="disabled" class="flex w-full items-center gap-2 px-2.5 py-1.5 text-left text-xs hover:bg-muted/50 disabled:opacity-50" @click="toggle(table)">
+      <button v-for="table in filteredTables" :key="table" type="button" :data-table-name="table" :disabled="disabled" class="flex w-full items-center gap-2 px-2.5 py-1.5 text-left text-xs hover:bg-muted/50 disabled:opacity-50" @click="toggle(table)">
         <CheckSquare v-if="selectedSet.has(table)" class="h-3.5 w-3.5 shrink-0 text-primary" />
         <Square v-else class="h-3.5 w-3.5 shrink-0 text-muted-foreground/40" />
         <span class="truncate">{{ table }}</span>

apps/desktop/src/components/diff/__tests__/DataCompareDialog.keyColumns.spec.ts

@@ -0,0 +1,399 @@
+// @vitest-environment happy-dom
+
+import { createApp, defineComponent, h, nextTick, type App } from "vue";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import i18n from "@/i18n";
+import DataCompareDialog from "@/components/diff/DataCompareDialog.vue";
+import type { DataCompareSession } from "@/composables/useDataCompareSession";
+
+const SOURCE_CONNECTION_ID = "oracle-11g";
+const TARGET_CONNECTION_ID = "oracle-jdbc-11g";
+
+interface MockColumn {
+  name: string;
+  data_type: string;
+  is_primary_key: boolean;
+}
+
+function column(name: string, isPrimaryKey = false): MockColumn {
+  return { name, data_type: "NUMBER", is_primary_key: isPrimaryKey };
+}
+
+const TABLES: Record<string, { source: MockColumn[]; target: MockColumn[] }> = {
+  USERS: {
+    source: [column("USER_ID", true), column("EMAIL"), column("NAME")],
+    target: [column("USER_ID", true), column("EMAIL"), column("NAME")],
+  },
+  ORDERS: {
+    source: [column("TENANT_ID", true), column("ORDER_ID", true), column("ORDER_NO"), column("AMOUNT")],
+    target: [column("TENANT_ID", true), column("ORDER_ID", true), column("ORDER_NO"), column("AMOUNT")],
+  },
+  PAYMENTS: {
+    source: [column("TENANT_ID"), column("PAYMENT_NO")],
+    target: [column("TENANT_ID"), column("PAYMENT_NO")],
+  },
+  CASES: {
+    source: [column("snid", true), column("username")],
+    target: [column("SNID", true), column("USERNAME")],
+  },
+};
+
+const TABLE_NAMES = Object.keys(TABLES);
+
+const mocks = vi.hoisted(() => ({
+  ensureConnected: vi.fn().mockResolvedValue(undefined),
+  listDatabases: vi.fn().mockResolvedValue([]),
+  listSchemas: vi.fn(),
+  listTables: vi.fn(),
+  getColumns: vi.fn(),
+  buildDataCompareSyncPlan: vi.fn(),
+  executeBatch: vi.fn(),
+  prepareDataCompareFromTables: vi.fn(),
+  prepareDataCompareMissingTarget: vi.fn(),
+}));
+
+const sessionMocks = vi.hoisted(() => ({
+  session: null as DataCompareSession | null,
+}));
+
+vi.mock("@/composables/useDataCompareSession", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("@/composables/useDataCompareSession")>();
+  return {
+    ...actual,
+    getDataCompareSession: (id: string | null | undefined) => (id && sessionMocks.session?.id === id ? sessionMocks.session : undefined),
+  };
+});
+
+vi.mock("@/stores/connectionStore", () => {
+  const connections = [
+    { id: "oracle-11g", name: "Oracle XE 11g", db_type: "oracle", driver_profile: "oracle", database: "XE" },
+    {
+      id: "oracle-jdbc-11g",
+      name: "Oracle JDBC 11g",
+      db_type: "jdbc",
+      driver_profile: "oracle",
+      connection_string: "jdbc:oracle:thin:@//localhost:1521/XE",
+      jdbc_driver_class: "oracle.jdbc.OracleDriver",
+    },
+  ];
+  return {
+    useConnectionStore: () => ({
+      connections,
+      sidebarLayout: {
+        groups: [{ id: "oracle", name: "Oracle", collapsed: false }],
+        order: [{ type: "group", id: "oracle", children: connections.map((connection) => ({ type: "connection", id: connection.id })) }],
+      },
+      getConfig: (id: string) => connections.find((connection) => connection.id === id),
+      ensureConnected: mocks.ensureConnected,
+    }),
+  };
+});
+
+vi.mock("@/lib/backend/api", () => ({
+  listDatabases: mocks.listDatabases,
+  listSchemas: mocks.listSchemas,
+  listTables: mocks.listTables,
+  getColumns: mocks.getColumns,
+  buildDataCompareSyncPlan: mocks.buildDataCompareSyncPlan,
+  executeBatch: mocks.executeBatch,
+  prepareDataCompareFromTables: mocks.prepareDataCompareFromTables,
+  prepareDataCompareMissingTarget: mocks.prepareDataCompareMissingTarget,
+}));
+
+const mountedApps: App[] = [];
+
+async function flushAsyncSetup() {
+  for (let index = 0; index < 12; index += 1) {
+    await nextTick();
+    await new Promise((resolve) => setTimeout(resolve, 0));
+  }
+}
+
+afterEach(() => {
+  for (const app of mountedApps.splice(0)) app.unmount();
+  document.body.textContent = "";
+  sessionMocks.session = null;
+  vi.clearAllMocks();
+});
+
+function mountDialog(props: Record<string, unknown>): App {
+  const container = document.createElement("div");
+  document.body.append(container);
+  const app = createApp(
+    defineComponent({
+      setup: () => () => h(DataCompareDialog, { open: true, ...props }),
+    }),
+  );
+  mountedApps.push(app);
+  app.use(i18n);
+  app.mount(container);
+  return app;
+}
+
+/** The dialog prefills only the source side; the target must be selected for a task to match. */
+async function selectMatchedTarget(): Promise<void> {
+  const targetConnectionTrigger = document.querySelectorAll<HTMLButtonElement>("button.dbx-diff-connection-trigger")[1];
+  targetConnectionTrigger?.click();
+  await flushAsyncSetup();
+  document.querySelector<HTMLButtonElement>(`[data-picker-connection="${TARGET_CONNECTION_ID}"]`)?.click();
+  await flushAsyncSetup();
+
+  const targetDatabaseTrigger = document.querySelectorAll<HTMLButtonElement>("button.dbx-searchable-select-trigger")[2];
+  targetDatabaseTrigger?.click();
+  await flushAsyncSetup();
+  const reporting = [...document.querySelectorAll<HTMLButtonElement>(".dbx-searchable-select-list button")].find((button) => button.textContent?.trim() === "REPORTING");
+  reporting?.click();
+  await flushAsyncSetup();
+}
+
+async function openWithPrefilledSource(table: string, options: { matchedTarget?: boolean } = {}) {
+  mountDialog({ prefillConnectionId: SOURCE_CONNECTION_ID, prefillDatabase: "DBX_TEST", prefillSchema: "DBX_TEST", prefillTable: table });
+  await flushAsyncSetup();
+  if (options.matchedTarget !== false) await selectMatchedTarget();
+}
+
+function keyColumnTrigger(): HTMLButtonElement | undefined {
+  return document.querySelector<HTMLButtonElement>("button.dbx-compare-key-columns-trigger") ?? undefined;
+}
+
+function keyColumnHint(): string {
+  return document.querySelector<HTMLElement>("[data-key-column-hint]")?.textContent?.trim() ?? "";
+}
+
+function keyColumnRow(table: string): HTMLButtonElement | undefined {
+  return document.querySelector<HTMLButtonElement>(`button.dbx-compare-key-table-row[data-key-column-table="${table}"]`) ?? undefined;
+}
+
+function tableSelectionRow(table: string): HTMLButtonElement | undefined {
+  return document.querySelector<HTMLButtonElement>(`[data-table-name="${table}"]`) ?? undefined;
+}
+
+function buttonWithText(text: string): HTMLButtonElement | undefined {
+  return [...document.querySelectorAll<HTMLButtonElement>("button")].find((button) => button.textContent?.includes(text)) ?? undefined;
+}
+
+function compareButton(): HTMLButtonElement | undefined {
+  return buttonWithText(i18n.global.t("dataCompare.compare"));
+}
+
+async function togglePickerColumn(columnName: string): Promise<void> {
+  // The popover stays open between toggles, so only open it when it is closed.
+  if (!document.querySelector(".dbx-compare-key-columns-list")) {
+    keyColumnTrigger()?.click();
+    await flushAsyncSetup();
+  }
+  const option = document.querySelector<HTMLButtonElement>(`.dbx-compare-key-columns-list button[data-column-name="${columnName}"]`);
+  expect(option).toBeDefined();
+  option?.click();
+  await flushAsyncSetup();
+}
+
+function comparedKeyColumns(): { table: string; keys: string[] }[] {
+  return mocks.prepareDataCompareFromTables.mock.calls.map((call) => ({
+    table: (call[0] as { targetTable: string }).targetTable,
+    keys: (call[0] as { keyColumns: string[] }).keyColumns,
+  }));
+}
+
+function preparationFor(table: string) {
+  return {
+    result: { added: [{ key: "1", keyValues: { USER_ID: 1 }, values: { USER_ID: 1, EMAIL: "a@b.c" } }], removed: [], modified: [] },
+    syncStatements: [`INSERT INTO ${table} (USER_ID) VALUES (1)`],
+    syncSql: `INSERT INTO ${table} (USER_ID) VALUES (1)`,
+    preSyncStatements: [],
+    sourceRowCount: 1,
+    targetRowCount: 0,
+    sourceTruncated: false,
+    targetTruncated: false,
+  };
+}
+
+function configuredSession(): DataCompareSession {
+  return {
+    id: "compare-session",
+    version: 1,
+    status: "completed",
+    config: {
+      sourceConnectionId: SOURCE_CONNECTION_ID,
+      sourceDatabase: "DBX_TEST",
+      sourceSchema: "DBX_TEST",
+      sourceDatabases: ["DBX_TEST"],
+      sourceSchemas: ["DBX_TEST"],
+      sourceTables: TABLE_NAMES,
+      selectedSourceTables: ["USERS", "ORDERS"],
+      targetConnectionId: TARGET_CONNECTION_ID,
+      targetDatabase: "REPORTING",
+      targetSchema: "DBX_TEST",
+      targetDatabases: ["REPORTING"],
+      targetSchemas: ["DBX_TEST"],
+      targetTables: TABLE_NAMES,
+      targetTable: "",
+      keyColumnsByTable: { ORDERS: ["ORDER_NO"] },
+      label: "Oracle source → Oracle target",
+    },
+    progress: null,
+    batchResults: [],
+    syncPlan: { insertCount: 0, updateCount: 0, deleteCount: 0, statementCount: 0, syncStatements: [], syncSql: "" },
+    error: null,
+    startedAt: 1,
+    finishedAt: 2,
+  };
+}
+
+function setupMocks() {
+  mocks.listSchemas.mockResolvedValue(["DBX_TEST", "REPORTING", "SYS"]);
+  mocks.listTables.mockResolvedValue(TABLE_NAMES.map((name) => ({ name, table_type: "TABLE" })));
+  mocks.getColumns.mockImplementation((connectionId: string, _database: string, _schema: string, table: string) => {
+    const entry = TABLES[table];
+    if (!entry) return Promise.resolve([]);
+    return Promise.resolve(connectionId === SOURCE_CONNECTION_ID ? entry.source : entry.target);
+  });
+  mocks.prepareDataCompareFromTables.mockImplementation((options: { targetTable: string }) => Promise.resolve(preparationFor(options.targetTable)));
+  mocks.prepareDataCompareMissingTarget.mockImplementation((options: { sourceTable: string }) => Promise.resolve(preparationFor(options.sourceTable)));
+  mocks.buildDataCompareSyncPlan.mockResolvedValue({ insertCount: 1, updateCount: 0, deleteCount: 0, statementCount: 1, syncStatements: ["INSERT"], syncSql: "INSERT" });
+}
+
+describe("DataCompareDialog match columns", () => {
+  it("auto-selects the single primary key of a single-table compare and labels it", async () => {
+    setupMocks();
+    await openWithPrefilledSource("USERS");
+
+    expect(keyColumnTrigger()?.textContent?.trim()).toBe("USER_ID");
+    expect(keyColumnHint()).toBe(i18n.global.t("dataCompare.keyColumnsAutoSummary", { columns: "USER_ID" }));
+    expect(document.body.textContent).toContain(i18n.global.t("dataCompare.keyColumnsStatusAuto"));
+    expect(compareButton()?.disabled).toBe(false);
+  });
+
+  it("keeps the declared order of a composite primary key", async () => {
+    setupMocks();
+    await openWithPrefilledSource("ORDERS");
+
+    expect(keyColumnTrigger()?.textContent?.trim()).toBe("TENANT_ID, ORDER_ID");
+    expect(keyColumnHint()).toBe(i18n.global.t("dataCompare.keyColumnsAutoSummary", { columns: "TENANT_ID, ORDER_ID" }));
+  });
+
+  it("offers the real database columns with a primary-key badge", async () => {
+    setupMocks();
+    await openWithPrefilledSource("ORDERS");
+
+    keyColumnTrigger()?.click();
+    await flushAsyncSetup();
+
+    const options = [...document.querySelectorAll<HTMLButtonElement>(".dbx-compare-key-columns-list button")];
+    expect(options.map((option) => option.dataset.columnName)).toEqual(["TENANT_ID", "ORDER_ID", "ORDER_NO", "AMOUNT"]);
+    expect(options.map((option) => option.dataset.primaryKey)).toEqual(["true", "true", "false", "false"]);
+    expect(document.body.textContent).toContain(i18n.global.t("dataCompare.keyColumnsSelectedCount", { selected: 2, total: 4 }));
+  });
+
+  it("requires an explicit selection when the table has no primary key", async () => {
+    setupMocks();
+    await openWithPrefilledSource("PAYMENTS");
+
+    expect(keyColumnTrigger()?.textContent?.trim()).toBe(i18n.global.t("dataCompare.keyColumnsSelectPlaceholder"));
+    expect(keyColumnHint()).toBe(i18n.global.t("dataCompare.keyColumnsNoPrimaryKey"));
+    expect(compareButton()?.disabled).toBe(true);
+
+    await togglePickerColumn("TENANT_ID");
+
+    expect(keyColumnTrigger()?.textContent?.trim()).toBe("TENANT_ID");
+    expect(compareButton()?.disabled).toBe(false);
+    expect(document.body.textContent).toContain(i18n.global.t("dataCompare.keyColumnsStatusManual"));
+  });
+
+  it("lets the user override the inferred primary key and reset back to it", async () => {
+    setupMocks();
+    await openWithPrefilledSource("ORDERS");
+
+    await togglePickerColumn("TENANT_ID");
+    await togglePickerColumn("ORDER_ID");
+    await togglePickerColumn("ORDER_NO");
+
+    expect(keyColumnTrigger()?.textContent?.trim()).toBe("ORDER_NO");
+    expect(keyColumnHint()).toBe(i18n.global.t("dataCompare.keyColumnsManualHint", { columns: "ORDER_NO" }));
+
+    buttonWithText(i18n.global.t("dataCompare.keyColumnsResetAuto"))?.click();
+    await flushAsyncSetup();
+
+    expect(keyColumnTrigger()?.textContent?.trim()).toBe("TENANT_ID, ORDER_ID");
+  });
+
+  it("configures match columns per table in a batch compare", async () => {
+    setupMocks();
+    await openWithPrefilledSource("USERS");
+
+    tableSelectionRow("ORDERS")?.click();
+    await flushAsyncSetup();
+
+    expect(keyColumnRow("USERS")?.textContent).toContain("USER_ID");
+    expect(keyColumnRow("ORDERS")?.textContent).toContain("TENANT_ID, ORDER_ID");
+    expect(keyColumnRow("USERS")?.textContent).toContain(i18n.global.t("dataCompare.keyColumnsStatusAuto"));
+    expect(keyColumnRow("ORDERS")?.textContent).toContain(i18n.global.t("dataCompare.keyColumnsStatusAuto"));
+
+    keyColumnRow("ORDERS")?.click();
+    await flushAsyncSetup();
+    await togglePickerColumn("TENANT_ID");
+    await togglePickerColumn("ORDER_ID");
+    await togglePickerColumn("ORDER_NO");
+
+    expect(keyColumnRow("ORDERS")?.textContent).toContain("ORDER_NO");
+    expect(keyColumnRow("ORDERS")?.textContent).toContain(i18n.global.t("dataCompare.keyColumnsStatusManual"));
+    // The other table must keep its own primary key.
+    expect(keyColumnRow("USERS")?.textContent).toContain("USER_ID");
+    expect(keyColumnRow("USERS")?.textContent).toContain(i18n.global.t("dataCompare.keyColumnsStatusAuto"));
+
+    compareButton()?.click();
+    await flushAsyncSetup();
+
+    expect(comparedKeyColumns()).toEqual([
+      { table: "USERS", keys: ["USER_ID"] },
+      { table: "ORDERS", keys: ["ORDER_NO"] },
+    ]);
+  });
+
+  it("restores per-table match columns when a session is reopened", async () => {
+    setupMocks();
+    sessionMocks.session = configuredSession();
+    mountDialog({ sessionId: "compare-session" });
+    await flushAsyncSetup();
+
+    expect(keyColumnRow("USERS")?.textContent).toContain("USER_ID");
+    expect(keyColumnRow("USERS")?.textContent).toContain(i18n.global.t("dataCompare.keyColumnsStatusAuto"));
+    expect(keyColumnRow("ORDERS")?.textContent).toContain("ORDER_NO");
+    expect(keyColumnRow("ORDERS")?.textContent).toContain(i18n.global.t("dataCompare.keyColumnsStatusManual"));
+  });
+
+  it("drops a stale manual selection after switching to another source table", async () => {
+    setupMocks();
+    await openWithPrefilledSource("USERS");
+
+    await togglePickerColumn("EMAIL");
+    expect(keyColumnTrigger()?.textContent?.trim()).toBe("USER_ID, EMAIL");
+    expect(document.body.textContent).toContain(i18n.global.t("dataCompare.keyColumnsStatusManual"));
+
+    tableSelectionRow("USERS")?.click();
+    await flushAsyncSetup();
+    tableSelectionRow("ORDERS")?.click();
+    await flushAsyncSetup();
+
+    expect(keyColumnTrigger()?.textContent?.trim()).toBe("TENANT_ID, ORDER_ID");
+    expect(keyColumnHint()).toBe(i18n.global.t("dataCompare.keyColumnsAutoSummary", { columns: "TENANT_ID, ORDER_ID" }));
+  });
+
+  it("sends the source spelling of a match column whose target column differs in case", async () => {
+    setupMocks();
+    await openWithPrefilledSource("CASES");
+
+    expect(keyColumnTrigger()?.textContent?.trim()).toBe("snid");
+    expect(compareButton()?.disabled).toBe(false);
+
+    compareButton()?.click();
+    await flushAsyncSetup();
+
+    const options = mocks.prepareDataCompareFromTables.mock.calls[0]?.[0] as { keyColumns: string[]; columns: string[]; sourceColumns: string[]; sourceTable: string };
+    expect(options.sourceTable).toBe("CASES");
+    expect(options.keyColumns).toEqual(["SNID"]);
+    expect(options.columns).toEqual(["SNID", "USERNAME"]);
+    expect(options.sourceColumns).toEqual(["snid", "username"]);
+    expect(document.body.textContent).not.toContain(i18n.global.t("dataCompare.missingKeyColumns", { columns: "snid" }));
+  });
+});