Allow user to configure database:import chunk size and concurrency (#5726)

* Allow user to configure database:import chunk size and concurrency

* Cap chunk size at NGINX max payload size

Author: tohhsinpei

src/commands/database-import.ts

@@ -21,8 +21,14 @@ interface DatabaseImportOptions extends Options {
   instanceDetails: DatabaseInstance;
   disableTriggers?: boolean;
   filter?: string;
+  chunkSize?: string;
+  concurrency?: string;
 }
 
+const MAX_CHUNK_SIZE_MB = 10;
+const MAX_PAYLOAD_SIZE_MB = 256;
+const CONCURRENCY_LIMIT = 5;
+
 export const command = new Command("database:import <path> [infile]")
   .description(
     "non-atomically import the contents of a JSON file to the specified path in Realtime Database"
@@ -41,6 +47,8 @@ export const command = new Command("database:import <path> [infile]")
     "--filter <dataPath>",
     "import only data at this path in the JSON file (if omitted, import entire file)"
   )
+  .option("--chunk-size <mb>", "max chunk size in megabytes, default to 10 MB")
+  .option("--concurrency <val>", "concurrency limit, default to 5")
   .before(requirePermissions, ["firebasedatabase.instances.update"])
   .before(requireDatabaseInstance)
   .before(populateInstanceDetails)
@@ -54,6 +62,11 @@ export const command = new Command("database:import <path> [infile]")
       throw new FirebaseError("No file supplied");
     }
 
+    const chunkMegabytes = options.chunkSize ? parseInt(options.chunkSize, 10) : MAX_CHUNK_SIZE_MB;
+    if (chunkMegabytes > MAX_PAYLOAD_SIZE_MB) {
+      throw new FirebaseError("Max chunk size cannot exceed 256 MB");
+    }
+
     const projectId = needProjectId(options);
     const origin = realtimeOriginOrEmulatorOrCustomUrl(options.instanceDetails.databaseUrl);
     const dbPath = utils.getDatabaseUrl(origin, options.instance, path);
@@ -77,7 +90,9 @@ export const command = new Command("database:import <path> [infile]")
 
     const inStream = fs.createReadStream(infile);
     const dataPath = options.filter || "";
-    const importer = new DatabaseImporter(dbUrl, inStream, dataPath);
+    const chunkBytes = chunkMegabytes * 1024 * 1024;
+    const concurrency = options.concurrency ? parseInt(options.concurrency, 10) : CONCURRENCY_LIMIT;
+    const importer = new DatabaseImporter(dbUrl, inStream, dataPath, chunkBytes, concurrency);
 
     let responses;
     try {

src/database/import.ts

@@ -9,9 +9,6 @@ import { Client, ClientResponse } from "../apiv2";
 import { FirebaseError } from "../error";
 import * as pLimit from "p-limit";
 
-const MAX_CHUNK_SIZE = 1024 * 1024 * 10;
-const CONCURRENCY_LIMIT = 5;
-
 type JsonType = { [key: string]: JsonType } | string | number | boolean;
 
 type Data = {
@@ -31,15 +28,17 @@ type ChunkedData = {
  */
 export default class DatabaseImporter {
   private client: Client;
-  private limit = pLimit(CONCURRENCY_LIMIT);
+  private limit: pLimit.Limit;
 
   constructor(
     private dbUrl: URL,
     private inStream: stream.Readable,
     private dataPath: string,
-    private chunkSize = MAX_CHUNK_SIZE
+    private chunkBytes: number,
+    concurrency: number
   ) {
     this.client = new Client({ urlPrefix: dbUrl.origin, auth: true });
+    this.limit = pLimit(concurrency);
   }
 
   /**
@@ -157,7 +156,7 @@ export default class DatabaseImporter {
         }
       }
 
-      if (hasChunkedChild || size >= this.chunkSize) {
+      if (hasChunkedChild || size >= this.chunkBytes) {
         return { chunks, size };
       } else {
         return { chunks: null, size };

src/test/database/import.spec.ts

@@ -7,6 +7,8 @@ import DatabaseImporter from "../../database/import";
 import { FirebaseError } from "../../error";
 
 const dbUrl = new URL("https://test-db.firebaseio.com/foo");
+const chunkSize = 1024 * 1024 * 10;
+const concurrencyLimit = 5;
 
 describe("DatabaseImporter", () => {
   const DATA = { a: 100, b: [true, "bar", { f: { g: 0, h: 1 }, i: "baz" }] };
@@ -22,7 +24,9 @@ describe("DatabaseImporter", () => {
     const importer = new DatabaseImporter(
       dbUrl,
       utils.stringToStream(INVALID_JSON)!,
-      /* importPath= */ "/"
+      /* importPath= */ "/",
+      chunkSize,
+      concurrencyLimit
     );
 
     await expect(importer.execute()).to.be.rejectedWith(
@@ -37,7 +41,13 @@ describe("DatabaseImporter", () => {
     nock("https://test-db.firebaseio.com")
       .put("/foo/b.json", JSON.stringify([true, "bar", { f: { g: 0, h: 1 }, i: "baz" }]))
       .reply(200);
-    const importer = new DatabaseImporter(dbUrl, DATA_STREAM, /* importPath= */ "/");
+    const importer = new DatabaseImporter(
+      dbUrl,
+      DATA_STREAM,
+      /* importPath= */ "/",
+      chunkSize,
+      concurrencyLimit
+    );
 
     const responses = await importer.execute();
 
@@ -58,7 +68,8 @@ describe("DatabaseImporter", () => {
       dbUrl,
       DATA_STREAM,
       /* importPath= */ "/",
-      /* chunkSize= */ 20
+      /* chunkSize= */ 20,
+      concurrencyLimit
     );
 
     const responses = await importer.execute();
@@ -72,7 +83,13 @@ describe("DatabaseImporter", () => {
     nock("https://test-db.firebaseio.com")
       .put("/foo/b.json", JSON.stringify([true, "bar", { f: { g: 0, h: 1 }, i: "baz" }]))
       .reply(200);
-    const importer = new DatabaseImporter(dbUrl, DATA_STREAM, /* importPath= */ "/b");
+    const importer = new DatabaseImporter(
+      dbUrl,
+      DATA_STREAM,
+      /* importPath= */ "/b",
+      chunkSize,
+      concurrencyLimit
+    );
 
     const responses = await importer.execute();
 
@@ -82,7 +99,13 @@ describe("DatabaseImporter", () => {
 
   it("throws FirebaseError when data location is nonempty", async () => {
     nock("https://test-db.firebaseio.com").get("/foo.json?shallow=true").reply(200, { a: "foo" });
-    const importer = new DatabaseImporter(dbUrl, DATA_STREAM, /* importPath= */ "/");
+    const importer = new DatabaseImporter(
+      dbUrl,
+      DATA_STREAM,
+      /* importPath= */ "/",
+      chunkSize,
+      concurrencyLimit
+    );
 
     await expect(importer.execute()).to.be.rejectedWith(
       FirebaseError,