refactor(FIR-35470): remove redundant account id resolution (#116)

ptiurin · web-flow · commit 0f4446208a5e · 2024-09-02T15:03:36.000+01:00
diff --git a/src/common/tokenCache.ts b/src/common/tokenCache.ts
@@ -1,4 +1,4 @@
-import { AccountInfo } from "../connection/base";
+import { AccountInfo } from "../connection/connection_v1";
 
 export type TokenKey = {
   clientId: string;
diff --git a/src/connection/base.ts b/src/connection/base.ts
@@ -6,7 +6,7 @@ import {
 } from "../types";
 import { Statement } from "../statement";
 import { generateUserAgent } from "../common/util";
-import { ConnectionError, CompositeError } from "../common/errors";
+import { CompositeError } from "../common/errors";
 import JSONbig from "json-bigint";
 
 const defaultQuerySettings = {
@@ -17,11 +17,6 @@ const defaultResponseSettings = {
   normalizeData: false
 };
 
-export interface AccountInfo {
-  id: string;
-  infraVersion: number;
-}
-
 const updateParametersHeader = "Firebolt-Update-Parameters";
 const allowedUpdateParameters = ["database"];
 const updateEndpointHeader = "Firebolt-Update-Endpoint";
@@ -34,7 +29,6 @@ export abstract class Connection {
   protected options: ConnectionOptions;
   protected userAgent: string;
   protected parameters: Record<string, string>;
-  protected accountInfo: AccountInfo | undefined;
   engineEndpoint!: string;
   activeRequests = new Set<{ abort: () => void }>();
 
@@ -53,13 +47,6 @@ export abstract class Connection {
 
   abstract resolveEngineEndpoint(): Promise<string>;
 
-  abstract resolveAccountInfo(): Promise<AccountInfo>;
-
-  async resolveAccountId() {
-    const accInfo = await this.resolveAccountInfo();
-    return accInfo.id;
-  }
-
   protected getRequestUrl(executeQueryOptions: ExecuteQueryOptions): string {
     const params = this.getBaseParameters(executeQueryOptions);
 
@@ -132,16 +119,6 @@ export abstract class Connection {
 
   private async handleUpdateEndpointHeader(headerValue: string): Promise<void> {
     const [endpoint, newParams] = this.splitEndpoint(headerValue);
-
-    // 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 = endpoint;
     this.parameters = {
diff --git a/src/connection/connection_v1.ts b/src/connection/connection_v1.ts
@@ -1,10 +1,14 @@
 import { ACCOUNT, ACCOUNT_BY_NAME } from "../common/api";
-import { AccountInfo, Connection as BaseConnection } from "./base";
+import { Connection as BaseConnection } from "./base";
 import { ResourceManager } from "../service";
 
-const INFRA_VERSION = 1;
+export interface AccountInfo {
+  id: string;
+}
 
 export class ConnectionV1 extends BaseConnection {
+  protected accountInfo: AccountInfo | undefined;
+
   async resolveEngineEndpoint() {
     const resourceManager = new ResourceManager({
       connection: this,
@@ -27,6 +31,11 @@ export class ConnectionV1 extends BaseConnection {
     return this.engineEndpoint;
   }
 
+  async resolveAccountId() {
+    const accInfo = await this.resolveAccountInfo();
+    return accInfo.id;
+  }
+
   async resolveAccountInfo(): Promise<AccountInfo> {
     if (this.accountInfo === undefined) {
       const { httpClient, apiEndpoint } = this.context;
@@ -37,15 +46,15 @@ export class ConnectionV1 extends BaseConnection {
         const { account_id } = await httpClient
           .request<{ account_id: string }>("GET", url)
           .ready();
-        this.accountInfo = { id: account_id, infraVersion: INFRA_VERSION };
+        this.accountInfo = { id: account_id };
       } else {
         const url = `${apiEndpoint}/${ACCOUNT}`;
         const {
           account: { id }
         } = await httpClient
           .request<{ account: { id: string } }>("GET", url)
           .ready();
-        this.accountInfo = { id, infraVersion: INFRA_VERSION };
+        this.accountInfo = { id };
       }
     }
     return this.accountInfo;
diff --git a/src/connection/connection_v2.ts b/src/connection/connection_v2.ts
@@ -1,11 +1,7 @@
 import { AccountNotFoundError, ApiError } from "../common/errors";
-import {
-  ACCOUNT_ID_BY_NAME,
-  ACCOUNT_SYSTEM_ENGINE,
-  QUERY_URL
-} from "../common/api";
+import { ACCOUNT_SYSTEM_ENGINE, QUERY_URL } from "../common/api";
 
-import { Connection as BaseConnection, AccountInfo } from "./base";
+import { Connection as BaseConnection } from "./base";
 import { Cache, inMemoryCache, noneCache } from "../common/tokenCache";
 
 export class ConnectionV2 extends BaseConnection {
@@ -56,44 +52,13 @@ export class ConnectionV2 extends BaseConnection {
     }
   }
 
-  async resolveAccountInfo(): Promise<AccountInfo> {
-    if (this.accountInfo === undefined) {
-      const cachedValue = this.cache.accountInfoStorage.get({
-        account: this.account,
-        apiEndpoint: this.context.apiEndpoint
-      });
-      if (cachedValue) {
-        this.accountInfo = cachedValue;
-      } else {
-        const { httpClient, apiEndpoint } = this.context;
-        const url = `${apiEndpoint}/${ACCOUNT_ID_BY_NAME(this.account)}`;
-        const { id, infraVersion } = await httpClient
-          .request<{ id: string; region: string; infraVersion: string }>(
-            "GET",
-            url
-          )
-          .ready();
-        this.accountInfo = { id, infraVersion: parseInt(infraVersion ?? "2") };
-        this.cache.accountInfoStorage.set(
-          {
-            account: this.account,
-            apiEndpoint: this.context.apiEndpoint
-          },
-          this.accountInfo
-        );
-      }
-    }
-    return this.accountInfo;
-  }
-
   async resolveEngineEndpoint() {
     const { engineName, database } = this.options;
     // Connect to system engine first
     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 (database) {
       await this.execute(`USE DATABASE "${database}"`);
diff --git a/src/service/database/v1/index.ts b/src/service/database/v1/index.ts
@@ -4,17 +4,16 @@ import {
   ACCOUNT_ENGINE_URL_BY_DATABASE_NAME,
   ResultsPage
 } from "../../../common/api";
-import { ResourceManagerContext } from "../../../types";
+import { ResourceManagerContextV1 } from "../../../types";
 import { DatabaseModel } from "./model";
 import { ID, Database } from "./types";
 import { CreateDatabaseOptions } from "../types";
 import { resolveRegionKey } from "../../utils";
 import { ResourceManager } from "../../index";
-
 export class DatabaseService {
-  private readonly context: ResourceManagerContext;
+  private readonly context: ResourceManagerContextV1;
 
-  constructor(context: ResourceManagerContext) {
+  constructor(context: ResourceManagerContextV1) {
     this.context = context;
   }
 
diff --git a/src/service/database/v1/model.ts b/src/service/database/v1/model.ts
@@ -2,17 +2,17 @@ import {
   ACCOUNT_ENGINE_URL_BY_DATABASE_NAME,
   ACCOUNT_DATABASE
 } from "../../../common/api";
-import { ResourceManagerContext } from "../../../types";
+import { ResourceManagerContextV1 } from "../../../types";
 import { ID, Database } from "./types";
 import { ResourceManager } from "../../index";
 
 export class DatabaseModel {
-  private readonly context: ResourceManagerContext;
+  private readonly context: ResourceManagerContextV1;
   id: ID;
   name: string;
   description: string;
 
-  constructor(context: ResourceManagerContext, database: Database) {
+  constructor(context: ResourceManagerContextV1, database: Database) {
     const { id, name, description } = database;
     this.id = id;
     this.name = name;
diff --git a/src/service/database/v1/types.ts b/src/service/database/v1/types.ts
@@ -4,4 +4,4 @@ export type Database = {
   id: ID;
   name: string;
   description: string;
-};
+};
diff --git a/src/service/engine/v1/index.ts b/src/service/engine/v1/index.ts
@@ -6,7 +6,7 @@ import {
   ENGINES_BY_IDS_URL,
   ResultsPage
 } from "../../../common/api";
-import { ResourceManagerContext } from "../../../types";
+import { ResourceManagerContextV1 } from "../../../types";
 import { EngineModel } from "./model";
 import { Engine, ID } from "./types";
 import { CreateEngineOptions } from "../types";
@@ -19,9 +19,9 @@ import {
 import { DatabaseService } from "../../database/v1";
 
 export class EngineService {
-  private readonly context: ResourceManagerContext;
+  private readonly context: ResourceManagerContextV1;
 
-  constructor(context: ResourceManagerContext) {
+  constructor(context: ResourceManagerContextV1) {
     this.context = context;
   }
 
diff --git a/src/service/engine/v1/model.ts b/src/service/engine/v1/model.ts
@@ -4,18 +4,18 @@ import {
   ACCOUNT_ENGINE_STOP,
   ACCOUNT_ENGINE_RESTART
 } from "../../../common/api";
-import { ResourceManagerContext } from "../../../types";
+import { ResourceManagerContextV1 } from "../../../types";
 import { ID, Engine, EngineStatusSummary } from "./types";
 
 export class EngineModel {
-  private readonly context: ResourceManagerContext;
+  private readonly context: ResourceManagerContextV1;
   id: ID;
   name: string;
   description: string;
   endpoint: string;
   current_status_summary: EngineStatusSummary;
 
-  constructor(context: ResourceManagerContext, engine: Engine) {
+  constructor(context: ResourceManagerContextV1, engine: Engine) {
     const { id, name, description, endpoint, current_status_summary } = engine;
     this.id = id;
     this.name = name;
diff --git a/src/service/index.ts b/src/service/index.ts
@@ -4,7 +4,7 @@ import { DatabaseService as DatabaseServiceV1 } from "./database/v1";
 import { EngineServiceInterface } from "./engine/types";
 import { EngineService as EngineServiceV2 } from "./engine";
 import { EngineService as EngineServiceV1 } from "./engine/v1";
-import { ResourceManagerContext } from "../types";
+import { ResourceManagerContext, ResourceManagerContextV1 } from "../types";
 
 export class ResourceManager {
   private context: ResourceManagerContext;
@@ -18,8 +18,12 @@ export class ResourceManager {
       this.engine = new EngineServiceV2(this.context);
       this.database = new DatabaseServiceV2(this.context);
     } else if (httpClient.authenticator.isUsernamePassword()) {
-      this.engine = new EngineServiceV1(this.context);
-      this.database = new DatabaseServiceV1(this.context);
+      this.engine = new EngineServiceV1(
+        this.context as ResourceManagerContextV1
+      );
+      this.database = new DatabaseServiceV1(
+        this.context as ResourceManagerContextV1
+      );
     } else {
       throw new Error(
         "Invalid auth credentials provided. Please check your credentials and try again."
diff --git a/src/types.ts b/src/types.ts
@@ -5,6 +5,7 @@ import { LoggerInterface, LoggerOptions } from "./logger";
 import { Meta } from "./meta";
 import { QueryFormatter } from "./formatter";
 import { Connection } from "./connection";
+import { ConnectionV1 } from "./connection/connection_v1";
 
 export type Statistics = {
   duration: number | BigNumber;
@@ -118,3 +119,7 @@ export type ResourceManagerContext = {
   connection: Connection;
   apiEndpoint: string;
 };
+
+export type ResourceManagerContextV1 = ResourceManagerContext & {
+  connection: ConnectionV1;
+};
diff --git a/test/integration/v2/engine.test.ts b/test/integration/v2/engine.test.ts
@@ -123,7 +123,6 @@ describe.each([
       );
       expect(engine.name == name);
 
-      const accountId = (await connection.resolveAccountInfo()).infraVersion;
       await firebolt.resourceManager.engine.attachToDatabase(engine, database);
       const attached_engines = await database.getAttachedEngines();
       expect(attached_engines.includes(engine));
diff --git a/test/integration/v2/systemEngine.test.ts b/test/integration/v2/systemEngine.test.ts
@@ -41,20 +41,9 @@ describe("system engine", () => {
         `create database if not exists "${databaseName}"`
       );
 
-      const acc_version = (await connection.resolveAccountInfo()).infraVersion;
-
-      if (acc_version === 1) {
-        await connection.execute(
-          `create engine if not exists "${engineName}" with SPEC = 'B1' SCALE = 1`
-        );
-        await connection.execute(
-          `attach engine "${engineName}" to "${databaseName}"`
-        );
-      } else {
-        await connection.execute(
-          `create engine if not exists "${engineName}" with TYPE=S NODES=1`
-        );
-      }
+      await connection.execute(
+        `create engine if not exists "${engineName}" with TYPE=S NODES=1`
+      );
     } catch (error) {
       console.log(error);
       expect(true).toEqual(false);
diff --git a/test/unit/v1/connection.test.ts b/test/unit/v1/connection.test.ts
@@ -1,6 +1,8 @@
 import { setupServer, SetupServerApi } from "msw/node";
 import { rest } from "msw";
 import { ConnectionOptions, Firebolt } from "../../../src";
+import { ConnectionV1 } from "../../../src/connection/connection_v1";
+
 
 const apiEndpoint = "fake.api.com";
 const accountName = "my_account";
@@ -126,11 +128,13 @@ describe("Connection v1", () => {
     };
 
     const connection = await firebolt.connect(connectionParams);
+    expect(connection).toBeInstanceOf(ConnectionV1);
     await connection.execute("SELECT 1");
     expect(connection.engineEndpoint).toBe(engineObject.endpoint);
-    const account_info = await connection.resolveAccountInfo();
+    const account_info = await (
+      connection as unknown as ConnectionV1
+    ).resolveAccountInfo();
     expect(account_info.id).toBe("some_account");
-    expect(account_info.infraVersion).toBe(1);
   });
   it("Can run set statements", async () => {
     const param = "my_var";
diff --git a/test/unit/v2/connection.test.ts b/test/unit/v2/connection.test.ts

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-import { AccountInfo } from "../connection/base";`
	`1`	`+import { AccountInfo } from "../connection/connection_v1";`
`2`	`2`
`3`	`3`	`export type TokenKey = {`
`4`	`4`	`clientId: string;`