feature: Databerry API Upload single endpoint

gmpetrov · Jun 9, 2023 · 30f33b7 · 30f33b7 · vercel · Jun 9, 2023
1 parent ffef6be
commit 30f33b7
Show file tree

Hide file tree

Showing 5 changed files with 354 additions and 43 deletions.
diff --git a/package.json b/package.json
@@ -76,6 +76,7 @@
     "mammoth": "^1.5.1",
     "mime-types": "^2.1.35",
     "msgpackr": "^1.9.3",
+    "multer": "1.4.5-lts.1",
     "nanoid": "^4.0.2",
     "next": "13.2.4",
     "next-auth": "^4.20.1",
@@ -114,6 +115,7 @@
     "@swc/core": "^1.3.51",
     "@types/cors": "^2.8.13",
     "@types/mime-types": "^2.1.1",
+    "@types/multer": "^1.4.7",
     "@types/nprogress": "^0.2.0",
     "@types/react-syntax-highlighter": "^15.5.6",
     "@types/uuid": "^9.0.1",

diff --git a/pages/api/external/datastores/file-upload/[id].ts b/pages/api/external/datastores/file-upload/[id].ts
@@ -1,60 +1,78 @@
 import {
   DatasourceStatus,
   DatasourceType,
-  DatastoreVisibility,
   SubscriptionPlan,
   Usage,
 } from '@prisma/client';
 import Cors from 'cors';
+import mime from 'mime-types';
+import multer from 'multer';
 import { NextApiRequest, NextApiResponse } from 'next';
 import { z } from 'zod';
 
-import { AppNextApiRequest } from '@app/types/index';
+import {
+  AcceptedDatasourceMimeTypes,
+  AppNextApiRequest,
+} from '@app/types/index';
+import accountConfig from '@app/utils/account-config';
 import { ApiError, ApiErrorType } from '@app/utils/api-error';
+import { s3 } from '@app/utils/aws';
 import { createApiHandler, respond } from '@app/utils/createa-api-handler';
+import generateFunId from '@app/utils/generate-fun-id';
 import guardDataProcessingUsage from '@app/utils/guard-data-processing-usage';
 import prisma from '@app/utils/prisma-client';
 import runMiddleware from '@app/utils/run-middleware';
 import triggerTaskLoadDatasource from '@app/utils/trigger-task-load-datasource';
-import validate from '@app/utils/validate';
 
 const cors = Cors({
   methods: ['POST', 'HEAD'],
 });
 
 const handler = createApiHandler();
 
-const Schema = z.object({
-  id: z.string().cuid(),
+const FileSchema = z.object({
+  mimetype: z.enum(AcceptedDatasourceMimeTypes),
+  fieldname: z.string(),
+  originalname: z.string(),
+  encoding: z.string(),
+  size: z.number(),
+  buffer: z.any(),
 });
 
-export const fileUpload = async (
-  req: AppNextApiRequest,
-  res: NextApiResponse
-) => {
-  const datastoreId = req.query.id as string;
+export const upload = async (req: AppNextApiRequest, res: NextApiResponse) => {
+  const file = (req as any).file as z.infer<typeof FileSchema>;
+
+  try {
+    await FileSchema.parseAsync(file);
+  } catch (err) {
+    console.log('Error File Upload', err);
+    throw new ApiError(ApiErrorType.INVALID_REQUEST);
+  }
 
-  const data = req.body as z.infer<typeof Schema>;
-  const datasourceId = data?.id;
+  const datastoreId = req.query.id as string;
 
   // get Bearer token from header
   const authHeader = req.headers.authorization;
-  const token = authHeader && authHeader.split(' ')?.[1];
+  const token = authHeader && authHeader?.split(' ')?.[1];
+
+  if (!token) {
+    throw new ApiError(ApiErrorType.UNAUTHORIZED);
+  }
 
   if (!datastoreId) {
     throw new ApiError(ApiErrorType.INVALID_REQUEST);
   }
 
-  const datasource = await prisma.appDatasource.findUnique({
+  const datastore = await prisma.datastore.findUnique({
     where: {
-      id: datasourceId,
+      id: datastoreId,
     },
     include: {
-      datastore: true,
+      apiKeys: true,
       owner: {
         include: {
-          usage: true,
           apiKeys: true,
+          usage: true,
           subscriptions: {
             where: {
               status: 'active',
@@ -65,42 +83,80 @@ export const fileUpload = async (
     },
   });
 
-  if (!datasource) {
+  if (!datastore) {
     throw new ApiError(ApiErrorType.NOT_FOUND);
   }
 
-  if (
-    datasource?.datastoreId !== datastoreId ||
-    (datasource?.datastore?.visibility === DatastoreVisibility.private &&
-      (!token ||
-        !datasource?.owner?.apiKeys.find((each) => each.key === token)))
-  ) {
+  if (!token || !datastore?.owner?.apiKeys.find((each) => each.key === token)) {
     throw new ApiError(ApiErrorType.UNAUTHORIZED);
   }
 
+  const plan =
+    datastore?.owner?.subscriptions?.[0]?.plan || SubscriptionPlan.level_0;
+
+  if (file.size > accountConfig[plan]?.limits?.maxFileSize) {
+    throw new ApiError(ApiErrorType.USAGE_LIMIT);
+  }
+
   guardDataProcessingUsage({
-    usage: datasource?.owner?.usage as Usage,
-    plan:
-      datasource?.owner?.subscriptions?.[0]?.plan || SubscriptionPlan.level_0,
+    usage: datastore?.owner?.usage as Usage,
+    plan,
   });
 
+  const datasource = await prisma.appDatasource.create({
+    data: {
+      type: DatasourceType.file,
+      name: file.originalname || generateFunId(),
+      config: {
+        type: file.mimetype,
+      },
+      status: DatasourceStatus.pending,
+      owner: {
+        connect: {
+          id: datastore?.ownerId!,
+        },
+      },
+      datastore: {
+        connect: {
+          id: datastoreId,
+        },
+      },
+    },
+  });
+
+  // Add to S3
+  const fileExt = mime.extension(file.mimetype);
+  const fileName = `${datasource.id}${fileExt ? `.${fileExt}` : ''}`;
+
+  const params = {
+    Bucket: process.env.NEXT_PUBLIC_S3_BUCKET_NAME!,
+    Key: `datastores/${datastore.id}/${fileName}`,
+    Body: file.buffer,
+    ContentType: file.mimetype,
+    ACL: 'public-read',
+  };
+
+  await s3.putObject(params).promise();
+
+  // Trigger processing
   await triggerTaskLoadDatasource([
     {
       userId: datasource.ownerId!,
-      datasourceId,
+      datasourceId: datasource.id,
       priority: 1,
     },
   ]);
 
   return datasource;
 };
 
-handler.post(
-  validate({
-    body: Schema,
-    handler: respond(fileUpload),
-  })
-);
+handler.use(multer().single('file')).post(respond(upload));
+
+export const config = {
+  api: {
+    bodyParser: false,
+  },
+};
 
 export default async function wrapper(
   req: NextApiRequest,

diff --git a/pages/api/external/datastores/file-upload/process/[id].ts b/pages/api/external/datastores/file-upload/process/[id].ts
@@ -0,0 +1,112 @@
+import {
+  DatasourceStatus,
+  DatasourceType,
+  DatastoreVisibility,
+  SubscriptionPlan,
+  Usage,
+} from '@prisma/client';
+import Cors from 'cors';
+import { NextApiRequest, NextApiResponse } from 'next';
+import { z } from 'zod';
+
+import { AppNextApiRequest } from '@app/types/index';
+import { ApiError, ApiErrorType } from '@app/utils/api-error';
+import { createApiHandler, respond } from '@app/utils/createa-api-handler';
+import guardDataProcessingUsage from '@app/utils/guard-data-processing-usage';
+import prisma from '@app/utils/prisma-client';
+import runMiddleware from '@app/utils/run-middleware';
+import triggerTaskLoadDatasource from '@app/utils/trigger-task-load-datasource';
+import validate from '@app/utils/validate';
+
+const cors = Cors({
+  methods: ['POST', 'HEAD'],
+});
+
+const handler = createApiHandler();
+
+const Schema = z.object({
+  id: z.string().cuid(),
+});
+
+export const processUpload = async (
+  req: AppNextApiRequest,
+  res: NextApiResponse
+) => {
+  const datastoreId = req.query.id as string;
+
+  const data = req.body as z.infer<typeof Schema>;
+  const datasourceId = data?.id;
+
+  // get Bearer token from header
+  const authHeader = req.headers.authorization;
+  const token = authHeader && authHeader.split(' ')?.[1];
+
+  if (!datastoreId) {
+    throw new ApiError(ApiErrorType.INVALID_REQUEST);
+  }
+
+  const datasource = await prisma.appDatasource.findUnique({
+    where: {
+      id: datasourceId,
+    },
+    include: {
+      datastore: true,
+      owner: {
+        include: {
+          usage: true,
+          apiKeys: true,
+          subscriptions: {
+            where: {
+              status: 'active',
+            },
+          },
+        },
+      },
+    },
+  });
+
+  if (!datasource) {
+    throw new ApiError(ApiErrorType.NOT_FOUND);
+  }
+
+  if (
+    datasource?.datastoreId !== datastoreId ||
+    (datasource?.datastore?.visibility === DatastoreVisibility.private &&
+      (!token ||
+        !datasource?.owner?.apiKeys.find((each) => each.key === token)))
+  ) {
+    throw new ApiError(ApiErrorType.UNAUTHORIZED);
+  }
+
+  guardDataProcessingUsage({
+    usage: datasource?.owner?.usage as Usage,
+    plan:
+      datasource?.owner?.subscriptions?.[0]?.plan || SubscriptionPlan.level_0,
+  });
+
+  await triggerTaskLoadDatasource([
+    {
+      userId: datasource.ownerId!,
+      datasourceId,
+      priority: 1,
+    },
+  ]);
+
+  return datasource;
+};
+
+handler.post(
+  validate({
+    body: Schema,
+    handler: respond(processUpload),
+  })
+);
+
+export default async function wrapper(
+  req: NextApiRequest,
+  res: NextApiResponse
+) {
+  await runMiddleware(req, res, cors);
+
+  return handler(req, res);
+}