From 56f015ffd5981862b2904d793497a4aac650dbea Mon Sep 17 00:00:00 2001
From: Jan Prochazka <prochazka@sprinx.cz>
Date: Tue, 17 Sep 2024 13:23:51 +0200
Subject: [PATCH] JSON import rootField support

---
 .../__tests__/import-formats.spec.js          | 29 +++++++++++++++++++
 packages/api/src/shell/jsonReader.js          | 26 ++++++++++++-----
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/integration-tests/__tests__/import-formats.spec.js b/integration-tests/__tests__/import-formats.spec.js
index 7bfcc905..21dc38e2 100644
--- a/integration-tests/__tests__/import-formats.spec.js
+++ b/integration-tests/__tests__/import-formats.spec.js
@@ -120,3 +120,32 @@ test('JSON object import test', async () => {
     { mykey: 'k2', id: 2, val: 'v2' },
   ]);
 });
+
+test('JSON filtered object import test', async () => {
+  const jsonFileName = tmp.tmpNameSync();
+
+  fs.writeFileSync(
+    jsonFileName,
+    JSON.stringify({
+      filtered: {
+        k1: { id: 1, val: 'v1' },
+        k2: { id: 2, val: 'v2' },
+      },
+    })
+  );
+
+  const reader = await dbgateApi.jsonReader({
+    fileName: jsonFileName,
+    jsonStyle: 'object',
+    keyField: 'mykey',
+    rootField: 'filtered',
+  });
+
+  const rows = await getReaderRows(reader);
+
+  expect(rows.length).toEqual(2);
+  expect(rows).toEqual([
+    { mykey: 'k1', id: 1, val: 'v1' },
+    { mykey: 'k2', id: 2, val: 'v2' },
+  ]);
+});
diff --git a/packages/api/src/shell/jsonReader.js b/packages/api/src/shell/jsonReader.js
index 74a76cee..1347595e 100644
--- a/packages/api/src/shell/jsonReader.js
+++ b/packages/api/src/shell/jsonReader.js
@@ -2,12 +2,13 @@ const fs = require('fs');
 const stream = require('stream');
 const byline = require('byline');
 const { getLogger } = require('dbgate-tools');
-const logger = getLogger('jsonReader');
 const { parser } = require('stream-json');
 const { pick } = require('stream-json/filters/Pick');
 const { streamArray } = require('stream-json/streamers/StreamArray');
 const { streamObject } = require('stream-json/streamers/StreamObject');
 
+const logger = getLogger('jsonReader');
+
 class ParseStream extends stream.Transform {
   constructor({ limitRows, jsonStyle, keyField }) {
     super({ objectMode: true });
@@ -42,7 +43,14 @@ class ParseStream extends stream.Transform {
   }
 }
 
-async function jsonReader({ fileName, jsonStyle, keyField = '_key', encoding = 'utf-8', limitRows = undefined }) {
+async function jsonReader({
+  fileName,
+  jsonStyle,
+  keyField = '_key',
+  rootField = null,
+  encoding = 'utf-8',
+  limitRows = undefined,
+}) {
   logger.info(`Reading file ${fileName}`);
 
   const fileStream = fs.createReadStream(
@@ -55,16 +63,18 @@ async function jsonReader({ fileName, jsonStyle, keyField = '_key', encoding = '
 
   const parseStream = new ParseStream({ limitRows, jsonStyle, keyField });
 
-  if (jsonStyle === 'object') {
-    const tramsformer = streamObject();
-    parseJsonStream.pipe(tramsformer);
-    tramsformer.pipe(parseStream);
+  const tramsformer = jsonStyle === 'object' ? streamObject() : streamArray();
+
+  if (rootField) {
+    const filterStream = pick({ filter: rootField });
+    parseJsonStream.pipe(filterStream);
+    filterStream.pipe(tramsformer);
   } else {
-    const tramsformer = streamArray();
     parseJsonStream.pipe(tramsformer);
-    tramsformer.pipe(parseStream);
   }
 
+  tramsformer.pipe(parseStream);
+
   return parseStream;
 }