fix(lib-storage): use input parameter object size information if available

Author: kuhe

lib/lib-storage/src/Upload.spec.ts

@@ -1,11 +1,23 @@
+import {
+  CompleteMultipartUploadCommand,
+  CompleteMultipartUploadCommandOutput,
+  CreateMultipartUploadCommand,
+  PutObjectCommand,
+  PutObjectTaggingCommand,
+  S3,
+  S3Client,
+  UploadPartCommand,
+} from "@aws-sdk/client-s3";
+import { AbortController } from "@smithy/abort-controller";
+import { EventEmitter, Readable } from "stream";
 import { afterAll, afterEach, beforeEach, describe, expect, test as it, vi } from "vitest";
 
+import { Progress, Upload } from "./index";
+
 /* eslint-disable no-var */
 var hostname = "s3.region.amazonaws.com";
 var port: number | undefined;
 
-import { EventEmitter, Readable } from "stream";
-
 vi.mock("@aws-sdk/client-s3", async () => {
   const sendMock = vi.fn().mockImplementation(async (x) => x);
   const endpointMock = vi.fn().mockImplementation(() => ({
@@ -61,24 +73,20 @@ vi.mock("@aws-sdk/client-s3", async () => {
   };
 });
 
-import {
-  CompleteMultipartUploadCommand,
-  CompleteMultipartUploadCommandOutput,
-  CreateMultipartUploadCommand,
-  PutObjectCommand,
-  PutObjectTaggingCommand,
-  S3,
-  S3Client,
-  UploadPartCommand,
-} from "@aws-sdk/client-s3";
-import { AbortController } from "@smithy/abort-controller";
-
-import { Progress, Upload } from "./index";
-
 const DEFAULT_PART_SIZE = 1024 * 1024 * 5;
 
+type Expose = {
+  totalBytes: number | undefined;
+};
+type VisibleForTesting = Omit<Upload, keyof Expose> & Expose;
+
 describe(Upload.name, () => {
-  const s3MockInstance = new S3Client();
+  const s3MockInstance = new S3Client({
+    credentials: {
+      accessKeyId: "UNIT",
+      secretAccessKey: "UNIT",
+    },
+  });
 
   beforeEach(() => {
     vi.clearAllMocks();
@@ -107,6 +115,43 @@ describe(Upload.name, () => {
     Body: "this-is-a-sample-payload",
   };
 
+  it("uses the input parameters for object length if provided", async () => {
+    const falseFileStream = Object.assign(Readable.from("abcd"), {
+      path: "/dev/null",
+    });
+    let upload = new Upload({
+      params: {
+        Bucket: "",
+        Key: "",
+        Body: falseFileStream,
+        MpuObjectSize: 6 * 1024 * 1024,
+      },
+      client: s3MockInstance,
+    }) as unknown as VisibleForTesting;
+    expect(upload.totalBytes).toEqual(6 * 1024 * 1024);
+
+    upload = new Upload({
+      params: {
+        Bucket: "",
+        Key: "",
+        Body: falseFileStream,
+        ContentLength: 6 * 1024 * 1024,
+      },
+      client: s3MockInstance,
+    }) as unknown as VisibleForTesting;
+    expect(upload.totalBytes).toEqual(6 * 1024 * 1024);
+
+    upload = new Upload({
+      params: {
+        Bucket: "",
+        Key: "",
+        Body: falseFileStream,
+      },
+      client: s3MockInstance,
+    }) as unknown as VisibleForTesting;
+    expect(upload.totalBytes).toEqual(0);
+  });
+
   it("correctly exposes the event emitter API", () => {
     const upload = new Upload({
       params,

lib/lib-storage/src/Upload.ts

@@ -3,15 +3,18 @@ import {
   ChecksumAlgorithm,
   CompletedPart,
   CompleteMultipartUploadCommand,
+  CompleteMultipartUploadCommandInput,
   CompleteMultipartUploadCommandOutput,
   CreateMultipartUploadCommand,
+  CreateMultipartUploadCommandInput,
   CreateMultipartUploadCommandOutput,
   PutObjectCommand,
   PutObjectCommandInput,
   PutObjectTaggingCommand,
   S3Client,
   Tag,
   UploadPartCommand,
+  UploadPartCommandInput,
 } from "@aws-sdk/client-s3";
 import { AbortController } from "@smithy/abort-controller";
 import {
@@ -52,7 +55,8 @@ export class Upload extends EventEmitter {
   private readonly tags: Tag[] = [];
 
   private readonly client: S3Client;
-  private readonly params: PutObjectCommandInput;
+  private readonly params: PutObjectCommandInput &
+    Partial<CreateMultipartUploadCommandInput & UploadPartCommandInput & CompleteMultipartUploadCommandInput>;
 
   // used for reporting progress.
   private totalBytes?: number;
@@ -89,12 +93,15 @@ export class Upload extends EventEmitter {
     this.params = options.params;
 
     // set progress defaults
-    this.totalBytes = byteLength(this.params.Body);
+    this.totalBytes = this.params.MpuObjectSize ?? this.params.ContentLength ?? byteLength(this.params.Body);
     this.bytesUploadedSoFar = 0;
     this.abortController = options.abortController ?? new AbortController();
 
     this.partSize = Math.max(Upload.MIN_PART_SIZE, Math.floor((this.totalBytes || 0) / this.MAX_PARTS));
-    this.expectedPartsCount = this.totalBytes !== undefined ? Math.ceil(this.totalBytes / this.partSize) : undefined;
+
+    if (this.totalBytes !== undefined) {
+      this.expectedPartsCount = Math.ceil(this.totalBytes / this.partSize);
+    }
 
     this.__validateInput();
   }