feat(FIR-42859): async query execution

Author: ptiurin

src/connection/base.ts

@@ -9,12 +9,13 @@ import { generateUserAgent } from "../common/util";
 import { CompositeError } from "../common/errors";
 import JSONbig from "json-bigint";
 import { QueryFormatter } from "../formatter/base";
+import { AsyncStatement } from "../statement/async";
 
 const defaultQuerySettings = {
   output_format: OutputFormat.COMPACT
 };
 
-const defaultResponseSettings = {
+export const defaultResponseSettings = {
   normalizeData: false
 };
 
@@ -160,6 +161,17 @@ export abstract class Connection {
     }
   }
 
+  abstract executeAsync(
+    query: string,
+    executeQueryOptions?: ExecuteQueryOptions
+  ): Promise<AsyncStatement>;
+
+  abstract isAsyncQueryRunning(token: string): Promise<boolean>;
+
+  abstract isAsyncQuerySuccessful(token: string): Promise<boolean | undefined>;
+
+  abstract cancelAsyncQuery(token: string): Promise<void>;
+
   async execute(
     query: string,
     executeQueryOptions: ExecuteQueryOptions = {}
@@ -220,7 +232,7 @@ export abstract class Connection {
     }
   }
 
-  private async throwErrorIfErrorBody(text: string, response: Response) {
+  protected async throwErrorIfErrorBody(text: string, response: Response) {
     // Hack, but looks like this is a limitation of the fetch API
     // In order to read the body here and elesewhere, we need to clone the response
     // since body can only be read once

src/connection/connection_v1.ts

@@ -1,6 +1,8 @@
 import { ACCOUNT, ACCOUNT_BY_NAME } from "../common/api";
 import { Connection as BaseConnection } from "./base";
 import { ResourceManager } from "../service";
+import { ExecuteQueryOptions } from "../types";
+import { AsyncStatement } from "../statement/async";
 
 export interface AccountInfo {
   id: string;
@@ -63,4 +65,32 @@ export class ConnectionV1 extends BaseConnection {
   async testConnection() {
     await this.execute("select 1");
   }
+
+  // Async methods
+  async isAsyncQueryRunning(token: string): Promise<boolean> {
+    throw new Error(
+      "Asynchronous query running check is not supported in this Firebolt version."
+    );
+  }
+
+  async isAsyncQuerySuccessful(token: string): Promise<boolean | undefined> {
+    throw new Error(
+      "Asynchronous query success check is not supported in this Firebolt version."
+    );
+  }
+
+  async cancelAsyncQuery(token: string): Promise<void> {
+    throw new Error(
+      "Asynchronous query cancellation is not supported in this Firebolt version."
+    );
+  }
+
+  async executeAsync(
+    query: string,
+    executeQueryOptions?: ExecuteQueryOptions
+  ): Promise<AsyncStatement> {
+    throw new Error(
+      "Asynchronous execution is not supported in this Firebolt version."
+    );
+  }
 }

src/connection/connection_v2.ts

@@ -1,8 +1,10 @@
 import { AccountNotFoundError, ApiError } from "../common/errors";
 import { ACCOUNT_SYSTEM_ENGINE, QUERY_URL } from "../common/api";
 
-import { Connection as BaseConnection } from "./base";
+import { Connection as BaseConnection, defaultResponseSettings } from "./base";
 import { Cache, inMemoryCache, noneCache } from "../common/tokenCache";
+import { ExecuteQueryOptions } from "../types";
+import { AsyncStatement } from "../statement/async";
 
 export class ConnectionV2 extends BaseConnection {
   private get account(): string {
@@ -70,6 +72,93 @@ export class ConnectionV2 extends BaseConnection {
     return this.engineEndpoint;
   }
 
+  // Async methods
+  async executeAsync(
+    query: string,
+    executeQueryOptions: ExecuteQueryOptions = {}
+  ): Promise<AsyncStatement> {
+    const { httpClient } = this.context;
+
+    executeQueryOptions.response = {
+      ...defaultResponseSettings,
+      ...(executeQueryOptions.response ?? {})
+    };
+
+    const { parameters, namedParameters } = executeQueryOptions;
+
+    let formattedQuery: string;
+    if (this.queryFormatter.isSetStatement(query)) {
+      // can't have an async set query
+      throw new Error("SET statements cannot be executed asynchronously.");
+    } else {
+      formattedQuery = this.queryFormatter.formatQuery(
+        query,
+        parameters,
+        namedParameters
+      );
+    }
+
+    const body = formattedQuery;
+    const asyncExecuteQueryOptions = {
+      ...executeQueryOptions,
+      settings: {
+        ...executeQueryOptions.settings,
+        async: true
+      }
+    };
+    const url = this.getRequestUrl(asyncExecuteQueryOptions);
+
+    const request = httpClient.request<Response>("POST", url, {
+      headers: { "user-agent": this.userAgent },
+      body,
+      raw: true
+    });
+
+    const response = await request.ready();
+    const text = await response.text();
+    await this.throwErrorIfErrorBody(text, response);
+    return new AsyncStatement(this.context, {
+      query: formattedQuery,
+      text,
+      executeQueryOptions: asyncExecuteQueryOptions
+    });
+  }
+
+  private async getAsyncQueryInfo(token: string) {
+    const query = `CALL fb_GetAsyncStatus('${token}')`;
+
+    const statement = await this.execute(query);
+    const { data, meta } = await statement.fetchResult();
+    const result: Record<string, any> = {};
+    if (data.length > 0) {
+      meta.forEach((field, index) => {
+        result[field.name] = data[0][index];
+      });
+    } else {
+      throw new Error("No data returned from fb_GetAsyncStatus");
+    }
+    return result;
+  }
+
+  async isAsyncQueryRunning(token: string): Promise<boolean> {
+    const info = await this.getAsyncQueryInfo(token);
+    return info["status"] === "RUNNING";
+  }
+
+  async isAsyncQuerySuccessful(token: string): Promise<boolean | undefined> {
+    const info = await this.getAsyncQueryInfo(token);
+    if (info["status"] === "RUNNING") {
+      return undefined;
+    }
+    return info["status"] === "ENDED_SUCCESSFULLY";
+  }
+
+  async cancelAsyncQuery(token: string): Promise<void> {
+    const info = await this.getAsyncQueryInfo(token);
+    const async_query_id = info["query_id"];
+    this.execute(`CANCEL QUERY WHERE query_id='${async_query_id}'`);
+  }
+
   async testConnection() {
     const settings = { internal: [{ auto_start_stop_control: "ignore" }] };
     await this.execute("select 1", { settings });

src/http/node.ts

@@ -26,7 +26,7 @@ const DEFAULT_ERROR = "Server error";
 const DEFAULT_USER_AGENT = systemInfoString();
 
 const PROTOCOL_VERSION_HEADER = "Firebolt-Protocol-Version";
-const PROTOCOL_VERSION = "2.1";
+const PROTOCOL_VERSION = "2.3";
 const createSocket = HttpsAgent.prototype.createSocket;
 
 const agentOptions = {

src/statement/async.ts

@@ -0,0 +1,54 @@
+import JSONbig from "json-bigint";
+import {
+  ExecuteQueryOptions,
+  StreamOptions,
+  Context,
+  Statistics
+} from "../types";
+import { Meta } from "../meta";
+import { BaseStatement } from "./baseStatement";
+
+export class AsyncStatement extends BaseStatement {
+  private asyncToken: string;
+
+  constructor(
+    context: Context,
+    {
+      query,
+      text,
+      executeQueryOptions
+    }: {
+      query: string;
+      text: string;
+      executeQueryOptions: ExecuteQueryOptions;
+    }
+  ) {
+    super(context, { query, text, executeQueryOptions });
+    this.asyncToken = this.parseResponse(this.text).token;
+  }
+
+  private parseResponse(response: string) {
+    const parsed = JSONbig.parse(response);
+    const { token, message, monitorSql } = parsed;
+    return {
+      token,
+      message,
+      monitorSql
+    };
+  }
+  get async_query_token(): string {
+    return this.asyncToken;
+  }
+
+  async streamResult(options?: StreamOptions) {
+    throw new Error("Method not allowed for async statements");
+  }
+
+  async fetchResult(): Promise<{
+    data: any;
+    meta: Meta[];
+    statistics: Statistics | null;
+  }> {
+    throw new Error("Method not allowed for async statements");
+  }
+}

src/statement/baseStatement.ts

@@ -0,0 +1,39 @@
+import {
+  ExecuteQueryOptions,
+  Context,
+  StreamOptions,
+  Statistics
+} from "../types";
+import { Meta } from "../meta";
+
+export abstract class BaseStatement {
+  protected context: Context;
+  protected query: string;
+  protected executeQueryOptions: ExecuteQueryOptions;
+  protected text: string;
+
+  constructor(
+    context: Context,
+    {
+      query,
+      text,
+      executeQueryOptions
+    }: {
+      query: string;
+      text: string;
+      executeQueryOptions: ExecuteQueryOptions;
+    }
+  ) {
+    this.context = context;
+    this.text = text;
+    this.query = query;
+    this.executeQueryOptions = executeQueryOptions;
+  }
+
+  abstract streamResult(options?: StreamOptions): Promise<any>;
+  abstract fetchResult(): Promise<{
+    data: any;
+    meta: Meta[];
+    statistics: Statistics | null;
+  }>;
+}

src/statement/index.ts

@@ -8,19 +8,11 @@ import {
 import { Meta } from "../meta";
 import { isDataQuery } from "../common/util";
 import { RowStream } from "./stream/rowStream";
-import { JSONStream } from "./stream/jsonStream";
-import {
-  normalizeResponse,
-  getNormalizedStatistics
-} from "./normalizeResponse";
+import { normalizeResponse } from "./normalizeResponse";
 import { CompositeError } from "../common/errors";
+import { BaseStatement } from "./baseStatement";
 
-export class Statement {
-  private context: Context;
-  private query: string;
-  private executeQueryOptions: ExecuteQueryOptions;
-
-  private text;
+export class Statement extends BaseStatement {
   private rowStream: RowStream;
 
   constructor(
@@ -35,10 +27,7 @@ export class Statement {
       executeQueryOptions: ExecuteQueryOptions;
     }
   ) {
-    this.context = context;
-    this.text = text;
-    this.query = query;
-    this.executeQueryOptions = executeQueryOptions;
+    super(context, { query, text, executeQueryOptions });
     this.rowStream = new RowStream();
   }
 

src/types.ts

@@ -35,6 +35,7 @@ export type QuerySettings = Record<
 > & {
   output_format?: OutputFormat;
   internal?: Record<string, string | number>[];
+  async?: boolean;
 };
 
 export type RowParser = (row: string, isLastRow: boolean) => any;