fix(canvas): reuse save transaction for image archive lookups

lib/db/queries.ts

@@ -1239,13 +1239,13 @@ export async function getMediaAssetByUserAndUrl(params: {
 export async function getMediaAssetsByUserAndUrls(params: {
   userId: string;
   urls: string[];
-}): Promise<MediaAsset[]> {
+}, executor: Pick<typeof db, "select" | "insert" | "delete"> = db): Promise<MediaAsset[]> {
   const normalizedUrls = Array.from(
     new Set(params.urls.map((url) => normalizeMediaUrl(url)).filter(Boolean)),
   );
   if (!normalizedUrls.length) return [];
   try {
-    return await db
+    return await executor
       .select()
       .from(mediaAsset)
       .where(
@@ -1398,9 +1398,9 @@ export async function createMediaAssetRelation(params: {
   sourceNodeId?: string | null;
   confidence?: number | null;
   metadata?: any;
-}): Promise<MediaAssetRelation | null> {
+}, executor: Pick<typeof db, "select" | "insert" | "delete"> = db): Promise<MediaAssetRelation | null> {
   try {
-    const [result] = await db
+    const [result] = await executor
       .insert(mediaAssetRelation)
       .values({
         userId: params.userId,
@@ -1590,8 +1590,8 @@ export async function syncCanvasNodesForProject(params: {
     edges?: any[];
     viewport?: { x: number; y: number; zoom: number };
   } | null | undefined;
-}) {
-  params = { ...params, canvasData: await resolveArchivedImageUrls(params.canvasData) };
+}, executor: Pick<typeof db, "select" | "insert" | "delete"> = db) {
+  params = { ...params, canvasData: await resolveArchivedImageUrls(params.canvasData, executor) };
   const nodes = Array.isArray(params.canvasData?.nodes) ? params.canvasData?.nodes : [];
   const assetUrls = Array.from(
     new Set(
@@ -1604,7 +1604,7 @@ export async function syncCanvasNodesForProject(params: {
     ? await getMediaAssetsByUserAndUrls({
         userId: params.userId,
         urls: assetUrls,
-      })
+      }, executor)
     : [];
   const assetByUrl = new Map<string, MediaAsset>();
   for (const asset of matchedAssets) {
@@ -1628,7 +1628,7 @@ export async function syncCanvasNodesForProject(params: {
     };
   });
 
-  await db.delete(mediaAssetRelation).where(
+  await executor.delete(mediaAssetRelation).where(
     and(
       eq(mediaAssetRelation.projectId, params.projectId),
       eq(mediaAssetRelation.relationType, "located_on_canvas"),
@@ -1666,7 +1666,7 @@ export async function syncCanvasNodesForProject(params: {
   });
 
   if (insertableNodes.length > 0) {
-    await db
+    await executor
       .insert(canvasNode)
       .values(insertableNodes)
       .onConflictDoUpdate({
@@ -1684,14 +1684,14 @@ export async function syncCanvasNodesForProject(params: {
       });
 
     const activeNodeIds = insertableNodes.map((node) => node.id);
-    await db.delete(canvasNode).where(
+    await executor.delete(canvasNode).where(
       and(
         eq(canvasNode.projectId, params.projectId),
         notInArray(canvasNode.id, activeNodeIds),
       ),
     );
   } else {
-    await db.delete(canvasNode).where(eq(canvasNode.projectId, params.projectId));
+    await executor.delete(canvasNode).where(eq(canvasNode.projectId, params.projectId));
   }
 
   for (const node of enrichedNodes) {
@@ -1712,7 +1712,7 @@ export async function syncCanvasNodesForProject(params: {
         position: node.position || null,
         size: node.size || null,
       },
-    });
+    }, executor);
   }
 
   return {

lib/storage/image-archive-db.ts

@@ -290,10 +290,13 @@ export async function recoverStaleImageArchiveTasks() {
     );
 }
 
-export async function getArchivedImageUrlMappings(urls: string[]) {
+export async function getArchivedImageUrlMappings(
+  urls: string[],
+  executor: Pick<typeof db, "select"> = db,
+) {
   if (!urls.length) return {};
   const sources = [...new Set(urls.filter(isGeneratedImageSource))];
-  const tasks = await db
+  const tasks = await executor
     .select({ sourceUrl: imageArchiveTask.sourceUrl, targetUrl: imageArchiveTask.targetUrl })
     .from(imageArchiveTask)
     .where(inArray(imageArchiveTask.sourceUrl, sources.map(imageArchiveSourceKey)));
@@ -309,10 +312,13 @@ export async function getArchivedImageUrlMappings(urls: string[]) {
 }
 
 // Normalize late writes too: an archive can finish before a task/canvas save arrives.
-export async function resolveArchivedImageUrls<T>(value: T): Promise<T> {
+export async function resolveArchivedImageUrls<T>(
+  value: T,
+  executor: Pick<typeof db, "select"> = db,
+): Promise<T> {
   const sources = collectImageArchiveSources(value);
   if (!sources.length) return value;
-  return replaceArchivedImageUrls(value, await getArchivedImageUrlMappings(sources));
+  return replaceArchivedImageUrls(value, await getArchivedImageUrlMappings(sources, executor));
 }
 
 // Download assets use an independent JSONB lease; original archiving/replacement

tests/integration/canvas-save-transactions.test.ts

@@ -0,0 +1,93 @@
+// Run with a disposable migrated database, POSTGRES_MAX_CONNECTIONS=1 and
+// node --conditions=react-server --import tsx --test <this file>.
+import assert from "node:assert/strict";
+import { randomUUID } from "node:crypto";
+import { after, before, test } from "node:test";
+import { and, eq, sql } from "drizzle-orm";
+import { db } from "@/lib/db";
+import { syncCanvasNodesForProject } from "@/lib/db/queries";
+import { canvasNode, canvasProject, imageArchiveTask, mediaAsset, mediaAssetRelation, user } from "@/lib/db/schema";
+import { resolveArchivedImageUrls } from "@/lib/storage/image-archive-db";
+
+const database = new URL(process.env.POSTGRES_URL || "http://missing");
+if (!["127.0.0.1", "localhost", "::1"].includes(database.hostname)) throw new Error("Use a disposable local database");
+if (!/_test$/.test(database.pathname)) throw new Error("Use a disposable test database");
+if (process.env.POSTGRES_MAX_CONNECTIONS !== "1") throw new Error("A single connection makes accidental pool re-entry observable");
+
+const userId = randomUUID();
+const projectId = randomUUID();
+const nodeId = randomUUID();
+const assetId = randomUUID();
+const sourceUrl = `https://provider.example/${randomUUID()}.png`;
+const targetUrl = `https://storage.example/${randomUUID()}.png`;
+const canvas = (x: number) => ({ nodes: [{ id: nodeId, type: "image", position: { x, y: 0 }, data: { src: sourceUrl } }], edges: [] });
+
+before(async () => {
+  await db.insert(user).values({ id: userId, email: `${userId}@transaction-test.invalid` });
+  await db.insert(canvasProject).values({ id: projectId, userId, title: "Transaction regression", canvasData: canvas(0) });
+  await db.insert(mediaAsset).values({ id: assetId, userId, projectId, mediaType: "image", sourceType: "generated", originalUrl: targetUrl, normalizedUrl: targetUrl });
+  await db.insert(imageArchiveTask).values({ sourceUrl, targetUrl, status: "stored", archiveAfterAt: new Date(), replaceAfterAt: new Date() });
+});
+
+after(async () => {
+  await db.delete(imageArchiveTask).where(eq(imageArchiveTask.sourceUrl, sourceUrl));
+  await db.delete(user).where(eq(user.id, userId));
+  await db.$client.end({ timeout: 2 });
+});
+
+test("archive URL lookup stays on the existing connection", { timeout: 8000 }, async () => {
+  const result = await db.transaction(async (tx) => {
+    await tx.execute(sql`set local idle_in_transaction_session_timeout = '2s'`);
+    await tx.select().from(canvasProject).where(eq(canvasProject.id, projectId)).for("update");
+    return resolveArchivedImageUrls(canvas(1), tx);
+  });
+  assert.equal(result.nodes[0].data.src, targetUrl);
+});
+
+async function save(x: number) {
+  return db.transaction(async (tx) => {
+    await tx.execute(sql`set local idle_in_transaction_session_timeout = '2s'`);
+    await tx.select().from(canvasProject).where(and(eq(canvasProject.id, projectId), eq(canvasProject.userId, userId))).for("update");
+    const next = await syncCanvasNodesForProject({ projectId, userId, canvasData: canvas(x) }, tx);
+    await tx.update(canvasProject).set({ canvasData: next }).where(eq(canvasProject.id, projectId));
+    return next;
+  });
+}
+
+test("canvas saves reuse the transaction for URL lookup, asset lookup and node relations", { timeout: 8000 }, async () => {
+  const result = await save(2);
+  assert.equal(result.nodes[0].data.src, targetUrl);
+  assert.equal(result.nodes[0].data.assetId, assetId);
+  const nodes = await db.select().from(canvasNode).where(eq(canvasNode.projectId, projectId));
+  assert.equal(nodes.length, 1);
+  assert.equal((nodes[0].data as { src: string }).src, targetUrl);
+  const relations = await db.select().from(mediaAssetRelation).where(eq(mediaAssetRelation.projectId, projectId));
+  assert.equal(relations.length, 1);
+  assert.equal(relations[0].fromAssetId, assetId);
+});
+
+test("ten overlapping saves complete with a pool of one connection", { timeout: 8000 }, async () => {
+  const results = await Promise.all(Array.from({ length: 10 }, (_, x) => save(x + 10)));
+  assert.equal(results.length, 10);
+  for (const result of results) assert.equal(result.nodes[0].data.src, targetUrl);
+  const relations = await db.select().from(mediaAssetRelation).where(eq(mediaAssetRelation.projectId, projectId));
+  assert.equal(relations.length, 1);
+});
+
+test("a failed save rolls back node changes and asset relations together", { timeout: 8000 }, async () => {
+  await save(30);
+  await assert.rejects(db.transaction(async (tx) => {
+    await tx.execute(sql`set local idle_in_transaction_session_timeout = '2s'`);
+    await tx.select().from(canvasProject).where(eq(canvasProject.id, projectId)).for("update");
+    await syncCanvasNodesForProject({ projectId, userId, canvasData: canvas(99) }, tx);
+    throw new Error("intentional rollback");
+  }), /intentional rollback/);
+  const [node] = await db.select().from(canvasNode).where(eq(canvasNode.id, nodeId));
+  assert.equal((node.position as { x: number }).x, 30);
+  const relations = await db.select().from(mediaAssetRelation).where(eq(mediaAssetRelation.projectId, projectId));
+  assert.equal(relations.length, 1);
+});
+
+test("standalone archive normalization remains compatible", async () => {
+  assert.equal(await resolveArchivedImageUrls(sourceUrl), targetUrl);
+});