From b45e667d516c2d4668838b3b336b8fe56c80db26 Mon Sep 17 00:00:00 2001
From: "Craig Boger (from Dev Box)" <bogercraig@microsoft.com>
Date: Fri, 16 Feb 2024 15:05:26 -0800
Subject: [PATCH] Reverting back to a single client for memory testing.

---
 src/Common/CosmosClient.ts | 54 ++++++++++++++++++++++----------------
 1 file changed, 31 insertions(+), 23 deletions(-)

diff --git a/src/Common/CosmosClient.ts b/src/Common/CosmosClient.ts
index daf00119e..cbefa2402 100644
--- a/src/Common/CosmosClient.ts
+++ b/src/Common/CosmosClient.ts
@@ -131,7 +131,7 @@ 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 _clients: Map<string, Cosmos.CosmosClient> = new Map();
 
 let _client: Cosmos.CosmosClient;
 
@@ -139,29 +139,33 @@ 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);
+      });
+  }
 
   const retrievedEndpoint = endpoint() || "https://cosmos.azure.com";
 
-  if (_clients.has(retrievedEndpoint)) {
-    console.log(`Current Client List: ${JSON.stringify(_clients)}`);
-    return _clients.get(retrievedEndpoint);
+  // if (_clients.has(retrievedEndpoint)) {
+  //   console.log(`Current Client List: ${JSON.stringify(_clients)}`);
+  //   return _clients.get(retrievedEndpoint);
+  // }
+
+  if (_client && currentUserContextDocumentEndpoint === mydatabaseAccountEndpoint) {
+    return _client;
   }
 
   let _defaultHeaders: Cosmos.CosmosHeaders = {};
@@ -225,7 +229,11 @@ export function client(): Cosmos.CosmosClient {
     (options as any).plugins = plugins;
   }
 
-  _clients.set(retrievedEndpoint, new Cosmos.CosmosClient(options));
+  _client = null;
+  _client = new Cosmos.CosmosClient(options);
+  return _client;
 
-  return _clients.get(retrievedEndpoint);
+  // _clients.set(retrievedEndpoint, new Cosmos.CosmosClient(options));
+
+  // return _clients.get(retrievedEndpoint);
 }