feat(datasets): create mutation for dataset creation (#5583)

Author: IcaroG

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

@@ -235,6 +235,26 @@ type Mutation {
   osoApp_saveNotebookPreview(
     input: SaveNotebookPreviewInput!
   ): SaveNotebookPreviewPayload!
+
+  """
+  Create a new dataset
+  """
+  osoApp_createDataset(input: CreateDatasetInput!): CreateDatasetPayload!
+}
+
+input CreateDatasetInput {
+  orgName: String!
+  name: String!
+  displayName: String!
+  description: String
+  datasetType: DatasetType!
+  isPublic: Boolean
+}
+
+type CreateDatasetPayload {
+  dataset: Dataset
+  message: String!
+  success: Boolean!
 }
 
 input CreateInvitationInput {

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

@@ -22,6 +22,9 @@ import { logger } from "@/lib/logger";
 import { assert } from "@opensource-observer/utils";
 import { createAdminClient } from "@/lib/supabase/admin";
 import { DatasetsRow } from "@/lib/types/schema-types";
+import z from "zod";
+import { validateInput } from "@/app/api/v1/osograph/utils/validation";
+import { v4 as uuidv4 } from "uuid";
 
 const TRINO_SCHEMA_TIMEOUT = 10000; // 10 seconds
 
@@ -215,6 +218,62 @@ export const datasetResolver: GraphQLResolverModule<GraphQLContext> = {
       return filteredResults;
     },
   },
+  Mutation: {
+    osoApp_createDataset: async (
+      _: unknown,
+      { input }: { input: z.infer<typeof CreateDatasetSchema> },
+      context: GraphQLContext,
+    ) => {
+      const authenticatedUser = requireAuthentication(context.user);
+      const validated = validateInput(CreateDatasetSchema, input);
+      const organization = await getOrganizationByName(validated.orgName);
+      await requireOrgMembership(authenticatedUser.userId, organization.id);
+
+      const supabase = createAdminClient();
+      const datasetId = uuidv4();
+      let catalog: string;
+      let schema: string;
+
+      switch (validated.datasetType) {
+        case "USER_MODEL":
+          catalog = "user_iceberg";
+          schema = `ds_${datasetId.replace(/-/g, "")}`;
+          break;
+        default:
+          throw new Error(
+            `Dataset type "${validated.datasetType}" is not supported yet.`,
+          );
+      }
+
+      const { data: dataset, error } = await supabase
+        .from("datasets")
+        .insert({
+          id: datasetId,
+          org_id: organization.id,
+          name: validated.name,
+          display_name: validated.displayName,
+          description: validated.description,
+          catalog,
+          schema,
+          created_by: authenticatedUser.userId,
+          is_public: validated.isPublic ?? false,
+          dataset_type: validated.datasetType,
+        })
+        .select()
+        .single();
+
+      if (error) {
+        logger.error("Failed to create dataset:", error);
+        throw ServerErrors.database("Failed to create dataset");
+      }
+
+      return {
+        dataset,
+        message: "Dataset created successfully",
+        success: true,
+      };
+    },
+  },
   Dataset: {
     id: (parent: DatasetsRow) => parent.id,
     orgId: (parent: DatasetsRow) => parent.org_id,
@@ -232,6 +291,21 @@ export const datasetResolver: GraphQLResolverModule<GraphQLContext> = {
   },
 };
 
+export const CreateDatasetSchema = z.object({
+  orgName: z.string().min(1, "Organization name is required"),
+  name: z
+    .string()
+    .min(1, "Dataset name is required")
+    .regex(
+      /^[a-zA-Z][a-zA-Z0-9_]+$/,
+      "Dataset name can only contain letters, numbers, and underscores",
+    ),
+  displayName: z.string().min(1, "Display name is required"),
+  description: z.string().optional(),
+  isPublic: z.boolean().optional(),
+  datasetType: z.enum(["USER_MODEL", "DATA_CONNECTOR", "DATA_INGESTION"]),
+});
+
 async function getOrganizationDatasets(orgId: string): Promise<DatasetsRow[]> {
   const supabase = createAdminClient();
 

apps/frontend/components/dataprovider/oso-global-context.tsx

@@ -97,6 +97,7 @@ const OsoGlobalActionNames: ExtractMethodNames<OsoAppClient>[] = _.sortBy([
   "deductOrganizationCredits",
   "updateOrganizationTier",
   "setOrganizationCredits",
+  "createDataset",
 ]);
 const OsoGlobalActions: Partial<ExtractMethods<OsoAppClient>> = _.fromPairs(
   OsoGlobalActionNames.map((name) => [

apps/frontend/lib/clients/oso-app/oso-app.ts

@@ -2442,6 +2442,83 @@ class OsoAppClient {
       "X-Supabase-Auth": `${sessionData.session.access_token}:${sessionData.session.refresh_token}`,
     };
   }
+
+  async createDataset(
+    args: Partial<{
+      orgName: string;
+      name: string;
+      displayName: string;
+      description: string;
+      datasetType: "USER_MODEL" | "DATA_CONNECTOR" | "DATA_INGESTION";
+      isPublic: boolean;
+    }>,
+  ) {
+    const { orgName, name, displayName, description, datasetType, isPublic } = {
+      orgName: ensure(args.orgName, "Missing orgName argument"),
+      name: ensure(args.name, "Missing name argument"),
+      displayName: ensure(args.displayName, "Missing displayName argument"),
+      description: args.description,
+      datasetType: ensure(args.datasetType, "Missing datasetType argument"),
+      isPublic: args.isPublic,
+    };
+
+    const CREATE_DATASET_MUTATION = gql(`
+      mutation CreateDataset($input: CreateDatasetInput!) {
+        osoApp_createDataset(input: $input) {
+          success
+          message
+          dataset {
+            id
+            name
+            displayName
+            description
+            catalog
+            schema
+            datasetType
+            isPublic
+          }
+        }
+      }
+    `);
+
+    const response = await fetch("/api/v1/osograph", {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({
+        query: print(CREATE_DATASET_MUTATION),
+        variables: {
+          input: {
+            orgName,
+            name,
+            displayName,
+            description,
+            datasetType,
+            isPublic,
+          },
+        },
+      }),
+    });
+
+    const result = await response.json();
+
+    if (result.errors) {
+      logger.error("Failed to create dataset:", result.errors[0].message);
+      throw new Error(`Failed to create dataset: ${result.errors[0].message}`);
+    }
+
+    const payload = result.data?.osoApp_createDataset;
+    if (!payload) {
+      throw new Error("No response data from create dataset mutation");
+    }
+
+    if (payload.success) {
+      logger.log(`Successfully created dataset "${displayName}"`);
+    }
+
+    return payload.dataset;
+  }
 }
 
 export { OsoAppClient };

apps/frontend/lib/graphql/generated/gql.ts

@@ -14,6 +14,7 @@ import { TypedDocumentNode as DocumentNode } from "@graphql-typed-document-node/
  * Learn more about it here: https://the-guild.dev/graphql/codegen/plugins/presets/preset-client#reducing-bundle-size
  */
 type Documents = {
+  "\n      mutation CreateDataset($input: CreateDatasetInput!) {\n        osoApp_createDataset(input: $input) {\n          success\n          message\n          dataset {\n            id\n            name\n            displayName\n            description\n            catalog\n            schema\n            datasetType\n            isPublic\n          }\n        }\n      }\n    ": typeof types.CreateDatasetDocument;
   "\n      mutation SavePreview($input: SaveNotebookPreviewInput!) {\n        osoApp_saveNotebookPreview(input: $input) {\n          success\n          message\n        }\n      }\n    ": typeof types.SavePreviewDocument;
   "\nquery AssetGraph {\n  assetNodes {\n    assetKey {\n      path\n    }\n    dependencyKeys {\n      path\n    }\n  }\n}": typeof types.AssetGraphDocument;
   '\nquery AssetMaterializedData($assetKeys: [AssetKeyInput!] = {path: ""}) {\n  assetNodes(assetKeys: $assetKeys) {\n    assetKey {\n      path\n    }\n    partitionStats {\n      numFailed\n      numMaterialized\n      numMaterializing\n      numPartitions\n    }\n    assetPartitionStatuses {\n      ... on TimePartitionStatuses {\n        __typename\n        ranges {\n          endKey\n          startKey\n          status\n        }\n      }\n    }\n    assetMaterializations(limit: 1) {\n      runOrError {\n        ... on Run {\n          endTime\n        }\n      }\n    }\n  }\n}': typeof types.AssetMaterializedDataDocument;
@@ -22,6 +23,8 @@ type Documents = {
   "\n  query TimeseriesMetricsByCollection(\n    $collectionIds: [String!],\n    $metricIds: [String!],\n    $startDate: Oso_Date!,\n    $endDate: Oso_Date!, \n  ) {\n    oso_timeseriesMetricsByCollectionV0(where: {\n      collectionId: {_in: $collectionIds},\n      metricId: {_in: $metricIds},\n      sampleDate: { _gte: $startDate, _lte: $endDate }\n    }) {\n      amount\n      metricId\n      collectionId\n      sampleDate\n      unit\n    }\n    oso_collectionsV1(where: { collectionId: { _in: $collectionIds }}) {\n      collectionId\n      collectionSource\n      collectionNamespace\n      collectionName\n      displayName\n      description\n    }\n    oso_metricsV0(where: {metricId: {_in: $metricIds}}) {\n      metricId\n      metricSource\n      metricNamespace\n      metricName\n      displayName\n      description\n    }\n  }\n": typeof types.TimeseriesMetricsByCollectionDocument;
 };
 const documents: Documents = {
+  "\n      mutation CreateDataset($input: CreateDatasetInput!) {\n        osoApp_createDataset(input: $input) {\n          success\n          message\n          dataset {\n            id\n            name\n            displayName\n            description\n            catalog\n            schema\n            datasetType\n            isPublic\n          }\n        }\n      }\n    ":
+    types.CreateDatasetDocument,
   "\n      mutation SavePreview($input: SaveNotebookPreviewInput!) {\n        osoApp_saveNotebookPreview(input: $input) {\n          success\n          message\n        }\n      }\n    ":
     types.SavePreviewDocument,
   "\nquery AssetGraph {\n  assetNodes {\n    assetKey {\n      path\n    }\n    dependencyKeys {\n      path\n    }\n  }\n}":
@@ -50,6 +53,12 @@ const documents: Documents = {
  */
 export function gql(source: string): unknown;
 
+/**
+ * The gql function is used to parse GraphQL queries into a document that can be used by GraphQL clients.
+ */
+export function gql(
+  source: "\n      mutation CreateDataset($input: CreateDatasetInput!) {\n        osoApp_createDataset(input: $input) {\n          success\n          message\n          dataset {\n            id\n            name\n            displayName\n            description\n            catalog\n            schema\n            datasetType\n            isPublic\n          }\n        }\n      }\n    ",
+): (typeof documents)["\n      mutation CreateDataset($input: CreateDatasetInput!) {\n        osoApp_createDataset(input: $input) {\n          success\n          message\n          dataset {\n            id\n            name\n            displayName\n            description\n            catalog\n            schema\n            datasetType\n            isPublic\n          }\n        }\n      }\n    "];
 /**
  * The gql function is used to parse GraphQL queries into a document that can be used by GraphQL clients.
  */