Add storage ID handling and optimize canvas storage URL retrieval

- Introduced `hasStorageId` function to check for valid storage IDs in canvas nodes.
- Updated `batchGetUrlsForCanvas` query to utilize helper functions for improved readability and maintainability.
- Implemented `assertCanvasOwner`, `listNodesForCanvas`, `collectStorageIds`, and `resolveStorageUrls` to streamline the process of fetching storage URLs associated with a canvas.
- Enhanced query logic to skip unnecessary database calls when no valid storage IDs are present.

components/canvas/canvas.tsx

@@ -152,6 +152,11 @@ function isLikelyTransientSyncError(error: unknown): boolean {
   );
 }
 
+function hasStorageId(node: Doc<"nodes">): boolean {
+  const data = node.data as Record<string, unknown> | undefined;
+  return typeof data?.storageId === "string" && data.storageId.length > 0;
+}
+
 function CanvasInner({ canvasId }: CanvasInnerProps) {
   const t = useTranslations('toasts');
   const { screenToFlowPosition } = useReactFlow();
@@ -209,9 +214,14 @@ function CanvasInner({ canvasId }: CanvasInnerProps) {
     api.edges.list,
     shouldSkipCanvasQueries ? "skip" : { canvasId },
   );
+  const shouldSkipStorageUrlQuery = useMemo(() => {
+    if (shouldSkipCanvasQueries) return true;
+    if (convexNodes === undefined) return true;
+    return !convexNodes.some(hasStorageId);
+  }, [convexNodes, shouldSkipCanvasQueries]);
   const storageUrlsById = useQuery(
     api.storage.batchGetUrlsForCanvas,
-    shouldSkipCanvasQueries ? "skip" : { canvasId },
+    shouldSkipStorageUrlQuery ? "skip" : { canvasId },
   );
   const canvas = useQuery(
     api.canvases.get,

convex/storage.ts

@@ -1,8 +1,57 @@
-import { mutation, query } from "./_generated/server";
+import { mutation, query, type QueryCtx } from "./_generated/server";
 import { v } from "convex/values";
 import { requireAuth } from "./helpers";
 import type { Id } from "./_generated/dataModel";
 
+type StorageUrlMap = Record<string, string | undefined>;
+
+async function assertCanvasOwner(
+  ctx: QueryCtx,
+  canvasId: Id<"canvases">,
+  userId: string,
+): Promise<void> {
+  const canvas = await ctx.db.get(canvasId);
+  if (!canvas || canvas.ownerId !== userId) {
+    throw new Error("Canvas not found");
+  }
+}
+
+async function listNodesForCanvas(ctx: QueryCtx, canvasId: Id<"canvases">) {
+  return await ctx.db
+    .query("nodes")
+    .withIndex("by_canvas", (q) => q.eq("canvasId", canvasId))
+    .collect();
+}
+
+function collectStorageIds(
+  nodes: Array<{ data: unknown }>,
+): Array<Id<"_storage">> {
+  const ids = new Set<Id<"_storage">>();
+
+  for (const node of nodes) {
+    const data = node.data as Record<string, unknown> | undefined;
+    const storageId = data?.storageId;
+    if (typeof storageId === "string" && storageId.length > 0) {
+      ids.add(storageId as Id<"_storage">);
+    }
+  }
+
+  return [...ids];
+}
+
+async function resolveStorageUrls(
+  ctx: QueryCtx,
+  storageIds: Array<Id<"_storage">>,
+): Promise<StorageUrlMap> {
+  const entries = await Promise.all(
+    storageIds.map(
+      async (id) => [id, (await ctx.storage.getUrl(id)) ?? undefined] as const,
+    ),
+  );
+
+  return Object.fromEntries(entries) as StorageUrlMap;
+}
+
 export const generateUploadUrl = mutation({
   args: {},
   handler: async (ctx) => {
@@ -19,32 +68,11 @@ export const batchGetUrlsForCanvas = query({
   args: { canvasId: v.id("canvases") },
   handler: async (ctx, { canvasId }) => {
     const user = await requireAuth(ctx);
-    const canvas = await ctx.db.get(canvasId);
+    await assertCanvasOwner(ctx, canvasId, user.userId);
-    if (!canvas || canvas.ownerId !== user.userId) {
-      throw new Error("Canvas not found");
-    }
 
-    const nodes = await ctx.db
+    const nodes = await listNodesForCanvas(ctx, canvasId);
-      .query("nodes")
+    const storageIds = collectStorageIds(nodes);
-      .withIndex("by_canvas", (q) => q.eq("canvasId", canvasId))
-      .collect();
 
-    const ids = new Set<Id<"_storage">>();
+    return await resolveStorageUrls(ctx, storageIds);
-    for (const node of nodes) {
-      const data = node.data as Record<string, unknown> | undefined;
-      const sid = data?.storageId;
-      if (typeof sid === "string" && sid.length > 0) {
-        ids.add(sid as Id<"_storage">);
-      }
-    }
-
-    const entries = await Promise.all(
-      [...ids].map(
-        async (id) =>
-          [id, (await ctx.storage.getUrl(id)) ?? undefined] as const,
-      ),
-    );
-
-    return Object.fromEntries(entries) as Record<string, string | undefined>;
   },
 });