Stream top-level JSON objects

Author: tohhsinpei

package.json

@@ -176,6 +176,7 @@
     "@types/glob": "^7.1.1",
     "@types/inquirer": "^8.1.3",
     "@types/js-yaml": "^3.12.2",
+    "@types/jsonstream": "^0.8.30",
     "@types/jsonwebtoken": "^8.3.8",
     "@types/libsodium-wrappers": "^0.7.9",
     "@types/lodash": "^4.14.149",

src/commands/database-import.ts

@@ -59,8 +59,8 @@ export const command = new Command("database:import <path> [infile]")
       throw new FirebaseError("Command aborted.");
     }
 
-    const inputString = await utils.streamToString(fs.createReadStream(infile));
-    const importer = new DatabaseImporter(dbUrl, inputString);
+    const inStream = fs.createReadStream(infile);
+    const importer = new DatabaseImporter(dbUrl, inStream);
     try {
       await importer.execute();
     } catch (err: any) {

src/database/import.ts

@@ -1,9 +1,12 @@
 import * as clc from "colorette";
+import * as stream from "stream";
 import pLimit from "p-limit";
 import { URL } from "url";
 import { Client } from "../apiv2";
 import { FirebaseError } from "../error";
 
+const JSONStream = require("JSONStream");
+
 const MAX_CHUNK_SIZE = 1024 * 1024;
 const CONCURRENCY_LIMIT = 5;
 
@@ -23,23 +26,14 @@ type ChunkedData = {
  * The data is parsed and chunked into subtrees of ~1 MB, to be subsequently written in parallel.
  */
 export default class DatabaseImporter {
-  chunks: Data[];
   private client: Client;
   private limit = pLimit(CONCURRENCY_LIMIT);
 
-  constructor(private dbUrl: URL, file: string, private chunkSize = MAX_CHUNK_SIZE) {
-    let data;
-    try {
-      data = { json: JSON.parse(file), pathname: dbUrl.pathname };
-    } catch (err: any) {
-      throw new FirebaseError("Invalid data; couldn't parse JSON object, array, or value.", {
-        original: err,
-        exit: 2,
-      });
-    }
-
-    const chunkedData = this.chunkData(data);
-    this.chunks = chunkedData.chunks || [data];
+  constructor(
+    private dbUrl: URL,
+    private inStream: NodeJS.ReadableStream,
+    private chunkSize = MAX_CHUNK_SIZE
+  ) {
     this.client = new Client({ urlPrefix: dbUrl.origin, auth: true });
   }
 
@@ -48,7 +42,7 @@ export default class DatabaseImporter {
    */
   async execute(): Promise<any> {
     await this.checkLocationIsEmpty();
-    return Promise.all(this.chunks.map(this.writeChunk.bind(this)));
+    return this.readAndWriteChunks(this.inStream);
   }
 
   private async checkLocationIsEmpty(): Promise<void> {
@@ -68,6 +62,48 @@ export default class DatabaseImporter {
     }
   }
 
+  private readAndWriteChunks(inStream: NodeJS.ReadableStream): Promise<any> {
+    const { dbUrl } = this;
+    const chunkData = this.chunkData.bind(this);
+    const writeChunk = this.writeChunk.bind(this);
+    const getJoinedPath = this.getJoinedPath.bind(this);
+
+    const readChunks = new stream.Transform({ objectMode: true });
+    readChunks._transform = function (chunk: { key: string; value: any }, _, done) {
+      const data = { json: chunk.value, pathname: getJoinedPath(dbUrl.pathname, chunk.key) };
+      const chunkedData = chunkData(data);
+      const chunks = chunkedData.chunks || [data];
+      chunks.forEach((chunk: Data) => this.push(chunk));
+      done();
+    };
+
+    const writeChunks = new stream.Transform({ objectMode: true });
+    writeChunks._transform = async function (chunk: Data, _, done) {
+      const res = await writeChunk(chunk);
+      this.push(res);
+      done();
+    };
+
+    return new Promise((resolve, reject) => {
+      const results: any[] = [];
+      inStream
+        .pipe(JSONStream.parse("$*"))
+        .on("error", (err: any) =>
+          reject(
+            new FirebaseError("Invalid data; couldn't parse JSON object, array, or value.", {
+              original: err,
+              exit: 2,
+            })
+          )
+        )
+        .pipe(readChunks)
+        .pipe(writeChunks)
+        .on("data", (res: any) => results.push(res))
+        .on("error", reject)
+        .once("end", () => resolve(results));
+    });
+  }
+
   private writeChunk(chunk: Data): Promise<any> {
     return this.limit(() =>
       this.client.request({
@@ -91,9 +127,9 @@ export default class DatabaseImporter {
       let hasChunkedChild = false;
 
       for (const key of Object.keys(json)) {
-        size += key.length + 3; // "[key]":
+        size += key.length + 3; // "":
 
-        const child = { json: json[key], pathname: [pathname, key].join("/").replace("//", "/") };
+        const child = { json: json[key], pathname: this.getJoinedPath(pathname, key) };
         const childChunks = this.chunkData(child);
         size += childChunks.size;
         if (childChunks.chunks) {
@@ -111,4 +147,8 @@ export default class DatabaseImporter {
       }
     }
   }
+
+  private getJoinedPath(root: string, key: string): string {
+    return [root, key].join("/").replace("//", "/");
+  }
 }

src/test/database/import.spec.ts

@@ -1,4 +1,5 @@
 import * as nock from "nock";
+import * as utils from "../../utils";
 import { expect } from "chai";
 
 import DatabaseImporter from "../../database/import";
@@ -8,33 +9,37 @@ const dbUrl = new URL("https://test-db.firebaseio.com/foo");
 
 describe("DatabaseImporter", () => {
   const DATA = { a: 100, b: [true, "bar", { f: { g: 0, h: 1 }, i: "baz" }] };
+  let DATA_STREAM: NodeJS.ReadableStream;
 
-  it("parses data as single chunk", () => {
-    const importer = new DatabaseImporter(dbUrl, JSON.stringify(DATA));
-    expect(importer.chunks.length).to.equal(1);
-    expect(importer.chunks[0].json).to.deep.equal(DATA);
-    expect(importer.chunks[0].pathname).to.equal("/foo");
+  beforeEach(() => {
+    DATA_STREAM = utils.stringToStream(JSON.stringify(DATA))!;
   });
 
-  it("parses data as multiple chunks", () => {
-    const importer = new DatabaseImporter(dbUrl, JSON.stringify(DATA), /* chunkSize= */ 20);
-    expect(importer.chunks.length).to.equal(5);
-    expect(importer.chunks).to.deep.include({ json: 100, pathname: "/foo/a" });
-    expect(importer.chunks).to.deep.include({ json: true, pathname: "/foo/b/0" });
-    expect(importer.chunks).to.deep.include({ json: "bar", pathname: "/foo/b/1" });
-    expect(importer.chunks).to.deep.include({ json: { g: 0, h: 1 }, pathname: "/foo/b/2/f" });
-    expect(importer.chunks).to.deep.include({ json: "baz", pathname: "/foo/b/2/i" });
-  });
+  it("throws FirebaseError when JSON is invalid", async () => {
+    nock("https://test-db.firebaseio.com").get("/foo.json?shallow=true").reply(200);
 
-  it("throws FirebaseError when JSON is invalid", () => {
-    const INVALID_JSON = '{"a": }';
-    expect(() => new DatabaseImporter(dbUrl, INVALID_JSON)).to.throw(
+    const INVALID_JSON = '{"a": {"b"}}';
+    const importer = new DatabaseImporter(dbUrl, utils.stringToStream(INVALID_JSON)!);
+    await expect(importer.execute()).to.be.rejectedWith(
       FirebaseError,
       "Invalid data; couldn't parse JSON object, array, or value."
     );
   });
 
-  it("sends multiple chunked requests", async () => {
+  it("chunks data in top-level objects", async () => {
+    nock("https://test-db.firebaseio.com").get("/foo.json?shallow=true").reply(200);
+    nock("https://test-db.firebaseio.com").put("/foo/a.json", "100").reply(200);
+    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);
+    const responses = await importer.execute();
+    expect(responses).to.have.length(2);
+    expect(nock.isDone()).to.be.true;
+  });
+
+  it("chunks data according to provided chunk size", async () => {
     nock("https://test-db.firebaseio.com").get("/foo.json?shallow=true").reply(200);
     nock("https://test-db.firebaseio.com").put("/foo/a.json", "100").reply(200);
     nock("https://test-db.firebaseio.com").put("/foo/b/0.json", "true").reply(200);
@@ -44,15 +49,15 @@ describe("DatabaseImporter", () => {
       .reply(200);
     nock("https://test-db.firebaseio.com").put("/foo/b/2/i.json", '"baz"').reply(200);
 
-    const importer = new DatabaseImporter(dbUrl, JSON.stringify(DATA), /* chunkSize= */ 20);
+    const importer = new DatabaseImporter(dbUrl, DATA_STREAM, /* chunkSize= */ 20);
     const responses = await importer.execute();
     expect(responses).to.have.length(5);
     expect(nock.isDone()).to.be.true;
   });
 
   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, JSON.stringify(DATA));
+    const importer = new DatabaseImporter(dbUrl, DATA_STREAM);
     await expect(importer.execute()).to.be.rejectedWith(
       FirebaseError,
       /Importing is only allowed for an empty location./