Use async/await synax

outline · Apr 3, 2024 · d520b03 · d520b03
1 parent 3c7adbd
commit d520b03
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 44 deletions.
diff --git a/plugins/storage/server/api/files.ts b/plugins/storage/server/api/files.ts
@@ -75,7 +75,7 @@ router.get(
           "application/octet-stream"
       );
       ctx.attachment(fileName);
-      ctx.body = FileStorage.getFileStream(key);
+      ctx.body = await FileStorage.getFileStream(key);
     } else {
       const attachment = await Attachment.findOne({
         where: { key },
@@ -86,7 +86,7 @@ router.get(
       ctx.set("Cache-Control", cacheHeader);
       ctx.set("Content-Type", attachment.contentType);
       ctx.attachment(attachment.name);
-      ctx.body = attachment.stream;
+      ctx.body = await attachment.stream;
     }
   }
 );

diff --git a/server/storage/files/BaseStorage.ts b/server/storage/files/BaseStorage.ts
@@ -31,7 +31,7 @@ export default abstract class BaseStorage {
    *
    * @param key The path to the file
    */
-  public abstract getFileStream(key: string): NodeJS.ReadableStream | null;
+  public abstract getFileStream(key: string): Promise<Readable>;
 
   /**
    * Returns the upload URL for the storage provider.
@@ -101,7 +101,7 @@ export default abstract class BaseStorage {
    * @param key The path to the file
    */
   public async getFileBuffer(key: string) {
-    const stream = this.getFileStream(key);
+    const stream = await this.getFileStream(key);
     return new Promise<Buffer>((resolve, reject) => {
       const chunks: Buffer[] = [];
       if (!stream) {

diff --git a/server/storage/files/LocalStorage.ts b/server/storage/files/LocalStorage.ts
@@ -131,8 +131,13 @@ export default class LocalStorage extends BaseStorage {
     };
   }
 
-  public getFileStream(key: string) {
-    return fs.createReadStream(this.getFilePath(key));
+  public getFileStream(key: string): Promise<Readable> {
+    return new Promise((resolve, reject) => {
+      const filePath = this.getFilePath(key);
+      const stream = fs.createReadStream(filePath);
+      stream.on("ready", () => resolve(stream));
+      stream.on("error", (error) => reject(error));
+    });
   }
 
   private getFilePath(key: string) {

diff --git a/server/storage/files/S3Storage.ts b/server/storage/files/S3Storage.ts
@@ -18,7 +18,6 @@ import invariant from "invariant";
 import compact from "lodash/compact";
 import tmp from "tmp";
 import env from "@server/env";
-import Logger from "@server/logging/Logger";
 import BaseStorage from "./BaseStorage";
 
 export default class S3Storage extends BaseStorage {
@@ -179,57 +178,57 @@ export default class S3Storage extends BaseStorage {
     }
   };
 
-  public getFileHandle(key: string): Promise<{
+  public async getFileHandle(key: string): Promise<{
     path: string;
     cleanup: () => Promise<void>;
   }> {
-    return new Promise((resolve, reject) => {
-      tmp.dir((err, tmpDir) => {
-        if (err) {
-          return reject(err);
-        }
-        const tmpFile = path.join(tmpDir, "tmp");
-        const dest = fs.createWriteStream(tmpFile);
+    try {
+      const tmpDir = await new Promise<string>((resolve, reject) => {
+        tmp.dir((err, p) => {
+          if (err) {
+            reject(err);
+          } else {
+            resolve(p);
+          }
+        });
+      });
+
+      const tmpFile = path.join(tmpDir, "tmp");
+      const dest = fs.createWriteStream(tmpFile);
+
+      const stream = await this.getFileStream(key);
+      if (!stream) {
+        throw new Error("No stream available");
+      }
+
+      return new Promise((resolve, reject) => {
         dest.on("error", reject);
         dest.on("finish", () =>
-          resolve({ path: tmpFile, cleanup: () => fs.rm(tmpFile) })
+          resolve({ path: tmpFile, cleanup: () => fs.promises.rm(tmpFile) })
         );
-
-        const stream = this.getFileStream(key);
-        if (!stream) {
-          return reject(new Error("No stream available"));
-        }
-
-        stream
-          .on("error", (err) => {
-            dest.end();
-            reject(err);
-          })
-          .pipe(dest);
+        stream.pipe(dest);
       });
-    });
+    } catch (err) {
+      return Promise.reject(err);
+    }
   }
 
-  public getFileStream(key: string): NodeJS.ReadableStream | null {
+  public async getFileStream(key: string) {
     invariant(
       env.AWS_S3_UPLOAD_BUCKET_NAME,
       "AWS_S3_UPLOAD_BUCKET_NAME is required"
     );
 
-    this.client
-      .send(
-        new GetObjectCommand({
-          Bucket: env.AWS_S3_UPLOAD_BUCKET_NAME,
-          Key: key,
-        })
-      )
-      .then((item) => item.Body as NodeJS.ReadableStream)
-      .catch((err) => {
-        Logger.error("Error getting file stream from S3 ", err, {
-          key,
-        });
-      });
-    return null;
+    const resp = await this.client.send(
+      new GetObjectCommand({
+        Bucket: env.AWS_S3_UPLOAD_BUCKET_NAME,
+        Key: key,
+      })
+    );
+    if (!resp.Body) {
+      throw new Error("Error getting file stream from S3");
+    }
+    return resp.Body as Readable;
   }
 
   private client: S3Client;