Split creation of different trainspaces into different API endpoints

DSGT-DLP · Oct 20, 2023 · 18df85b · 18df85b
1 parent 57f1adc
commit 18df85b
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 34 deletions.
diff --git a/...tions/src/trainspace/create_trainspace.ts → ...src/trainspace/create_image_trainspace.ts b/...tions/src/trainspace/create_trainspace.ts → ...src/trainspace/create_image_trainspace.ts
@@ -6,38 +6,7 @@ import { DynamoDBClient, PutItemCommand } from '@aws-sdk/client-dynamodb';
 import { DynamoDBDocumentClient } from '@aws-sdk/lib-dynamodb';
 
 function validateRequestBody(eventBody: any, trainspace_id: string, uid: string) : TrainspaceData | null {
-    switch(eventBody['data_source'])
-    {
-        case "TABULAR":
-            return new TrainspaceData(trainspace_id, uid, "TABULAR", eventBody['default'], eventBody['name'], 
-                { 
-                    criterion: eventBody['criterion'],
-                    optimizer_name: eventBody['optimizer_name'],
-                    shuffle: eventBody['shuffle'],
-                    epochs: eventBody['epochs'],
-                    batch_size: eventBody['batch_size'],
-                    user_arch: eventBody['user_arch']
-                }, "");
-        case "PRETRAINED":
-
-            return null;
-        case "IMAGE":
-            return new TrainspaceData(trainspace_id, uid, "IMAGE", eventBody['dataset_data'], eventBody['name'], eventBody['parameters_data'], eventBody['review_data']['notification_email']);
-        case "AUDIO":
-
-            return null;
-        case "TEXTUAL":
-
-            return null;
-        case "CLASSICAL_ML":
-
-            return null;
-        case "OBJECT_DETECTION":
-
-            return null;
-        default:
-            return null;
-    }
+    return new TrainspaceData(trainspace_id, uid, "IMAGE", eventBody['dataset_data'], eventBody['name'], eventBody['parameters_data'], eventBody['review_data']['notification_email']);
 }
 
 export const handler: APIGatewayProxyHandlerV2 = async (event) => {

diff --git a/serverless/packages/functions/src/trainspace/create_tabular_trainspace.ts b/serverless/packages/functions/src/trainspace/create_tabular_trainspace.ts
@@ -0,0 +1,59 @@
+import { APIGatewayProxyHandlerV2 } from "aws-lambda";
+import parseJwt from "@dlp-sst-app/core/parseJwt";
+import TrainspaceData from './trainspace-data';
+import { v4 as uuidv4 } from 'uuid';
+import { DynamoDBClient, PutItemCommand } from '@aws-sdk/client-dynamodb';
+import { DynamoDBDocumentClient } from '@aws-sdk/lib-dynamodb';
+
+function validateRequestBody(eventBody: any, trainspace_id: string, uid: string) : TrainspaceData | null {
+    return new TrainspaceData(trainspace_id, uid, "TABULAR", eventBody['default'], eventBody['name'], 
+        { 
+            criterion: eventBody['criterion'],
+            optimizer_name: eventBody['optimizer_name'],
+            shuffle: eventBody['shuffle'],
+            epochs: eventBody['epochs'],
+            batch_size: eventBody['batch_size'],
+            user_arch: eventBody['user_arch']
+        }, "");
+}
+
+export const handler: APIGatewayProxyHandlerV2 = async (event) => {
+    if (event) {
+        const uid: string = parseJwt(event.headers.authorization ?? "")["user_id"];
+        const eventBody = JSON.parse(event.body? event.body : "");
+
+        const trainspaceId = uuidv4();
+        const trainspaceData = validateRequestBody(eventBody, trainspaceId, uid);
+
+        if (trainspaceData == null)
+        {
+            return {
+                statusCode: 400,
+                body: JSON.stringify({ message: "Invalid request body" })
+            }
+        }
+
+
+        const client = new DynamoDBClient({});
+        const docClient = DynamoDBDocumentClient.from(client);
+
+        const command = new PutItemCommand(trainspaceData.convertToDynamoItemInput("trainspace"));
+
+        const response = await docClient.send(command);
+        if (response.$metadata.httpStatusCode != 200) {
+            return {
+                statusCode: 500,
+                body: JSON.stringify({ message: "Internal server error."})
+            };
+        }
+
+        return {
+            statusCode: 200,
+            body: JSON.stringify({ trainspaceId: trainspaceId, message: "Successfully created a new trainspace."})
+        };
+      }
+    return {
+        statusCode: 404,
+        body: JSON.stringify({ message: "Not Found" }),
+    };
+};
diff --git a/serverless/stacks/AppStack.ts b/serverless/stacks/AppStack.ts
@@ -59,9 +59,15 @@ export function AppStack({ stack }: StackContext) {
         "packages/functions/src/datasets/user/columns.handler",
       "DELETE /dataset/user/{type}/{filename}" :
         "packages/functions/src/datasets/user/delete_userdata.handler",
-      "POST /trainspace": {
+      "POST /trainspace/tabular": {
         function: {
-          handler: "packages/functions/src/trainspace/create_trainspace.handler",
+          handler: "packages/functions/src/trainspace/create_tabular_trainspace.handler",
+          permissions: ["dynamodb:PutItem"]
+        }
+      },
+      "POST /trainspace/image": {
+        function: {
+          handler: "packages/functions/src/trainspace/create_image_trainspace.handler",
           permissions: ["dynamodb:PutItem"]
         }
       },