From f14c786b21d3e7946ba30e38aa83ea0abe81990e Mon Sep 17 00:00:00 2001
From: "Craig Boger (from Dev Box)" <bogercraig@microsoft.com>
Date: Mon, 12 Feb 2024 14:37:09 -0800
Subject: [PATCH] Rough implementation of a client map.  Need to figure out how
 to keep read client config in sync when other connection policy settings
 change.  If number of retry change, etc.

---
 src/Common/CosmosClient.ts | 50 +++++++++++++++++++++-----------------
 1 file changed, 28 insertions(+), 22 deletions(-)

diff --git a/src/Common/CosmosClient.ts b/src/Common/CosmosClient.ts
index 1cae31f20..daf00119e 100644
--- a/src/Common/CosmosClient.ts
+++ b/src/Common/CosmosClient.ts
@@ -131,32 +131,37 @@ enum SDKSupportedCapabilities {
 // Need to put in some kind of function here to recreate the CosmosClient with a new endpoint.
 // changeClientEndpoint.......
 
+let _clients: Map<string, Cosmos.CosmosClient> = new Map();
+
 let _client: Cosmos.CosmosClient;
 
 export function client(): Cosmos.CosmosClient {
   console.log(`Called primary client`);
   const currentUserContextDocumentEndpoint = userContext?.databaseAccount?.properties?.documentEndpoint;
   console.log(`Current selected endpoint in userContext: ${currentUserContextDocumentEndpoint}`);
-  let mydatabaseAccountEndpoint = "Ahhhhhhhhh";
-  if (_client) {
-    _client
-      .getDatabaseAccount()
-      .then((databaseAccount) => {
-        console.log(
-          `Current primary client endpoint contacted: ${JSON.stringify(
-            databaseAccount.diagnostics.clientSideRequestStatistics.locationEndpointsContacted,
-          )}`,
-        );
-        mydatabaseAccountEndpoint =
-          databaseAccount.diagnostics.clientSideRequestStatistics.locationEndpointsContacted[0];
-      })
-      .catch((error) => {
-        console.error("Error getting database account:", error);
-      });
-  }
+  // let mydatabaseAccountEndpoint = "Ahhhhhhhhh";
+  // if (_client) {
+  //   _client
+  //     .getDatabaseAccount()
+  //     .then((databaseAccount) => {
+  //       console.log(
+  //         `Current primary client endpoint contacted: ${JSON.stringify(
+  //           databaseAccount.diagnostics.clientSideRequestStatistics.locationEndpointsContacted,
+  //         )}`,
+  //       );
+  //       mydatabaseAccountEndpoint =
+  //         databaseAccount.diagnostics.clientSideRequestStatistics.locationEndpointsContacted[0];
+  //     })
+  //     .catch((error) => {
+  //       console.error("Error getting database account:", error);
+  //     });
+  // }
 
-  if (_client && currentUserContextDocumentEndpoint === mydatabaseAccountEndpoint) {
-    return _client;
+  const retrievedEndpoint = endpoint() || "https://cosmos.azure.com";
+
+  if (_clients.has(retrievedEndpoint)) {
+    console.log(`Current Client List: ${JSON.stringify(_clients)}`);
+    return _clients.get(retrievedEndpoint);
   }
 
   let _defaultHeaders: Cosmos.CosmosHeaders = {};
@@ -176,7 +181,7 @@ export function client(): Cosmos.CosmosClient {
   }
 
   const options: Cosmos.CosmosClientOptions = {
-    endpoint: endpoint() || "https://cosmos.azure.com", // CosmosClient gets upset if we pass a bad URL. This should never actually get called
+    endpoint: retrievedEndpoint, // CosmosClient gets upset if we pass a bad URL. This should never actually get called
     key: userContext.masterKey,
     tokenProvider,
     userAgentSuffix: "Azure Portal",
@@ -220,6 +225,7 @@ export function client(): Cosmos.CosmosClient {
     (options as any).plugins = plugins;
   }
 
-  _client = new Cosmos.CosmosClient(options);
-  return _client;
+  _clients.set(retrievedEndpoint, new Cosmos.CosmosClient(options));
+
+  return _clients.get(retrievedEndpoint);
 }