feat: FIR-43724 implement streaming in node sdk

Author: bogdantruta-firebolt

.gitignore

@@ -9,3 +9,4 @@ npm-debug.log*
 .#*
 coverage/
 allure-results/
+.idea

package-lock.json

@@ -10,6 +10,8 @@ import { CompositeError } from "../common/errors";
 import JSONbig from "json-bigint";
 import { QueryFormatter } from "../formatter/base";
 import { AsyncStatement } from "../statement/async";
+import { StreamStatement } from "../statement/stream";
+import { Response, Headers } from "node-fetch";
 
 const defaultQuerySettings = {
   output_format: OutputFormat.COMPACT
@@ -166,6 +168,11 @@ export abstract class Connection {
     executeQueryOptions?: ExecuteQueryOptions
   ): Promise<AsyncStatement>;
 
+  abstract executeStream(
+    query: string,
+    executeQueryOptions?: ExecuteQueryOptions
+  ): Promise<StreamStatement>;
+
   abstract isAsyncQueryRunning(token: string): Promise<boolean>;
 
   abstract isAsyncQuerySuccessful(token: string): Promise<boolean | undefined>;

src/connection/base.ts

@@ -3,6 +3,7 @@ import { Connection as BaseConnection } from "./base";
 import { ResourceManager } from "../service";
 import { ExecuteQueryOptions } from "../types";
 import { AsyncStatement } from "../statement/async";
+import { StreamStatement } from "../statement/stream";
 
 export interface AccountInfo {
   id: string;
@@ -93,4 +94,13 @@ export class ConnectionV1 extends BaseConnection {
       "Asynchronous execution is not supported in this Firebolt version."
     );
   }
+
+  async executeStream(
+    query: string,
+    executeQueryOptions?: ExecuteQueryOptions
+  ): Promise<StreamStatement> {
+    throw new Error(
+      "Asynchronous execution is not supported in this Firebolt version."
+    );
+  }
 }

src/connection/connection_v1.ts

@@ -3,8 +3,9 @@ import { ACCOUNT_SYSTEM_ENGINE, QUERY_URL } from "../common/api";
 
 import { Connection as BaseConnection } from "./base";
 import { Cache, inMemoryCache, noneCache } from "../common/tokenCache";
-import { ExecuteQueryOptions } from "../types";
+import { ExecuteQueryOptions, OutputFormat } from "../types";
 import { AsyncStatement } from "../statement/async";
+import { StreamStatement } from "../statement/stream";
 
 export class ConnectionV2 extends BaseConnection {
   private get account(): string {
@@ -103,6 +104,24 @@ export class ConnectionV2 extends BaseConnection {
     });
   }
 
+  async executeStream(
+    query: string,
+    executeQueryOptions: ExecuteQueryOptions = {}
+  ): Promise<StreamStatement> {
+    const { response } = await this.prepareAndExecuteQuery(query, {
+      ...executeQueryOptions,
+      settings: {
+        ...executeQueryOptions?.settings,
+        output_format: OutputFormat.JSON_LINES
+      }
+    });
+
+    return new StreamStatement({
+      response,
+      executeQueryOptions
+    });
+  }
+
   private async getAsyncQueryInfo(token: string) {
     const query = `CALL fb_GetAsyncStatus('${token}')`;
 

src/connection/connection_v2.ts

@@ -1,7 +1,7 @@
 import {
-  QueryResponse,
-  OutputFormat,
   ExecuteQueryOptions,
+  OutputFormat,
+  QueryResponse,
   Row
 } from "../types";
 import { Meta } from "../meta";
@@ -35,11 +35,11 @@ export const normalizeColumn = (column: { name: string; type: string }) => {
   return new Meta(column);
 };
 
-const getNormalizedMeta = (response: ParsedResponse): Meta[] => {
-  if (!response.meta) {
+export const getNormalizedMeta = (meta: any[]): Meta[] => {
+  if (!meta) {
     return [];
   }
-  return response.meta.map(normalizeColumn);
+  return meta.map(normalizeColumn);
 };
 
 export const getNormalizedStatistics = (response: ParsedResponse) => {
@@ -71,7 +71,7 @@ export const normalizeResponse = (
 
   const hydrate = executeQueryOptions?.response?.hydrateRow || hydrateRow;
 
-  const meta = getNormalizedMeta(response);
+  const meta = getNormalizedMeta(response.meta);
 
   const statistics = getNormalizedStatistics(response);
 
@@ -96,3 +96,21 @@ export const normalizeResponse = (
     statistics
   };
 };
+
+export const normalizeResponseRowStreaming = (
+  data: any[],
+  executeQueryOptions: ExecuteQueryOptions,
+  meta: Meta[]
+): Row[] => {
+  const { response: { normalizeData = false } = {} } = executeQueryOptions;
+
+  const hydrate = executeQueryOptions?.response?.hydrateRow || hydrateRow;
+
+  return data.map((row: Row) => {
+    const hydratedRow = hydrate(row, meta, executeQueryOptions);
+    if (normalizeData) {
+      return normalizeRow(hydratedRow, meta, executeQueryOptions);
+    }
+    return hydratedRow;
+  });
+};

src/statement/normalizeResponse.ts

@@ -0,0 +1,28 @@
+import { ExecuteQueryOptions } from "../types";
+import { ServerSideStream } from "./stream/serverSideStream";
+import { Response } from "node-fetch";
+
+export class StreamStatement {
+  private readonly executeQueryOptions: ExecuteQueryOptions;
+  private readonly response: Response;
+
+  constructor({
+    response,
+    executeQueryOptions
+  }: {
+    response: Response;
+    executeQueryOptions: ExecuteQueryOptions;
+  }) {
+    this.response = response;
+    this.executeQueryOptions = executeQueryOptions;
+  }
+
+  async streamResult() {
+    if (!this.response.body) {
+      throw new Error("Response body is not readable");
+    }
+    return {
+      data: new ServerSideStream(this.response, this.executeQueryOptions)
+    };
+  }
+}

src/statement/stream.ts

@@ -0,0 +1,65 @@
+import { Readable } from "stream";
+import JSONbig from "json-bigint";
+import readline from "readline";
+import {
+  getNormalizedMeta,
+  normalizeResponseRowStreaming
+} from "../normalizeResponse";
+import { Response } from "node-fetch";
+import { ExecuteQueryOptions } from "../../types";
+import { Meta } from "../../meta";
+
+export class ServerSideStream extends Readable {
+  private meta: Meta[] = [];
+  constructor(
+    private readonly response: Response,
+    private readonly executeQueryOptions: ExecuteQueryOptions
+  ) {
+    super({ objectMode: true });
+    const readLine = readline.createInterface({
+      input: response.body,
+      crlfDelay: Infinity
+    });
+
+    const lineParser = (line: string) => {
+      try {
+        if (line.trim()) {
+          const parsed = JSONbig.parse(line);
+          if (parsed) {
+            if (parsed.message_type === "DATA") {
+              if (parsed.data) {
+                const normalizedData = normalizeResponseRowStreaming(
+                  parsed.data,
+                  executeQueryOptions,
+                  this.meta
+                );
+                for (const data of normalizedData) {
+                  this.emit("data", data);
+                }
+              }
+            } else if (parsed.message_type === "START") {
+              const metaDataJson = JSONbig.parse(line);
+              this.meta = getNormalizedMeta(metaDataJson.result_columns);
+              this.emit("meta", this.meta);
+            } else if (parsed.message_type === "FINISH_SUCCESSFULLY") {
+              this.push(null);
+            } else if (parsed.message_type === "FINISH_WITH_ERROR") {
+              this.destroy(new Error(`Result encountered an error: ${line}`));
+            }
+          } else {
+            this.destroy(new Error(`Result row could not be parsed: ${line}`));
+          }
+        }
+      } catch (err) {
+        this.destroy(err);
+      }
+    };
+    readLine.on("line", lineParser);
+
+    readLine.on("close", () => {
+      this.push(null);
+    });
+  }
+
+  _read() {}
+}

src/statement/stream/serverSideStream.ts

@@ -26,7 +26,8 @@ export enum OutputFormat {
   // JSON_COMPACT_LIMITED = "FB_JSONCompactLimited",
   COMPACT = "JSON_Compact", // supported in v3 of packdb
   JSON_COMPACT = "JSONCompact", // to be as above, after ensure all clients has v3
-  JSON = "JSON"
+  JSON = "JSON",
+  JSON_LINES = "JSONLines_Compact"
 }
 
 export type QuerySettings = Record<

src/types.ts

@@ -1,53 +1,43 @@
-import stream, { TransformCallback } from "stream";
-import { Firebolt } from "../../../src/index";
+import { Firebolt } from "../../../src";
 
 const connectionParams = {
   auth: {
-    username: process.env.FIREBOLT_USERNAME as string,
-    password: process.env.FIREBOLT_PASSWORD as string
+    client_id: process.env.FIREBOLT_CLIENT_ID as string,
+    client_secret: process.env.FIREBOLT_CLIENT_SECRET as string
   },
+  account: process.env.FIREBOLT_ACCOUNT as string,
   database: process.env.FIREBOLT_DATABASE as string,
   engineName: process.env.FIREBOLT_ENGINE_NAME as string
 };
 
-jest.setTimeout(20000);
+jest.setTimeout(250000);
 
 describe("streams", () => {
-  it("stream transformters", async () => {
-    class SerializeRowStream extends stream.Transform {
-      public constructor() {
-        super({
-          objectMode: true,
-          transform(
-            row: any,
-            encoding: BufferEncoding,
-            callback: TransformCallback
-          ) {
-            const transformed = JSON.stringify(row);
-            this.push(transformed);
-            this.push("\n");
-            callback();
-          }
-        });
-      }
-    }
-
+  it("stream transformers", async () => {
     const firebolt = Firebolt({
       apiEndpoint: process.env.FIREBOLT_API_ENDPOINT as string
     });
 
-    const serializedStream = new SerializeRowStream();
-
     const connection = await firebolt.connect(connectionParams);
 
-    const statement = await connection.execute("select 1 union all select 2");
+    const statement = await connection.executeStream(
+      `select 1 from generate_series(1, 250000000)` //~1 GB response
+    );
 
     const { data } = await statement.streamResult();
-
-    data.pipe(serializedStream).pipe(process.stdout);
+    let sum = 0;
+
+    data
+      .on("meta", meta => {
+        console.log("Meta:", meta);
+      })
+      .on("data", row => {
+        sum += row[0];
+      });
 
     await new Promise(resolve => {
       data.on("end", () => {
+        expect(sum).toEqual(250000000);
         resolve(null);
       });
     });