feat: Support multiple organizations API key. (#68)

.eslintrc.json

@@ -20,6 +20,7 @@
     ],
     "rules": {
         "no-useless-constructor": "off",
-        "@typescript-eslint/ban-ts-comment": "off"
+        "@typescript-eslint/ban-ts-comment": "off",
+        "@typescript-eslint/no-non-null-assertion": "off"
       }
 }

src/api-client.ts

@@ -2,26 +2,27 @@ import axios, { AxiosInstance } from "axios";
 import { ENV0_API_URL } from "./common";
 import { EnvironmentModel } from "./get-environments";
 import { DeploymentStepLogsResponse, DeploymentStepResponse } from "./types";
+import { AuthService } from "./auth";
 
 class ApiClient {
-  private credentials?: { username: string; password: string };
   private readonly instance: AxiosInstance;
+  private authService?: AuthService;
   constructor() {
     this.instance = axios.create({ baseURL: `https://${ENV0_API_URL}` });
     this.instance.interceptors.request.use((config) => {
-      if (this.credentials) {
-        config.auth = this.credentials;
+      const credentials = this.authService?.getCredentials();
+      if (credentials) {
+        config.auth = {
+          username: credentials.username,
+          password: credentials.password,
+        };
       }
       return config;
     });
   }
 
-  public init(credentials: { username: string; password: string }) {
-    this.credentials = credentials;
-  }
-
-  public clearCredentials() {
-    this.credentials = undefined;
+  public init(authService: AuthService) {
+    this.authService = authService;
   }
 
   public async abortDeployment(deploymentId: string) {
@@ -62,11 +63,14 @@ class ApiClient {
     return response.data;
   }
 
-  public async getEnvironments(organizationId: string) {
+  public async getEnvironments() {
     const response = await this.instance.get<EnvironmentModel[]>(
       `/environments`,
       {
-        params: { organizationId, isActive: true },
+        params: {
+          organizationId: this.authService?.getCredentials().selectedOrgId,
+          isActive: true,
+        },
       }
     );
 

src/auth.ts

@@ -6,11 +6,19 @@ import {
   showUnexpectedErrorMessage,
 } from "./notification-messages";
 
-const env0KeyIdKey = "env0.keyId";
-const env0SecretKey = "env0.secret";
+const env0KeyIdStoreKey = "env0.keyId";
+const env0SecretStoreKey = "env0.secret";
+const selectedOrgIdStoreKey = "env0.selectedOrgId";
 
 export class AuthService {
   constructor(private readonly context: vscode.ExtensionContext) {}
+
+  private credentials?: {
+    username: string;
+    password: string;
+    selectedOrgId: string;
+  };
+
   public registerLoginCommand(onLogin: () => void) {
     const disposable = vscode.commands.registerCommand(
       "env0.login",
@@ -47,10 +55,9 @@ export class AuthService {
           },
         });
 
-        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-        if (await this.validateUserCredentials(keyId!, secret!)) {
-          // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-          await this.storeAuthData(keyId!, secret!);
+        const selectedOrgId = await this.pickOrganization(keyId!, secret!);
+        if (selectedOrgId) {
+          await this.storeAuthData(keyId!, secret!, selectedOrgId);
           await onLogin();
         }
       }
@@ -70,58 +77,95 @@ export class AuthService {
   }
 
   public async isLoggedIn() {
-    const { secret, keyId } = await this.getAuthData();
-    return !!(secret && keyId);
+    const { secret, keyId, selectedOrgId } = await this.getAuthData();
+    if (!(secret && keyId && selectedOrgId)) {
+      return false;
+    }
+    this.credentials = { username: keyId, password: secret, selectedOrgId };
+    return true;
   }
 
-  public async getApiKeyCredentials() {
-    const { secret, keyId } = await this.getAuthData();
-    if (!secret || !keyId) {
-      throw new Error("Could not read env0 api key values");
+  public getCredentials() {
+    if (!this.credentials) {
+      // this should happen only if the user is logged out
+      throw new Error("Could not read credentials");
     }
-    return { username: keyId, password: secret };
+
+    return this.credentials;
   }
 
   private async getAuthData() {
     return {
-      keyId: await this.context.secrets.get(env0KeyIdKey),
-      secret: await this.context.secrets.get(env0SecretKey),
+      keyId: await this.context.secrets.get(env0KeyIdStoreKey),
+      secret: await this.context.secrets.get(env0SecretStoreKey),
+      selectedOrgId: await this.context.secrets.get(selectedOrgIdStoreKey),
     };
   }
 
-  private async validateUserCredentials(keyId: string, secret: string) {
+  private async pickOrganization(keyId: string, secret: string) {
     // Displaying a loading indicator to inform the user that something is happening
     return await vscode.window.withProgress(
       { location: { viewId: ENV0_ENVIRONMENTS_VIEW_ID } },
       async () => {
         try {
-          await axios.get(`https://${ENV0_API_URL}/organizations`, {
-            auth: { username: keyId, password: secret },
-            validateStatus: function (status) {
-              return status >= 200 && status < 300;
-            },
+          const orgsRes = await axios.get(
+            `https://${ENV0_API_URL}/organizations`,
+            {
+              auth: { username: keyId, password: secret },
+              validateStatus: function (status) {
+                return status >= 200 && status < 300;
+              },
+            }
+          );
+          if (orgsRes.data.length === 1) {
+            return orgsRes.data[0].id;
+          }
+          const orgs = orgsRes.data.map((org: any) => ({
+            name: org.name,
+            id: org.id,
+          }));
+          const items: vscode.QuickPickItem[] = orgs.map((org: any) => ({
+            label: org.name,
+            description: org.id,
+          }));
+
+          const selectedItem = await vscode.window.showQuickPick(items, {
+            placeHolder: "Select an organization",
+            ignoreFocusOut: true,
           });
-          return true;
+          const selectedOrgId = selectedItem?.description;
+          if (!selectedOrgId) {
+            vscode.window.showErrorMessage("No organization selected");
+            return undefined;
+          }
+          return selectedOrgId;
           // eslint-disable-next-line @typescript-eslint/no-explicit-any
         } catch (e: any) {
           if (e?.response?.status >= 400 && e?.response?.status < 500) {
             showInvalidCredentialsMessage();
           } else {
             showUnexpectedErrorMessage();
           }
-          return false;
+          return undefined;
         }
       }
     );
   }
 
-  private async storeAuthData(keyId: string, secret: string) {
-    await this.context.secrets.store(env0KeyIdKey, keyId);
-    await this.context.secrets.store(env0SecretKey, secret);
+  private async storeAuthData(
+    keyId: string,
+    secret: string,
+    selectedOrgId: string
+  ) {
+    this.credentials = { username: keyId, password: secret, selectedOrgId };
+    await this.context.secrets.store(env0KeyIdStoreKey, keyId);
+    await this.context.secrets.store(env0SecretStoreKey, secret);
+    await this.context.secrets.store(selectedOrgIdStoreKey, selectedOrgId);
   }
 
   private async clearAuthData() {
-    await this.context.secrets.delete(env0KeyIdKey);
-    await this.context.secrets.delete(env0SecretKey);
+    await this.context.secrets.delete(env0KeyIdStoreKey);
+    await this.context.secrets.delete(env0SecretStoreKey);
+    await this.context.secrets.delete(selectedOrgIdStoreKey);
   }
 }

src/extension.ts

@@ -70,7 +70,7 @@ const init = async (
   environmentsTree: vscode.TreeView<Environment>,
   authService: AuthService
 ) => {
-  apiClient.init(await authService.getApiKeyCredentials());
+  apiClient.init(authService);
   extensionState.setLoggedIn(true);
   await loadEnvironments(environmentsDataProvider);
 };
@@ -81,7 +81,6 @@ const onLogOut = async () => {
   }
   stopEnvironmentPolling();
   environmentsDataProvider.clear();
-  apiClient.clearCredentials();
   extensionState.setLoggedIn(false);
 };
 

src/get-environments.ts

@@ -38,11 +38,7 @@ function repositoriesEqual(rep1: string, rep2: string): boolean {
 export async function getEnvironmentsForBranch() {
   let environments: EnvironmentModel[] = [];
 
-  const organizationId = await getOrganizationId();
-
-  if (organizationId) {
-    environments = await getEnvironments(organizationId);
-  }
+  environments = await apiClient.getEnvironments();
 
   if (environments.length > 0) {
     const { currentBranch, repository } = getGitRepoAndBranch();
@@ -63,18 +59,3 @@ export async function getEnvironmentsForBranch() {
   }
   return environments;
 }
-
-async function getEnvironments(organizationId: string) {
-  try {
-    return apiClient.getEnvironments(organizationId);
-  } catch (e) {
-    console.log(e);
-  }
-
-  return [];
-}
-
-async function getOrganizationId() {
-  const organizations = await apiClient.getOrganizations();
-  return organizations[0]?.id;
-}

src/test/integration/mocks/quick-pick.ts

@@ -0,0 +1,12 @@
+import * as jestMock from "jest-mock";
+import * as vscode from "vscode";
+
+let quickPickSpy: jestMock.SpyInstance<any>;
+export const spyOnQuickPick = () => {
+  quickPickSpy = jestMock.spyOn(vscode.window, "showQuickPick");
+  return quickPickSpy;
+};
+
+export const resetSpyOnQuickPick = () => {
+  quickPickSpy?.mockReset?.();
+};

src/test/integration/mocks/server.ts

@@ -28,16 +28,16 @@ const assertAuth = (credentials: Credentials, authHeader: string | null) => {
   assert.strictEqual(secret, credentials.secret);
 };
 
-export const mockGetOrganization = (
-  organizationId: string,
+export const mockGetOrganizations = (
+  orgs: { id: string; name: string }[],
   credentials?: Credentials
 ) => {
   server.use(
     rest.get(`https://${ENV0_API_URL}/organizations`, (req, res, ctx) => {
       if (credentials) {
         assertAuth(credentials, req.headers.get("Authorization"));
       }
-      return res(ctx.json([{ id: organizationId }]));
+      return res(ctx.json(orgs));
     })
   );
 };