feat: Fir 31841 dont append account id for system engine url in node sdk (#82)

Author: stepansergeevitch

src/connection/base.ts

@@ -115,11 +115,16 @@ export abstract class Connection {
     this.parameters = remainingParameters;
   }
 
-  private async handleUpdateEndpointHeader(headerValue: string): Promise<void> {
+  protected splitEndpoint(endpoint: string): [string, Record<string, string>] {
     const url = new URL(
-      headerValue.startsWith("http") ? headerValue : `https://${headerValue}`
+      endpoint.startsWith("http") ? endpoint : `https://${endpoint}`.trim()
     );
-    const newParams = Object.fromEntries(url.searchParams.entries());
+    const params = Object.fromEntries(url.searchParams.entries());
+    return [url.toString().replace(url.search, ""), params];
+  }
+
+  private async handleUpdateEndpointHeader(headerValue: string): Promise<void> {
+    const [endpoint, newParams] = this.splitEndpoint(headerValue);
 
     // Validate account_id if present
     const currentAccountId =
@@ -131,7 +136,7 @@ export abstract class Connection {
     }
 
     // Remove url parameters and update engineEndpoint
-    this.engineEndpoint = url.toString().replace(url.search, "");
+    this.engineEndpoint = endpoint;
     this.parameters = {
       ...this.parameters,
       ...newParams

src/connection/connection_v2.ts

@@ -22,16 +22,17 @@ export class ConnectionV2 extends BaseConnection {
     return this.options.account;
   }
 
-  private async getSystemEngineEndpoint(): Promise<string> {
+  private async getSystemEngineEndpointAndParameters(): Promise<
+    [string, Record<string, string>]
+  > {
     const { apiEndpoint, httpClient } = this.context;
     const accountName = this.account;
     const url = `${apiEndpoint}/${ACCOUNT_SYSTEM_ENGINE(accountName)}`;
     try {
       const { engineUrl } = await httpClient
         .request<{ engineUrl: string }>("GET", url)
         .ready();
-      // cut off query parameters that go after ?
-      return engineUrl.split("?")[0];
+      return this.splitEndpoint(engineUrl);
     } catch (e) {
       if (e instanceof ApiError && e.status == 404) {
         throw new AccountNotFoundError({ account_name: accountName });
@@ -142,8 +143,10 @@ export class ConnectionV2 extends BaseConnection {
   async resolveEngineEndpoint() {
     const { engineName, database } = this.options;
     // Connect to system engine first
-    const systemUrl = await this.getSystemEngineEndpoint();
-    this.engineEndpoint = path.join(systemUrl, QUERY_URL);
+    const [systemUrl, systemParameters] =
+      await this.getSystemEngineEndpointAndParameters();
+    this.engineEndpoint = new URL(QUERY_URL, systemUrl).href;
+    this.parameters = { ...this.parameters, ...systemParameters };
     this.accountInfo = await this.resolveAccountInfo();
 
     if (this.accountInfo.infraVersion >= 2) {
@@ -189,9 +192,12 @@ export class ConnectionV2 extends BaseConnection {
   protected getBaseParameters(
     executeQueryOptions: ExecuteQueryOptions
   ): Record<string, string | undefined> {
-    return {
-      account_id: this.accountInfo?.id,
-      ...super.getBaseParameters(executeQueryOptions)
-    };
+    if (this.accountInfo?.infraVersion == 1) {
+      return {
+        account_id: this.accountInfo?.id,
+        ...super.getBaseParameters(executeQueryOptions)
+      };
+    }
+    return super.getBaseParameters(executeQueryOptions);
   }
 }

test/unit/v2/connection.test.ts

@@ -210,4 +210,63 @@ describe("Connection V2", () => {
     const accountInfo = await (connection as ConnectionV2).resolveAccountInfo();
     expect(accountInfo.infraVersion).toBe(1);
   });
+  it("respects system engine query parameters for account version 2", async () => {
+    let systemEngineParamsUsed = {};
+    server.use(
+      rest.post(`https://id.fake.firebolt.io/oauth/token`, (req, res, ctx) => {
+        return res(
+          ctx.json({
+            access_token: "fake_access_token"
+          })
+        );
+      }),
+      rest.get(
+        `https://api.fake.firebolt.io/web/v3/account/my_account/resolve`,
+        (req, res, ctx) => {
+          return res(
+            ctx.json({
+              id: "123",
+              infraVersion: 2
+            })
+          );
+        }
+      ),
+      rest.get(
+        `https://api.fake.firebolt.io/web/v3/account/my_account/engineUrl`,
+        (req, res, ctx) => {
+          return res(
+            ctx.json({
+              engineUrl: "https://some_system_engine.com?param=value"
+            })
+          );
+        }
+      ),
+      rest.post(
+        `https://some_system_engine.com/${QUERY_URL}`,
+        (req, res, ctx) => {
+          systemEngineParamsUsed = Object.fromEntries(
+            req.url.searchParams.entries()
+          );
+          return res(ctx.json(engineUrlResponse));
+        }
+      )
+    );
+    const firebolt = Firebolt({
+      apiEndpoint
+    });
+
+    const connectionParams: ConnectionOptions = {
+      auth: {
+        client_id: "dummy",
+        client_secret: "dummy"
+      },
+      database: "dummy",
+      account: "my_account"
+    };
+
+    const connection = await firebolt.connect(connectionParams);
+    await connection.execute("SELECT 1");
+    expect(systemEngineParamsUsed).toHaveProperty("param");
+    expect(systemEngineParamsUsed).not.toHaveProperty("account_id");
+  });
 });