feat: Fir 30153 support use engine in node sdk (#78)

Author: stepansergeevitch

.github/workflows/integration-tests-v2.yaml

@@ -28,13 +28,14 @@ jobs:
         with:
           firebolt-client-id: ${{ secrets.FIREBOLT_CLIENT_ID_STG_NEW_IDN }}
           firebolt-client-secret: ${{ secrets.FIREBOLT_CLIENT_SECRET_STG_NEW_IDN }}
-          account: ${{ vars.FIREBOLT_ACCOUNT }}
+          account: ${{ vars.FIREBOLT_ACCOUNT_V1 }}
           api-endpoint: "api.staging.firebolt.io"
           region: "us-east-1"
 
       - name: Run integration tests
         env:
-          FIREBOLT_ACCOUNT: ${{ vars.FIREBOLT_ACCOUNT }}
+          FIREBOLT_ACCOUNT_V1: ${{ vars.FIREBOLT_ACCOUNT_V1 }}
+          FIREBOLT_ACCOUNT_V2: ${{ vars.FIREBOLT_ACCOUNT_V2 }}
           FIREBOLT_DATABASE: ${{ steps.setup.outputs.database_name }}
           FIREBOLT_ENGINE_NAME: ${{ steps.setup.outputs.engine_name }}
           FIREBOLT_API_ENDPOINT: "api.staging.firebolt.io"

src/connection/base.ts

@@ -6,6 +6,7 @@ import {
 } from "../types";
 import { Statement } from "../statement";
 import { generateUserAgent } from "../common/util";
+import { ConnectionError } from "../common/errors";
 
 const defaultQuerySettings = {
   output_format: OutputFormat.COMPACT
@@ -15,19 +16,33 @@ const defaultResponseSettings = {
   normalizeData: false
 };
 
+interface AccountInfo {
+  id: string;
+  infraVersion: number;
+}
+
 const updateParametersHeader = "Firebolt-Update-Parameters";
 const allowedUpdateParameters = ["database"];
+const updateEndpointHeader = "Firebolt-Update-Endpoint";
+const resetSessionHeader = "Firebolt-Reset-Session";
+const immutableParameters = ["database", "account_id", "output_format"];
 
 export abstract class Connection {
   protected context: Context;
   protected options: ConnectionOptions;
   protected userAgent: string;
+  protected parameters: Record<string, string>;
+  protected accountInfo: AccountInfo | undefined;
   engineEndpoint!: string;
   activeRequests = new Set<{ abort: () => void }>();
 
   constructor(context: Context, options: ConnectionOptions) {
     this.context = context;
     this.options = options;
+    this.parameters = {
+      ...(options.database ? { database: options.database } : {}),
+      ...defaultQuerySettings
+    };
     this.userAgent = generateUserAgent(
       options.additionalParameters?.userClients,
       options.additionalParameters?.userDrivers
@@ -60,31 +75,82 @@ export abstract class Connection {
     executeQueryOptions: ExecuteQueryOptions
   ): Record<string, string | undefined> {
     const { settings } = executeQueryOptions;
-    const { database } = this.options;
-    return { database, ...settings };
+
+    // convert all settings values to string
+    const strSettings = Object.entries(settings ?? {}).reduce<
+      Record<string, string>
+    >((acc, [key, value]) => {
+      if (value !== undefined) {
+        acc[key] = value.toString();
+      }
+      return acc;
+    }, {});
+
+    return { ...this.parameters, ...strSettings };
+  }
+
+  private handleUpdateParametersHeader(headerValue: string) {
+    const updateParameters = headerValue
+      .split(",")
+      .reduce((acc: Record<string, string>, param) => {
+        const [key, value] = param.split("=");
+        if (allowedUpdateParameters.includes(key)) {
+          acc[key] = value.trim();
+        }
+        return acc;
+      }, {});
+    this.parameters = {
+      ...this.parameters,
+      ...updateParameters
+    };
+  }
+
+  private handleResetSessionHeader() {
+    const remainingParameters: Record<string, string> = {};
+    for (const key in this.parameters) {
+      if (immutableParameters.includes(key)) {
+        remainingParameters[key] = this.parameters[key];
+      }
+    }
+    this.parameters = remainingParameters;
   }
 
-  private processHeaders(headers: Headers) {
+  private async handleUpdateEndpointHeader(headerValue: string): Promise<void> {
+    const url = new URL(
+      headerValue.startsWith("http") ? headerValue : `https://${headerValue}`
+    );
+    const newParams = Object.fromEntries(url.searchParams.entries());
+
+    // Validate account_id if present
+    const currentAccountId =
+      this.accountInfo?.id ?? (await this.resolveAccountId());
+    if (newParams.account_id && currentAccountId !== newParams.account_id) {
+      throw new ConnectionError({
+        message: `Failed to execute USE ENGINE command. Account parameter mismatch. Contact support.`
+      });
+    }
+
+    // Remove url parameters and update engineEndpoint
+    this.engineEndpoint = url.toString().replace(url.search, "");
+    this.parameters = {
+      ...this.parameters,
+      ...newParams
+    };
+  }
+
+  private async processHeaders(headers: Headers) {
     const updateHeaderValue = headers.get(updateParametersHeader);
     if (updateHeaderValue) {
-      const updateParameters = updateHeaderValue
-        .split(",")
-        .reduce((acc: Record<string, string>, param) => {
-          const [key, value] = param.split("=");
-          if (allowedUpdateParameters.includes(key)) {
-            acc[key] = value.trim();
-          }
-          return acc;
-        }, {});
-
-      if (updateParameters.database) {
-        this.options.database = updateParameters.database;
-        delete updateParameters.database;
-      }
-      this.options.additionalParameters = {
-        ...this.options.additionalParameters,
-        ...updateParameters
-      };
+      this.handleUpdateParametersHeader(updateHeaderValue);
+    }
+
+    if (headers.has(resetSessionHeader)) {
+      this.handleResetSessionHeader();
+    }
+
+    const updateEndpointValue = headers.get(updateEndpointHeader);
+    if (updateEndpointValue) {
+      await this.handleUpdateEndpointHeader(updateEndpointValue);
     }
   }
 
@@ -94,11 +160,6 @@ export abstract class Connection {
   ): Promise<Statement> {
     const { httpClient, queryFormatter } = this.context;
 
-    executeQueryOptions.settings = {
-      ...defaultQuerySettings,
-      ...(executeQueryOptions.settings ?? {})
-    };
-
     executeQueryOptions.response = {
       ...defaultResponseSettings,
       ...(executeQueryOptions.response ?? {})
@@ -124,7 +185,7 @@ export abstract class Connection {
 
     try {
       const response = await request.ready();
-      this.processHeaders(response.headers);
+      await this.processHeaders(response.headers);
       const statement = new Statement(this.context, {
         query: formattedQuery,
         request,

src/connection/connection_v2.ts

@@ -12,10 +12,9 @@ import {
 } from "../common/api";
 
 import { Connection as BaseConnection } from "./base";
+import * as path from "path";
 
 export class ConnectionV2 extends BaseConnection {
-  private accountId: string | undefined;
-
   private get account(): string {
     if (!this.options.account) {
       throw new Error("Account name is required");
@@ -31,7 +30,8 @@ export class ConnectionV2 extends BaseConnection {
       const { engineUrl } = await httpClient
         .request<{ engineUrl: string }>("GET", url)
         .ready();
-      return engineUrl;
+      // cut off query parameters that go after ?
+      return engineUrl.split("?")[0];
     } catch (e) {
       if (e instanceof ApiError && e.status == 404) {
         throw new AccountNotFoundError({ account_name: accountName });
@@ -125,57 +125,72 @@ export class ConnectionV2 extends BaseConnection {
     return res[0];
   }
 
-  async resolveAccountId() {
+  async resolveAccountInfo() {
     const { httpClient, apiEndpoint } = this.context;
     const url = `${apiEndpoint}/${ACCOUNT_ID_BY_NAME(this.account)}`;
-    const { id } = await httpClient
-      .request<{ id: string; region: string }>("GET", url)
+    const { id, infraVersion } = await httpClient
+      .request<{ id: string; region: string; infraVersion: string }>("GET", url)
       .ready();
-    return id;
+    return { id, infraVersion: parseInt(infraVersion ?? "1") };
+  }
+
+  async resolveAccountId() {
+    const accInfo = await this.resolveAccountInfo();
+    return accInfo.id;
   }
 
   async resolveEngineEndpoint() {
     const { engineName, database } = this.options;
     // Connect to system engine first
     const systemUrl = await this.getSystemEngineEndpoint();
-    this.engineEndpoint = `${systemUrl}/${QUERY_URL}`;
-    this.accountId = await this.resolveAccountId();
-    if (engineName && database) {
-      const engineEndpoint = await this.getEngineByNameAndDb(
-        engineName,
-        database
-      );
-      this.engineEndpoint = engineEndpoint;
-      // Account id is no longer needed
-      this.accountId = undefined;
-      return this.engineEndpoint;
-    }
-    if (engineName) {
-      const database = await this.getEngineDatabase(engineName);
-      if (!database) {
-        throw new AccessError({
-          message: `Engine ${engineName} is attached to a database that current user can not access.`
-        });
+    this.engineEndpoint = path.join(systemUrl, QUERY_URL);
+    this.accountInfo = await this.resolveAccountInfo();
+
+    if (this.accountInfo.infraVersion >= 2) {
+      if (database) {
+        await this.execute(`USE DATABASE ${database}`);
       }
-      const engineEndpoint = await this.getEngineByNameAndDb(
-        engineName,
-        database
-      );
-      this.options.database = database;
-      this.engineEndpoint = engineEndpoint;
-      // Account id is no longer needed
-      this.accountId = undefined;
-      return this.engineEndpoint;
+      if (engineName) {
+        await this.execute(`USE ENGINE ${engineName}`);
+      }
+    } else {
+      if (engineName && database) {
+        const engineEndpoint = await this.getEngineByNameAndDb(
+          engineName,
+          database
+        );
+        this.engineEndpoint = engineEndpoint;
+        // Account id is no longer needed
+        this.accountInfo = undefined;
+        return this.engineEndpoint;
+      }
+      if (engineName) {
+        const database = await this.getEngineDatabase(engineName);
+        if (!database) {
+          throw new AccessError({
+            message: `Engine ${engineName} is attached to a database that current user can not access.`
+          });
+        }
+        const engineEndpoint = await this.getEngineByNameAndDb(
+          engineName,
+          database
+        );
+        this.parameters["database"] = database;
+        this.engineEndpoint = engineEndpoint;
+        // Account id is no longer needed
+        this.accountInfo = undefined;
+        return this.engineEndpoint;
+      }
+      // If nothing specified connect to generic system engine
     }
-    // If nothing specified connect to generic system engine
     return this.engineEndpoint;
   }
 
   protected getBaseParameters(
     executeQueryOptions: ExecuteQueryOptions
   ): Record<string, string | undefined> {
     return {
-      account_id: this.accountId,
+      account_id: this.accountInfo?.id,
       ...super.getBaseParameters(executeQueryOptions)
     };
   }

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.0";
+const PROTOCOL_VERSION = "2.1";
 const createSocket = HttpsAgent.prototype.createSocket;
 
 const agentOptions = {

src/types.ts

@@ -29,7 +29,10 @@ export enum OutputFormat {
   JSON = "JSON"
 }
 
-export type QuerySettings = Record<string, unknown> & {
+export type QuerySettings = Record<
+  string,
+  string | number | boolean | undefined
+> & {
   output_format?: OutputFormat;
 };
 

test/integration/v2/account.test.ts

@@ -7,7 +7,7 @@ const connectionOptions = {
   },
   database: process.env.FIREBOLT_DATABASE as string,
   engineName: process.env.FIREBOLT_ENGINE_NAME as string,
-  account: process.env.FIREBOLT_ACCOUNT as string
+  account: process.env.FIREBOLT_ACCOUNT_V1 as string
 };
 
 jest.setTimeout(20000);

test/integration/v2/auth.test.ts

@@ -7,7 +7,7 @@ const auth = {
 
 const connectionOptions = {
   database: process.env.FIREBOLT_DATABASE as string,
-  account: process.env.FIREBOLT_ACCOUNT as string
+  account: process.env.FIREBOLT_ACCOUNT_V1 as string
 };
 
 jest.setTimeout(20000);

test/integration/v2/boolean.test.ts

@@ -5,7 +5,7 @@ const connectionParams = {
     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,
+  account: process.env.FIREBOLT_ACCOUNT_V1 as string,
   database: process.env.FIREBOLT_DATABASE as string,
   engineName: process.env.FIREBOLT_ENGINE_NAME as string
 };

test/integration/v2/bytea.test.ts

@@ -5,7 +5,7 @@ const connectionParams = {
     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,
+  account: process.env.FIREBOLT_ACCOUNT_V1 as string,
   database: process.env.FIREBOLT_DATABASE as string,
   engineName: process.env.FIREBOLT_ENGINE_NAME as string
 };

test/integration/v2/database.test.ts

@@ -5,7 +5,7 @@ const connectionOptions = {
     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,
+  account: process.env.FIREBOLT_ACCOUNT_V1 as string,
   database: process.env.FIREBOLT_DATABASE as string
 };