From 18df85b70e1d0c8c786f8d5ac1ed37437939d3fa Mon Sep 17 00:00:00 2001
From: "le_hooligan\\Yvng Alan" <alantao912@protonmail.ch>
Date: Fri, 20 Oct 2023 12:45:10 -0400
Subject: [PATCH] Split creation of different trainspaces into different API
 endpoints

---
 ...ainspace.ts => create_image_trainspace.ts} | 33 +----------
 .../trainspace/create_tabular_trainspace.ts   | 59 +++++++++++++++++++
 serverless/stacks/AppStack.ts                 | 10 +++-
 3 files changed, 68 insertions(+), 34 deletions(-)
 rename serverless/packages/functions/src/trainspace/{create_trainspace.ts => create_image_trainspace.ts} (60%)
 create mode 100644 serverless/packages/functions/src/trainspace/create_tabular_trainspace.ts

diff --git a/serverless/packages/functions/src/trainspace/create_trainspace.ts b/serverless/packages/functions/src/trainspace/create_image_trainspace.ts
similarity index 60%
rename from serverless/packages/functions/src/trainspace/create_trainspace.ts
rename to serverless/packages/functions/src/trainspace/create_image_trainspace.ts
index a151384db..b9153a1d4 100644
--- a/serverless/packages/functions/src/trainspace/create_trainspace.ts
+++ b/serverless/packages/functions/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
new file mode 100644
index 000000000..3b739a353
--- /dev/null
+++ 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" }),
+    };
+};
\ No newline at end of file
diff --git a/serverless/stacks/AppStack.ts b/serverless/stacks/AppStack.ts
index ba8516f3b..9fe043387 100644
--- 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"]
         }
       },