diff --git a/lib/db/queries.ts b/lib/db/queries.ts index d4e6c62c0..4012d8d18 100644 --- a/lib/db/queries.ts +++ b/lib/db/queries.ts @@ -1239,13 +1239,13 @@ export async function getMediaAssetByUserAndUrl(params: { export async function getMediaAssetsByUserAndUrls(params: { userId: string; urls: string[]; -}): Promise { +}, executor: Pick = db): Promise { 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 { +}, executor: Pick = db): Promise { 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 = 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(); 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 { diff --git a/lib/storage/image-archive-db.ts b/lib/storage/image-archive-db.ts index 6c34a9c30..74391e1e5 100644 --- a/lib/storage/image-archive-db.ts +++ b/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 = 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(value: T): Promise { +export async function resolveArchivedImageUrls( + value: T, + executor: Pick = db, +): Promise { 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 diff --git a/tests/integration/canvas-save-transactions.test.ts b/tests/integration/canvas-save-transactions.test.ts new file mode 100644 index 000000000..2a6cced29 --- /dev/null +++ b/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 . +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); +});