From 1c38b3e561ab6bb4bb2eea8a1c614d8660ef6e6e Mon Sep 17 00:00:00 2001
From: Karthikeyan C <karthikeyan@hasura.io>
Date: Wed, 18 Sep 2024 14:03:49 +0530
Subject: [PATCH] minor refactors

---
 src/cli/config.ts            |  4 ++--
 src/connector/connector.ts   | 25 +++++++++++++++++++++----
 src/connector/db/cosmosDb.ts | 27 ++++++++++++++++++++-------
 3 files changed, 43 insertions(+), 13 deletions(-)

diff --git a/src/cli/config.ts b/src/cli/config.ts
index 4df0127..0e021f3 100644
--- a/src/cli/config.ts
+++ b/src/cli/config.ts
@@ -292,7 +292,7 @@ type UpdateConfig = {
   connection: {
     endpoint: string;
     databaseName: string;
-    authenticationConfig: AzureCosmosAuthenticationConfig;
+    authentication: AzureCosmosAuthenticationConfig;
   };
   schema: CollectionsSchema;
 };
@@ -321,7 +321,7 @@ export async function generateConnectorConfig(outputConfigDir: string) {
       version: 2,
       connection: {
         endpoint: cosmosEndpoint,
-        authenticationConfig: connectionConfig,
+        authentication: connectionConfig,
         databaseName: cosmosDbName,
       },
       schema,
diff --git a/src/connector/connector.ts b/src/connector/connector.ts
index 35884ed..11a48df 100644
--- a/src/connector/connector.ts
+++ b/src/connector/connector.ts
@@ -1,6 +1,9 @@
 import * as sdk from "@hasura/ndc-sdk-typescript";
 import { CollectionsSchema, getNdcSchemaResponse } from "./schema";
-import { constructCosmosDbClient } from "./db/cosmosDb";
+import {
+  AzureCosmosAuthenticationConfig,
+  getCosmosDbClient,
+} from "./db/cosmosDb";
 import { Database } from "@azure/cosmos";
 import { executeQuery } from "./execution";
 import { readFileSync } from "fs";
@@ -10,7 +13,7 @@ export type Configuration = ConnectorConfig;
 export type ConnectorConfig = {
   connection: {
     endpoint: string;
-    key: string;
+    authentication: AzureCosmosAuthenticationConfig;
     databaseName: string;
   };
   schema: CollectionsSchema;
@@ -40,11 +43,25 @@ export function createConnector(): sdk.Connector<Configuration, State> {
     },
 
     tryInitState: async function (
-      _: Configuration,
+      config: Configuration,
       __: unknown,
     ): Promise<State> {
       try {
-        const databaseClient = constructCosmosDbClient().dbClient;
+        const {
+          databaseName,
+          authentication: authenticationConfig,
+          endpoint,
+        } = config.connection;
+        console.log(
+          "Initializing the state of the connector",
+          authenticationConfig,
+        );
+        const databaseClient = getCosmosDbClient(
+          endpoint,
+          databaseName,
+          authenticationConfig,
+        );
+
         return Promise.resolve({
           databaseClient,
         });
diff --git a/src/connector/db/cosmosDb.ts b/src/connector/db/cosmosDb.ts
index f2b96f7..d01ece7 100644
--- a/src/connector/db/cosmosDb.ts
+++ b/src/connector/db/cosmosDb.ts
@@ -4,12 +4,14 @@ import { throwError } from "../../utils";
 
 export type ManagedIdentityConfig = {
   type: "ManagedIdentity";
-  clientId: string;
+  // Name of the ENV var where the key can be found
+  fromEnvVar: string;
 };
 
 export type CosmosKeyConfig = {
   type: "Key";
-  key: string;
+  // Name of the ENV var where the key can be found
+  fromEnvVar: string;
 };
 
 export type AzureCosmosAuthenticationConfig =
@@ -23,7 +25,7 @@ export type RawCosmosDbConfig = {
 };
 
 /* Creates a new cosmos DB client with which the specified database can be queried. */
-function getCosmosDbClient(
+export function getCosmosDbClient(
   endpoint: string,
   databaseName: string,
   connectionConfig: AzureCosmosAuthenticationConfig,
@@ -31,14 +33,25 @@ function getCosmosDbClient(
   let dbClient: CosmosClient;
   switch (connectionConfig.type) {
     case "Key":
+      const key =
+        getEnvVariable(connectionConfig.fromEnvVar, true) ??
+        throwError(
+          `Azure Cosmos Key not found in the env var "${connectionConfig.fromEnvVar}"`,
+        );
       dbClient = new CosmosClient({
-        key: connectionConfig.key,
+        key,
         endpoint,
       });
       break;
     case "ManagedIdentity":
+      const managedIdentityClientId =
+        getEnvVariable(connectionConfig.fromEnvVar, true) ??
+        throwError(
+          `Azure Cosmos Key not found in the env var "${connectionConfig.fromEnvVar}"`,
+        );
+
       let credentials = new DefaultAzureCredential({
-        managedIdentityClientId: connectionConfig.clientId,
+        managedIdentityClientId,
       });
       dbClient = new CosmosClient({
         endpoint,
@@ -83,12 +96,12 @@ function getConnectionConfig(): AzureCosmosAuthenticationConfig | null {
     if (key) {
       return {
         type: "Key",
-        key,
+        fromEnvVar: "AZURE_COSMOS_KEY",
       };
     } else if (managed_identity_client_id) {
       return {
         type: "ManagedIdentity",
-        clientId: managed_identity_client_id,
+        fromEnvVar: "AZURE_COSMOS_MANAGED_CLIENT_ID",
       };
     }
   }