v4.6 -1 (#459)

packages/service/common/file/gridfs/controller.ts

@@ -0,0 +1,111 @@
+import { Types, connectionMongo } from '../../mongo';
+import { BucketNameEnum } from '@fastgpt/global/common/file/constants';
+import fsp from 'fs/promises';
+import fs from 'fs';
+import { DatasetFileSchema } from '@fastgpt/global/core/dataset/type';
+
+export function getGFSCollection(bucket: `${BucketNameEnum}`) {
+  return connectionMongo.connection.db.collection(`${bucket}.files`);
+}
+export function getGridBucket(bucket: `${BucketNameEnum}`) {
+  return new connectionMongo.mongo.GridFSBucket(connectionMongo.connection.db, {
+    bucketName: bucket
+  });
+}
+
+/* crud  file */
+export async function uploadFile({
+  bucketName,
+  teamId,
+  tmbId,
+  path,
+  filename,
+  metadata = {}
+}: {
+  bucketName: `${BucketNameEnum}`;
+  teamId: string;
+  tmbId: string;
+  path: string;
+  filename: string;
+  metadata?: Record<string, any>;
+}) {
+  if (!path) return Promise.reject(`filePath is empty`);
+  if (!filename) return Promise.reject(`filename is empty`);
+
+  const stats = await fsp.stat(path);
+  if (!stats.isFile()) return Promise.reject(`${path} is not a file`);
+
+  metadata.teamId = teamId;
+  metadata.tmbId = tmbId;
+
+  // create a gridfs bucket
+  const bucket = getGridBucket(bucketName);
+
+  const stream = bucket.openUploadStream(filename, {
+    metadata,
+    contentType: metadata?.contentType
+  });
+
+  // save to gridfs
+  await new Promise((resolve, reject) => {
+    fs.createReadStream(path)
+      .pipe(stream as any)
+      .on('finish', resolve)
+      .on('error', reject);
+  });
+
+  return String(stream.id);
+}
+
+export async function getFileById({
+  bucketName,
+  fileId
+}: {
+  bucketName: `${BucketNameEnum}`;
+  fileId: string;
+}) {
+  const db = getGFSCollection(bucketName);
+  const file = await db.findOne<DatasetFileSchema>({
+    _id: new Types.ObjectId(fileId)
+  });
+
+  if (!file) {
+    return Promise.reject('File not found');
+  }
+
+  return file;
+}
+
+export async function delFileById({
+  bucketName,
+  fileId
+}: {
+  bucketName: `${BucketNameEnum}`;
+  fileId: string;
+}) {
+  const bucket = getGridBucket(bucketName);
+
+  await bucket.delete(new Types.ObjectId(fileId));
+  return true;
+}
+
+export async function getDownloadBuf({
+  bucketName,
+  fileId
+}: {
+  bucketName: `${BucketNameEnum}`;
+  fileId: string;
+}) {
+  const bucket = getGridBucket(bucketName);
+
+  const stream = bucket.openDownloadStream(new Types.ObjectId(fileId));
+
+  const buf: Buffer = await new Promise((resolve, reject) => {
+    const buffers: Buffer[] = [];
+    stream.on('data', (data) => buffers.push(data));
+    stream.on('error', reject);
+    stream.on('end', () => resolve(Buffer.concat(buffers)));
+  });
+
+  return buf;
+}