feat: save preview notebooks to r2 (#5546)

* chore: regenerate `lock` file

* chore: upgrade aws `deps` and install `presigner`

* feat: implement `notebook` preview graphql `query` and `mutations`

* fix: remove `explicit` casts

* chore: register `saveNotebookPreview` callback

* feat: implement `r2` image save

* chore: run `codegen` scripts

* chore: run `codegen` scripts

* chore: address `PR` feedback

* fix: eslint `config` issue

* fix: `eslint` warning issue

Author: Jabolol

.lintstagedrc

@@ -1,6 +1,6 @@
 {
-    "**/!(.eslintrc)*.{js,jsx,ts,tsx,sol}": [
-        "eslint --fix --max-warnings 0 --no-ignore",
+    "**/*.{js,jsx,ts,tsx,sol}": [
+        "eslint --flag v10_config_lookup_from_file --fix --max-warnings 0",
         "prettier --ignore-path .gitignore --write",
         "prettier --ignore-path .gitignore --log-level warn --check"
     ],
@@ -14,4 +14,4 @@
         "uv run isort",
         "pnpm pyright"
     ]
-}
+}

apps/frontend/app/api/v1/osograph/schema.graphql

@@ -80,6 +80,23 @@ type Invitation @key(fields: "id") {
   acceptedBy: User
 }
 
+"""
+Notebook entity - represents a computational notebook
+"""
+type Notebook @key(fields: "id") {
+  id: ID!
+  notebookName: String!
+  description: String
+  createdAt: DateTime!
+  updatedAt: DateTime!
+  organization: Organization!
+
+  """
+  Get preview image with signed URL
+  """
+  preview: NotebookPreview
+}
+
 """
 Dataset entity - represents a dataset within an organization
 """
@@ -136,6 +153,16 @@ type Query {
   """
   osoApp_invitation(id: ID!): Invitation
 
+  """
+  Get notebooks, optionally filtered by organization name
+  """
+  osoApp_notebooks(orgName: String): [Notebook!]!
+
+  """
+  Get a specific notebook by ID
+  """
+  osoApp_notebook(notebookId: ID!): Notebook
+
   """
   Get list of datasets for a given organization
   """
@@ -194,6 +221,13 @@ type Mutation {
     email: String!
     role: MemberRole!
   ): AddUserByEmailPayload!
+
+  """
+  Save notebook preview PNG image
+  """
+  osoApp_saveNotebookPreview(
+    input: SaveNotebookPreviewInput!
+  ): SaveNotebookPreviewPayload!
 }
 
 input CreateInvitationInput {
@@ -262,3 +296,23 @@ type UpdateProfilePayload {
   message: String!
   success: Boolean!
 }
+
+type NotebookPreview {
+  notebookId: ID!
+  signedUrl: String!
+  expiresAt: DateTime!
+}
+
+input SaveNotebookPreviewInput {
+  notebookId: ID!
+  orgName: String!
+  """
+  PNG image as data URL (data:image/png;base64,...)
+  """
+  previewImage: String!
+}
+
+type SaveNotebookPreviewPayload {
+  success: Boolean!
+  message: String!
+}

apps/frontend/app/api/v1/osograph/schema/resolvers/index.ts

@@ -5,6 +5,7 @@ import { organizationResolvers } from "@/app/api/v1/osograph/schema/resolvers/or
 import { userResolvers } from "@/app/api/v1/osograph/schema/resolvers/user";
 import { memberResolvers } from "@/app/api/v1/osograph/schema/resolvers/member";
 import { invitationResolvers } from "@/app/api/v1/osograph/schema/resolvers/invitation";
+import { notebookResolvers } from "@/app/api/v1/osograph/schema/resolvers/notebook";
 import { datasetResolver } from "@/app/api/v1/osograph/schema/resolvers/dataset";
 
 const dateTimeScalar = new GraphQLScalarType({
@@ -39,16 +40,19 @@ export const resolvers: GraphQLResolverMap<GraphQLContext> = {
     ...userResolvers.Query,
     ...organizationResolvers.Query,
     ...invitationResolvers.Query,
+    ...notebookResolvers.Query,
     ...datasetResolver.Query,
   },
   Mutation: {
     ...userResolvers.Mutation,
     ...memberResolvers.Mutation,
     ...invitationResolvers.Mutation,
+    ...notebookResolvers.Mutation,
   },
   User: userResolvers.User,
   Organization: organizationResolvers.Organization,
   OrganizationMember: memberResolvers.OrganizationMember,
   Invitation: invitationResolvers.Invitation,
+  Notebook: notebookResolvers.Notebook,
   Dataset: datasetResolver.Dataset,
 };

apps/frontend/app/api/v1/osograph/schema/resolvers/notebook.ts

@@ -0,0 +1,194 @@
+import { createAdminClient } from "@/lib/supabase/admin";
+import type { GraphQLResolverModule } from "@/app/api/v1/osograph/utils/types";
+import {
+  requireAuthentication,
+  requireOrgMembership,
+  getOrganizationByName,
+  getOrganization,
+  type GraphQLContext,
+} from "@/app/api/v1/osograph/utils/auth";
+import {
+  ServerErrors,
+  NotebookErrors,
+  OrganizationErrors,
+} from "@/app/api/v1/osograph/utils/errors";
+import { validateBase64PngImage } from "@/app/api/v1/osograph/utils/validation";
+import {
+  putBase64Image,
+  getPreviewSignedUrl,
+} from "@/lib/clients/cloudflare-r2";
+import { logger } from "@/lib/logger";
+
+const PREVIEWS_BUCKET = "notebook-previews";
+const SIGNED_URL_EXPIRY = 900;
+
+export const notebookResolvers: GraphQLResolverModule<GraphQLContext> = {
+  Query: {
+    osoApp_notebooks: async (
+      _: unknown,
+      args: { orgName?: string },
+      context: GraphQLContext,
+    ) => {
+      const authenticatedUser = requireAuthentication(context.user);
+      const supabase = createAdminClient();
+
+      let query = supabase
+        .from("notebooks")
+        .select(
+          "id, notebook_name, description, created_at, updated_at, org_id",
+        )
+        .is("deleted_at", null);
+
+      if (args.orgName) {
+        const org = await getOrganizationByName(args.orgName);
+        if (!org) {
+          throw OrganizationErrors.notFound();
+        }
+        await requireOrgMembership(authenticatedUser.userId, org.id);
+        query = query.eq("org_id", org.id);
+      } else {
+        const { data: memberships } = await supabase
+          .from("users_by_organization")
+          .select("org_id")
+          .eq("user_id", authenticatedUser.userId)
+          .is("deleted_at", null);
+
+        const orgIds = memberships?.map((m) => m.org_id) || [];
+        if (orgIds.length === 0) {
+          return [];
+        }
+        query = query.in("org_id", orgIds);
+      }
+
+      const { data: notebooks, error } = await query;
+
+      if (error) {
+        logger.error(`Failed to fetch notebooks: ${error}`);
+        throw ServerErrors.database("Failed to fetch notebooks");
+      }
+
+      return notebooks || [];
+    },
+
+    osoApp_notebook: async (
+      _: unknown,
+      args: { notebookId: string },
+      context: GraphQLContext,
+    ) => {
+      const authenticatedUser = requireAuthentication(context.user);
+      const supabase = createAdminClient();
+
+      const { data: notebook, error } = await supabase
+        .from("notebooks")
+        .select(
+          "id, notebook_name, description, created_at, updated_at, org_id",
+        )
+        .eq("id", args.notebookId)
+        .is("deleted_at", null)
+        .single();
+
+      if (error || !notebook) {
+        throw NotebookErrors.notFound();
+      }
+
+      await requireOrgMembership(authenticatedUser.userId, notebook.org_id);
+
+      return notebook;
+    },
+  },
+
+  Mutation: {
+    osoApp_saveNotebookPreview: async (
+      _: unknown,
+      args: {
+        input: {
+          notebookId: string;
+          orgName: string;
+          previewImage: string;
+        };
+      },
+      context: GraphQLContext,
+    ) => {
+      const authenticatedUser = requireAuthentication(context.user);
+      const { notebookId, orgName, previewImage } = args.input;
+
+      const org = await getOrganizationByName(orgName);
+      if (!org) {
+        throw OrganizationErrors.notFound();
+      }
+
+      await requireOrgMembership(authenticatedUser.userId, org.id);
+
+      try {
+        logger.log(`Validating base64 PNG image for notebook ${notebookId}`);
+        validateBase64PngImage(previewImage);
+
+        const objectKey = `${notebookId}.png`;
+        logger.log(
+          `Uploading notebook preview for ${notebookId} to bucket "${PREVIEWS_BUCKET}" with key "${objectKey}". Image size: ${previewImage.length} bytes`,
+        );
+
+        await putBase64Image(PREVIEWS_BUCKET, objectKey, previewImage);
+
+        const response = {
+          success: true,
+          message: "Notebook preview saved successfully",
+        };
+
+        logger.log(
+          `Successfully saved notebook preview for ${notebookId} to bucket "${PREVIEWS_BUCKET}"`,
+        );
+
+        return response;
+      } catch (error) {
+        logger.error(
+          `Failed to save notebook preview for notebook ${notebookId}: ${error}`,
+        );
+        throw ServerErrors.storage("Failed to save notebook preview");
+      }
+    },
+  },
+
+  Notebook: {
+    __resolveReference: async (reference: { id: string }) => {
+      const supabase = createAdminClient();
+      const { data } = await supabase
+        .from("notebooks")
+        .select(
+          "id, notebook_name, description, created_at, updated_at, org_id",
+        )
+        .eq("id", reference.id)
+        .is("deleted_at", null)
+        .single();
+      return data;
+    },
+
+    notebookName: (parent: { notebook_name: string }) => parent.notebook_name,
+    createdAt: (parent: { created_at: string }) => parent.created_at,
+    updatedAt: (parent: { updated_at: string }) => parent.updated_at,
+    organization: (parent: { org_id: string }) =>
+      getOrganization(parent.org_id),
+
+    preview: async (parent: { id: string }) => {
+      try {
+        const objectKey = `${parent.id}.png`;
+        const signedUrl = await getPreviewSignedUrl(
+          PREVIEWS_BUCKET,
+          objectKey,
+          SIGNED_URL_EXPIRY,
+        );
+
+        return {
+          notebookId: parent.id,
+          signedUrl,
+          expiresAt: new Date(Date.now() + SIGNED_URL_EXPIRY * 1000),
+        };
+      } catch (error) {
+        logger.error(
+          `Failed to generate preview URL for notebook ${parent.id}: ${error}`,
+        );
+        return null;
+      }
+    },
+  },
+};

apps/frontend/app/api/v1/osograph/utils/errors.ts

@@ -126,6 +126,9 @@ export const ServerErrors = {
 
   externalService: (rawErrorMessage?: string) =>
     createError(ErrorCode.EXTERNAL_SERVICE_ERROR, rawErrorMessage),
+
+  storage: (rawErrorMessage?: string) =>
+    createError(ErrorCode.EXTERNAL_SERVICE_ERROR, rawErrorMessage),
 } as const;
 
 export const OrganizationErrors = {
@@ -191,4 +194,8 @@ export const UserErrors = {
     createError(ErrorCode.BAD_USER_INPUT, "No fields provided to update"),
 } as const;
 
+export const NotebookErrors = {
+  notFound: () => createError(ErrorCode.NOT_FOUND, "Notebook not found"),
+} as const;
+
 export const CatalogErrors = {};

apps/frontend/app/api/v1/osograph/utils/validation.ts

@@ -1,6 +1,40 @@
 import { z } from "zod";
 import { ValidationErrors } from "@/app/api/v1/osograph/utils/errors";
 
+const MAX_PREVIEW_SIZE_MB = 1 * 1024 * 1024;
+const PNG_HEADER = Buffer.from([
+  0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a,
+]); // PNG magic bytes
+
+export function validateBase64PngImage(base64Data: string): void {
+  if (!base64Data.startsWith("data:image/png;base64,")) {
+    throw new Error("Invalid image format. Expected PNG data URL.");
+  }
+
+  const base64Content = base64Data.replace(/^data:[^;]+;base64,/, "");
+
+  if (base64Content.length > (MAX_PREVIEW_SIZE_MB * 4) / 3) {
+    throw new Error("Image is too large.");
+  }
+
+  try {
+    const buffer = Buffer.from(base64Content, "base64");
+
+    if (!buffer.subarray(0, 8).equals(PNG_HEADER)) {
+      throw new Error("Invalid PNG header.");
+    }
+
+    if (buffer.length > MAX_PREVIEW_SIZE_MB) {
+      throw new Error("Image is too large.");
+    }
+  } catch (error) {
+    if (error instanceof Error && error.message.includes("Invalid PNG")) {
+      throw error;
+    }
+    throw new Error("Invalid base64 encoding or corrupted image data.");
+  }
+}
+
 export const CreateInvitationSchema = z.object({
   email: z.string().email("Invalid email address"),
   orgName: z.string().min(1, "Organization name is required"),