Support importing at data path

Author: tohhsinpei

npm-shrinkwrap.json

@@ -26,6 +26,10 @@ export const command = new Command("database:import <path> [infile]")
     "suppress any Cloud functions triggered by this operation, default to true",
     true
   )
+  .option(
+    "--only <dataPath>",
+    "import only data at this path in the JSON file (if omitted, use root)"
+  )
   .before(requirePermissions, ["firebasedatabase.instances.update"])
   .before(requireDatabaseInstance)
   .before(populateInstanceDetails)
@@ -35,6 +39,10 @@ export const command = new Command("database:import <path> [infile]")
       throw new FirebaseError("Path must begin with /");
     }
 
+    if (options.only && !options.only.startsWith("/")) {
+      throw new FirebaseError("Data path must begin with /");
+    }
+
     if (!infile) {
       throw new FirebaseError("No file supplied");
     }
@@ -60,9 +68,12 @@ export const command = new Command("database:import <path> [infile]")
     }
 
     const inStream = fs.createReadStream(infile);
-    const importer = new DatabaseImporter(dbUrl, inStream);
+    const dataPath = options.only || "/";
+    const importer = new DatabaseImporter(dbUrl, inStream, dataPath);
+
+    let responses;
     try {
-      await importer.execute();
+      responses = await importer.execute();
     } catch (err: any) {
       if (err instanceof FirebaseError) {
         throw err;
@@ -71,7 +82,12 @@ export const command = new Command("database:import <path> [infile]")
       throw new FirebaseError(`Unexpected error while importing data: ${err}`, { exit: 2 });
     }
 
-    utils.logSuccess("Data persisted successfully");
+    if (responses.length) {
+      utils.logSuccess("Data persisted successfully");
+    } else {
+      utils.logWarning("No data was persisted. Check the data path supplied.");
+    }
+
     logger.info();
     logger.info(
       clc.bold("View data at:"),

src/commands/database-import.ts

@@ -1,6 +1,7 @@
 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";
@@ -26,14 +27,17 @@ type ChunkedData = {
  * The data is parsed and chunked into subtrees of ~1 MB, to be subsequently written in parallel.
  */
 export default class DatabaseImporter {
+  private jsonPath: string;
   private client: Client;
   private limit = pLimit(CONCURRENCY_LIMIT);
 
   constructor(
     private dbUrl: URL,
     private inStream: NodeJS.ReadableStream,
+    dataPath: string,
     private chunkSize = MAX_CHUNK_SIZE
   ) {
+    this.jsonPath = this.computeJsonPath(dataPath);
     this.client = new Client({ urlPrefix: dbUrl.origin, auth: true });
   }
 
@@ -66,7 +70,7 @@ export default class DatabaseImporter {
     const { dbUrl } = this;
     const chunkData = this.chunkData.bind(this);
     const writeChunk = this.writeChunk.bind(this);
-    const getJoinedPath = this.getJoinedPath.bind(this);
+    const getJoinedPath = this.joinPath.bind(this);
 
     const readChunks = new stream.Transform({ objectMode: true });
     readChunks._transform = function (chunk: { key: string; value: any }, _, done) {
@@ -85,9 +89,9 @@ export default class DatabaseImporter {
     };
 
     return new Promise((resolve, reject) => {
-      const results: any[] = [];
+      const responses: any[] = [];
       inStream
-        .pipe(JSONStream.parse("$*"))
+        .pipe(JSONStream.parse(this.jsonPath))
         .on("error", (err: any) =>
           reject(
             new FirebaseError("Invalid data; couldn't parse JSON object, array, or value.", {
@@ -98,9 +102,9 @@ export default class DatabaseImporter {
         )
         .pipe(readChunks)
         .pipe(writeChunks)
-        .on("data", (res: any) => results.push(res))
+        .on("data", (res: any) => responses.push(res))
         .on("error", reject)
-        .once("end", () => resolve(results));
+        .once("end", () => resolve(responses));
     });
   }
 
@@ -129,7 +133,7 @@ export default class DatabaseImporter {
       for (const key of Object.keys(json)) {
         size += key.length + 3; // "":
 
-        const child = { json: json[key], pathname: this.getJoinedPath(pathname, key) };
+        const child = { json: json[key], pathname: this.joinPath(pathname, key) };
         const childChunks = this.chunkData(child);
         size += childChunks.size;
         if (childChunks.chunks) {
@@ -148,7 +152,15 @@ export default class DatabaseImporter {
     }
   }
 
-  private getJoinedPath(root: string, key: string): string {
+  private computeJsonPath(dataPath: string): string {
+    if (dataPath === "/") {
+      return "$*";
+    } else {
+      return `${dataPath.split("/").slice(1).join(".")}.$*`;
+    }
+  }
+
+  private joinPath(root: string, key: string): string {
     return [root, key].join("/").replace("//", "/");
   }
 }

src/database/import.ts

@@ -7,7 +7,7 @@ import { FirebaseError } from "../../error";
 
 const dbUrl = new URL("https://test-db.firebaseio.com/foo");
 
-describe("DatabaseImporter", () => {
+describe.only("DatabaseImporter", () => {
   const DATA = { a: 100, b: [true, "bar", { f: { g: 0, h: 1 }, i: "baz" }] };
   let DATA_STREAM: NodeJS.ReadableStream;
 
@@ -19,7 +19,11 @@ describe("DatabaseImporter", () => {
     nock("https://test-db.firebaseio.com").get("/foo.json?shallow=true").reply(200);
 
     const INVALID_JSON = '{"a": {"b"}}';
-    const importer = new DatabaseImporter(dbUrl, utils.stringToStream(INVALID_JSON)!);
+    const importer = new DatabaseImporter(
+      dbUrl,
+      utils.stringToStream(INVALID_JSON)!,
+      /* importPath= */ "/"
+    );
     await expect(importer.execute()).to.be.rejectedWith(
       FirebaseError,
       "Invalid data; couldn't parse JSON object, array, or value."
@@ -33,7 +37,7 @@ describe("DatabaseImporter", () => {
       .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 importer = new DatabaseImporter(dbUrl, DATA_STREAM, /* importPath= */ "/");
     const responses = await importer.execute();
     expect(responses).to.have.length(2);
     expect(nock.isDone()).to.be.true;
@@ -49,15 +53,34 @@ describe("DatabaseImporter", () => {
       .reply(200);
     nock("https://test-db.firebaseio.com").put("/foo/b/2/i.json", '"baz"').reply(200);
 
-    const importer = new DatabaseImporter(dbUrl, DATA_STREAM, /* chunkSize= */ 20);
+    const importer = new DatabaseImporter(
+      dbUrl,
+      DATA_STREAM,
+      /* importPath= */ "/",
+      /* chunkSize= */ 20
+    );
     const responses = await importer.execute();
     expect(responses).to.have.length(5);
     expect(nock.isDone()).to.be.true;
   });
 
+  it("imports from data path", async () => {
+    nock("https://test-db.firebaseio.com").get("/foo.json?shallow=true").reply(200);
+    nock("https://test-db.firebaseio.com").put("/foo/0.json", "true").reply(200);
+    nock("https://test-db.firebaseio.com").put("/foo/1.json", '"bar"').reply(200);
+    nock("https://test-db.firebaseio.com")
+      .put("/foo/2.json", JSON.stringify({ f: { g: 0, h: 1 }, i: "baz" }))
+      .reply(200);
+
+    const importer = new DatabaseImporter(dbUrl, DATA_STREAM, /* importPath= */ "/b");
+    const responses = await importer.execute();
+    expect(responses).to.have.length(3);
+    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, DATA_STREAM);
+    const importer = new DatabaseImporter(dbUrl, DATA_STREAM, /* importPath= */ "/");
     await expect(importer.execute()).to.be.rejectedWith(
       FirebaseError,
       /Importing is only allowed for an empty location./