fix(query): enable editing multi-table JOIN result sets

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

@@ -627,6 +627,8 @@ describe("useDataGridEditor saveChanges reload", () => {
 
   function createSaveTestEditor(
     options: {
+      joinedWriteTargets?: import("@/types/database").QueryTab["queryWriteTargets"];
+      queryResult?: { columns: string[]; rows: CellValue[][] };
       currentPage?: Ref<number>;
       prepareFullReload?: () => void;
       customSaveHandler?: { save: ReturnType<typeof vi.fn> };
@@ -644,6 +646,7 @@ describe("useDataGridEditor saveChanges reload", () => {
         [2, "pending"],
       ],
     });
+    if (options.queryResult) result.value = options.queryResult;
     const editor = useDataGridEditor({
       result: computed(() => result.value),
       editable: computed(() => true),
@@ -659,6 +662,7 @@ describe("useDataGridEditor saveChanges reload", () => {
         primaryKeys: ["id"],
       })),
       sourceColumns: computed(() => undefined),
+      joinedWriteTargets: computed(() => options.joinedWriteTargets),
       onExecuteSql: computed(() => undefined),
       customSaveHandler: computed(() => options.customSaveHandler),
       manualTransactionSessionId: computed(() => options.manualTransactionSessionId),
@@ -694,6 +698,55 @@ describe("useDataGridEditor saveChanges reload", () => {
     expect(emit).toHaveBeenCalledWith("reload", undefined, "", undefined, undefined, 100, 0);
   });
 
+  it("previews and saves edits to both joined tables in a single transaction", async () => {
+    const refreshSavedRows = vi.fn();
+    const targets = [
+      { tableMeta: { tableName: "users", primaryKeys: ["id"], columns: [] }, sourceColumns: ["id", "name", undefined, undefined] },
+      { tableMeta: { tableName: "papers", primaryKeys: ["id"], columns: [] }, sourceColumns: [undefined, undefined, "id", "title"] },
+    ];
+    mocks.prepareDataGridSave.mockImplementation(async (options) => ({ statements: ["update " + options.tableMeta.tableName], rollbackStatements: ["undo " + options.tableMeta.tableName] }));
+    mocks.executeInTransaction.mockResolvedValue({ affected_rows: 2 });
+    const { editor, emit } = createSaveTestEditor({ joinedWriteTargets: targets, queryResult: { columns: ["id", "name", "paper_id", "title"], rows: [[1, "old", 20, "old"]] }, refreshSavedRows });
+    editor.dirtyRows.value.set(
+      0,
+      new Map([
+        [1, "new name"],
+        [3, "new title"],
+      ]),
+    );
+    expect(await editor.previewChanges()).toEqual(["update users", "update papers"]);
+    expect(mocks.executeInTransaction).not.toHaveBeenCalled();
+    await editor.saveChanges();
+    expect(mocks.executeInTransaction).toHaveBeenCalledWith("connection-1", "app", ["update users", "update papers"], undefined);
+    expect(mocks.executeBatch).not.toHaveBeenCalled();
+    expect(refreshSavedRows).not.toHaveBeenCalled();
+    expect(emit).toHaveBeenCalledWith("reload", undefined, "", undefined, undefined, 100, 0);
+    expect(mocks.prepareDataGridSave.mock.calls[0]![0].dirtyRows).toEqual([[0, [[1, "new name"]]]]);
+    expect(mocks.prepareDataGridSave.mock.calls[1]![0].dirtyRows).toEqual([[0, [[3, "new title"]]]]);
+  });
+
+  it("keeps both joined edits pending if the transaction fails", async () => {
+    const targets = [
+      { tableMeta: { tableName: "users", primaryKeys: ["id"], columns: [] }, sourceColumns: ["id", "name", undefined, undefined] },
+      { tableMeta: { tableName: "papers", primaryKeys: ["id"], columns: [] }, sourceColumns: [undefined, undefined, "id", "title"] },
+    ];
+    mocks.prepareDataGridSave.mockImplementation(async (options) => ({ statements: ["update " + options.tableMeta.tableName], rollbackStatements: [] }));
+    mocks.executeInTransaction.mockRejectedValue(new Error("second update failed"));
+    const { editor, emit } = createSaveTestEditor({ joinedWriteTargets: targets, queryResult: { columns: ["id", "name", "paper_id", "title"], rows: [[1, "old", 20, "old"]] } });
+    editor.dirtyRows.value.set(
+      0,
+      new Map([
+        [1, "new name"],
+        [3, "new title"],
+      ]),
+    );
+    await editor.saveChanges();
+    expect(editor.dirtyRows.value.get(0)?.size).toBe(2);
+    expect(editor.saveError.value).toContain("second update failed");
+    expect(mocks.executeBatch).not.toHaveBeenCalled();
+    expect(emit.mock.calls.some(([event]) => event === "reload")).toBe(false);
+  });
+
   it("saves query-result edits through the active manual transaction session", async () => {
     const statement = "UPDATE orders_test SET status='shipped' WHERE id=1";
     const refreshSavedRows = vi.fn().mockResolvedValue(true);

apps/desktop/src/composables/useDataGridEditor.ts

@@ -1,3 +1,4 @@
+import { joinedSaveOptions } from "@/lib/dataGrid/joinedRowChanges";
 import { ref, shallowRef, triggerRef, computed, nextTick, watch, getCurrentInstance, onActivated, onBeforeUnmount, onDeactivated, onMounted, toRaw, type ComputedRef, type Ref } from "vue";
 import * as api from "@/lib/backend/api";
 import type { CellValue } from "@/lib/dataGrid/cellValue";
@@ -99,6 +100,7 @@ export interface UseDataGridEditorOptions {
     | undefined
   >;
   sourceColumns?: ComputedRef<Array<string | undefined> | undefined>;
+  joinedWriteTargets?: ComputedRef<import("@/types/database").QueryTab["queryWriteTargets"]>;
   readonlyColumnIndexes?: ComputedRef<ReadonlySet<number> | undefined>;
   canEditExistingRows?: ComputedRef<boolean>;
   onExecuteSql: ComputedRef<((sql: string) => Promise<void>) | undefined>;
@@ -217,6 +219,7 @@ export function useDataGridEditor(options: UseDataGridEditorOptions) {
     database,
     tableMeta,
     sourceColumns = computed(() => undefined),
+    joinedWriteTargets = computed(() => undefined),
     readonlyColumnIndexes = computed(() => undefined),
     canEditExistingRows = computed(() => true),
     onExecuteSql,
@@ -666,7 +669,9 @@ export function useDataGridEditor(options: UseDataGridEditorOptions) {
       }
     }
     const columnName = sourceColumns.value?.[columnIndex] ?? result.value.columns[columnIndex];
-    const info = columnName ? tableMeta.value?.columns.find((column) => column.name.toLowerCase() === columnName.toLowerCase()) : undefined;
+    const sourceTarget = joinedWriteTargets.value?.find((target) => target.sourceColumns[columnIndex] !== undefined);
+    const metadata = sourceTarget?.tableMeta ?? tableMeta.value;
+    const info = columnName ? metadata?.columns.find((column) => column.name.toLowerCase() === columnName.toLowerCase()) : undefined;
     if (isBatching && batchColumnInfoCache) {
       batchColumnInfoCache.set(columnIndex, info);
     }
@@ -1494,12 +1499,33 @@ export function useDataGridEditor(options: UseDataGridEditorOptions) {
     };
   }
 
+  async function prepareSaveStatements(stmtOptions: NonNullable<ReturnType<typeof saveStatementOptions>>) {
+    const targets = joinedWriteTargets.value;
+    if (!targets?.length) return api.prepareDataGridSave(stmtOptions, saveDriverProfile());
+    const groups = joinedSaveOptions(targets, {
+      ...stmtOptions,
+      dirtyRows: new Map(stmtOptions.dirtyRows.map(([index, changes]) => [index, new Map(changes)])),
+      deletedRows: new Set(stmtOptions.deletedRows),
+    });
+    const prepared: Awaited<ReturnType<typeof api.prepareDataGridSave>> = { statements: [], rollbackStatements: [] };
+    for (const group of groups) {
+      const part = await api.prepareDataGridSave(group, saveDriverProfile());
+      if (part.validationError) return part;
+      prepared.statements.push(...part.statements);
+      prepared.rollbackStatements.unshift(...part.rollbackStatements);
+      if (prepared.executionSchema && part.executionSchema && prepared.executionSchema !== part.executionSchema) throw new Error("Joined source tables require incompatible execution schemas.");
+      prepared.executionSchema ??= part.executionSchema;
+    }
+    return prepared;
+  }
+
   function saveDriverProfile() {
     const id = connectionId.value;
     return id ? connectionStore.getConfig(id)?.driver_profile : undefined;
   }
 
   function tableHistoryTarget() {
+    if (joinedWriteTargets.value?.length) return [...new Set(joinedWriteTargets.value.map(({ tableMeta: target }) => [target.database, target.schema, target.tableName].filter(Boolean).join(".")))].join(", ");
     if (!tableMeta.value) return "";
     return [tableMeta.value.schema, tableMeta.value.tableName].filter(Boolean).join(".");
   }
@@ -1799,7 +1825,7 @@ export function useDataGridEditor(options: UseDataGridEditorOptions) {
     let preparedSave: Awaited<ReturnType<typeof api.prepareDataGridSave>> | undefined;
     if (stmtOptions) {
       try {
-        preparedSave = await api.prepareDataGridSave(stmtOptions, saveDriverProfile());
+        preparedSave = await prepareSaveStatements(stmtOptions);
       } catch (e: any) {
         saveError.value = normalizeDataGridSaveError(databaseType.value, e);
         await finishInterruptedSaveChanges(snapshot);
@@ -1892,7 +1918,7 @@ export function useDataGridEditor(options: UseDataGridEditorOptions) {
     applyDirtyRowsToResult(snapshot);
     options.onResultPayloadMutated?.();
     let savedRowsRefreshed = false;
-    if (!manualTransactionSessionId.value && !shouldReloadAfterSave && snapshot.dirtyRows.size > 0 && options.refreshSavedRows) {
+    if (!joinedWriteTargets.value?.length && !manualTransactionSessionId.value && !shouldReloadAfterSave && snapshot.dirtyRows.size > 0 && options.refreshSavedRows) {
       try {
         savedRowsRefreshed = await options.refreshSavedRows({
           dirtyRows: snapshot.dirtyRows,
@@ -2029,7 +2055,7 @@ export function useDataGridEditor(options: UseDataGridEditorOptions) {
       }
       const stmtOptions = saveStatementOptions();
       if (!stmtOptions) return [];
-      const prepared = await api.prepareDataGridSave(stmtOptions, saveDriverProfile());
+      const prepared = await prepareSaveStatements(stmtOptions);
       if (prepared?.validationError) {
         saveError.value = prepared.validationError;
         return [];