Updated to have a write client and multiple read clients. Added enum to help with selection of client based on needed operation.

Need to modify endpoint selection in CosmosClient to select endpoints based on client called.

src/Common/CosmosClient.ts

@@ -78,10 +78,11 @@ export const tokenProvider = async (requestInfo: Cosmos.RequestInfo) => {
   return decodeURIComponent(result.PrimaryReadWriteToken);
 };
 
+// TODO
+// Need to create separate plugins or change endpoint logic to return the correct write or read endpoint.
 export const requestPlugin: Cosmos.Plugin<any> = async (requestContext, diagnosticNode, next) => {
   requestContext.endpoint = new URL(configContext.PROXY_PATH, window.location.href).href;
   requestContext.headers["x-ms-proxy-target"] = endpoint();
-  // console.log(`Request context: ${JSON.stringify(requestContext)}`);
   return next(requestContext);
 };
 
@@ -128,36 +129,89 @@ enum SDKSupportedCapabilities {
   PartitionMerge = 1 << 0,
 }
 
-// Need to put in some kind of function here to recreate the CosmosClient with a new endpoint.
+// Client Management
-// changeClientEndpoint.......
 
 let _client: Cosmos.CosmosClient;
-let _currentClientEndpoint: string;
+let _readClients: Map<string, Cosmos.CosmosClient> = new Map();
 
-export function client(): Cosmos.CosmosClient {
+export enum ClientOperationType {
+  READ,
+  WRITE,
+}
+
+export function client(clientOperationType: ClientOperationType): Cosmos.CosmosClient {
+  switch (clientOperationType) {
+    case ClientOperationType.READ:
+      return readClients();
+    case ClientOperationType.WRITE:
+      return writeClient();
+    default:
+      throw new Error("Invalid operation type");
+  }
+}
+
+export function writeClient(): 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);
-  //     });
-  // }
 
-  if (_client && currentUserContextDocumentEndpoint === _currentClientEndpoint) {
+  if (_client) return _client;
+
+  let _defaultHeaders: Cosmos.CosmosHeaders = {};
+  _defaultHeaders["x-ms-cosmos-sdk-supportedcapabilities"] =
+    SDKSupportedCapabilities.None | SDKSupportedCapabilities.PartitionMerge;
+
+  if (
+    userContext.authType === AuthType.ConnectionString ||
+    userContext.authType === AuthType.EncryptedToken ||
+    userContext.authType === AuthType.ResourceToken
+  ) {
+    // Default to low priority. Needed for non-AAD-auth scenarios
+    // where we cannot use RP API, and thus, cannot detect whether priority
+    // based execution is enabled.
+    // The header will be ignored if priority based execution is disabled on the account.
+    _defaultHeaders["x-ms-cosmos-priority-level"] = PriorityLevel.Default;
+  }
+
+  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
+    key: userContext.masterKey,
+    tokenProvider,
+    userAgentSuffix: "Azure Portal",
+    defaultHeaders: _defaultHeaders,
+    connectionPolicy: {
+      retryOptions: {
+        maxRetryAttemptCount: LocalStorageUtility.getEntryNumber(StorageKey.RetryAttempts),
+        fixedRetryIntervalInMilliseconds: LocalStorageUtility.getEntryNumber(StorageKey.RetryInterval),
+        maxWaitTimeInSeconds: LocalStorageUtility.getEntryNumber(StorageKey.MaxWaitTimeInSeconds),
+      },
+    },
+  };
+
+  if (configContext.PROXY_PATH !== undefined) {
+    (options as any).plugins = [{ on: "request", plugin: requestPlugin }];
+  }
+
+  if (PriorityBasedExecutionUtils.isFeatureEnabled()) {
+    const plugins = (options as any).plugins || [];
+    plugins.push({ on: "request", plugin: PriorityBasedExecutionUtils.requestPlugin });
+    (options as any).plugins = plugins;
+  }
+
+  _client = new Cosmos.CosmosClient(options);
   return _client;
+}
+
+export function readClients(): Cosmos.CosmosClient {
+  console.log(`Called read only client`);
+  const currentUserContextDocumentEndpoint = userContext?.databaseAccount?.properties?.documentEndpoint;
+  console.log(`Current selected read endpoint in userContext: ${currentUserContextDocumentEndpoint}`);
+  3;
+
+  const selectedEndpoint = endpoint() || "https://cosmos.azure.com";
+
+  if (_readClients.has(selectedEndpoint)) {
+    return _readClients.get(selectedEndpoint);
   }
 
   let _defaultHeaders: Cosmos.CosmosHeaders = {};
@@ -176,11 +230,8 @@ export function client(): Cosmos.CosmosClient {
     _defaultHeaders["x-ms-cosmos-priority-level"] = PriorityLevel.Default;
   }
 
-  const clientEndpoint = endpoint() || "https://cosmos.azure.com";
-  _currentClientEndpoint = clientEndpoint;
-
   const options: Cosmos.CosmosClientOptions = {
-    endpoint: clientEndpoint, // CosmosClient gets upset if we pass a bad URL. This should never actually get called
+    endpoint: selectedEndpoint, // CosmosClient gets upset if we pass a bad URL. This should never actually get called
     key: userContext.masterKey,
     tokenProvider,
     userAgentSuffix: "Azure Portal",
@@ -194,26 +245,6 @@ export function client(): Cosmos.CosmosClient {
     },
   };
 
-  // Account details from userContext.
-  // console.log(`userContext details: ${JSON.stringify(userContext)}`);
-  // console.log(`userContext.databaseaccount details: ${JSON.stringify(userContext.databaseAccount)}`);
-  console.log(
-    `userContext?.databaseAccount?.properties?.documentEndpoint details: ${JSON.stringify(
-      userContext?.databaseAccount?.properties?.documentEndpoint,
-    )}`,
-  );
-  // console.log(`userContext?.endpoint details: ${JSON.stringify(userContext?.endpoint)}`);
-  // console.log(
-  //   `userContext?.databaseAccount?.properties?.readLocations details: ${JSON.stringify(
-  //     userContext?.databaseAccount?.properties?.readLocations,
-  //   )}`,
-  // );
-  // console.log(
-  //   `userContext?.databaseAccount?.properties?.writeLocations details: ${JSON.stringify(
-  //     userContext?.databaseAccount?.properties?.writeLocations,
-  //   )}`,
-  // );
-
   if (configContext.PROXY_PATH !== undefined) {
     (options as any).plugins = [{ on: "request", plugin: requestPlugin }];
   }
@@ -224,6 +255,7 @@ export function client(): Cosmos.CosmosClient {
     (options as any).plugins = plugins;
   }
 
-  _client = new Cosmos.CosmosClient(options);
+  _readClients.set(selectedEndpoint, new Cosmos.CosmosClient(options));
-  return _client;
+
+  return _readClients.get(selectedEndpoint);
 }

src/Common/ReadRegionCosmosClient.ts

@@ -1,200 +0,0 @@
-import * as Cosmos from "@azure/cosmos";
-import { sendCachedDataMessage } from "Common/MessageHandler";
-import { getAuthorizationTokenUsingResourceTokens } from "Common/getAuthorizationTokenUsingResourceTokens";
-import { AuthorizationToken, MessageTypes } from "Contracts/MessageTypes";
-import { checkDatabaseResourceTokensValidity } from "Platform/Fabric/FabricUtil";
-import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
-import { AuthType } from "../AuthType";
-import { PriorityLevel } from "../Common/Constants";
-import { Platform, configContext } from "../ConfigContext";
-import { userContext } from "../UserContext";
-import { logConsoleError } from "../Utils/NotificationConsoleUtils";
-import { EmulatorMasterKey, HttpHeaders } from "./Constants";
-import { getErrorMessage } from "./ErrorHandlingUtils";
-
-const _global = typeof self === "undefined" ? window : self;
-
-export const tokenProvider2 = async (requestInfo: Cosmos.RequestInfo) => {
-  const { verb, resourceId, resourceType, headers } = requestInfo;
-
-  if (userContext.features.enableAadDataPlane && userContext.aadToken) {
-    const AUTH_PREFIX = `type=aad&ver=1.0&sig=`;
-    const authorizationToken = `${AUTH_PREFIX}${userContext.aadToken}`;
-    return authorizationToken;
-  }
-
-  if (configContext.platform === Platform.Emulator) {
-    // TODO This SDK method mutates the headers object. Find a better one or fix the SDK.
-    await Cosmos.setAuthorizationTokenHeaderUsingMasterKey(verb, resourceId, resourceType, headers, EmulatorMasterKey);
-    return decodeURIComponent(headers.authorization);
-  }
-
-  if (configContext.platform === Platform.Fabric) {
-    switch (requestInfo.resourceType) {
-      case Cosmos.ResourceType.conflicts:
-      case Cosmos.ResourceType.container:
-      case Cosmos.ResourceType.sproc:
-      case Cosmos.ResourceType.udf:
-      case Cosmos.ResourceType.trigger:
-      case Cosmos.ResourceType.item:
-      case Cosmos.ResourceType.pkranges:
-        // User resource tokens
-        // TODO userContext.fabricContext.databaseConnectionInfo can be undefined
-        headers[HttpHeaders.msDate] = new Date().toUTCString();
-        const resourceTokens = userContext.fabricContext.databaseConnectionInfo.resourceTokens;
-        checkDatabaseResourceTokensValidity(userContext.fabricContext.databaseConnectionInfo.resourceTokensTimestamp);
-        return getAuthorizationTokenUsingResourceTokens(resourceTokens, requestInfo.path, requestInfo.resourceId);
-
-      case Cosmos.ResourceType.none:
-      case Cosmos.ResourceType.database:
-      case Cosmos.ResourceType.offer:
-      case Cosmos.ResourceType.user:
-      case Cosmos.ResourceType.permission:
-        // User master tokens
-        const authorizationToken = await sendCachedDataMessage<AuthorizationToken>(
-          MessageTypes.GetAuthorizationToken,
-          [requestInfo],
-          userContext.fabricContext.connectionId,
-        );
-        console.log("Response from Fabric: ", authorizationToken);
-        headers[HttpHeaders.msDate] = authorizationToken.XDate;
-        return decodeURIComponent(authorizationToken.PrimaryReadWriteToken);
-    }
-  }
-
-  if (userContext.masterKey) {
-    // TODO This SDK method mutates the headers object. Find a better one or fix the SDK.
-    await Cosmos.setAuthorizationTokenHeaderUsingMasterKey(verb, resourceId, resourceType, headers, EmulatorMasterKey);
-    return decodeURIComponent(headers.authorization);
-  }
-
-  if (userContext.resourceToken) {
-    return userContext.resourceToken;
-  }
-
-  const result = await getTokenFromAuthService2(verb, resourceType, resourceId);
-  headers[HttpHeaders.msDate] = result.XDate;
-  return decodeURIComponent(result.PrimaryReadWriteToken);
-};
-
-export const requestPlugin2: Cosmos.Plugin<any> = async (requestContext, diagnosticNode, next) => {
-  requestContext.endpoint = new URL(configContext.PROXY_PATH, window.location.href).href;
-  requestContext.headers["x-ms-proxy-target"] = endpoint2();
-  console.log(`Client2 request context: ${JSON.stringify(requestContext)}`);
-  return next(requestContext);
-};
-
-export const endpoint2 = () => {
-  if (configContext.platform === Platform.Emulator) {
-    // In worker scope, _global(self).parent does not exist
-    const location = _global.parent ? _global.parent.location : _global.location;
-    return configContext.EMULATOR_ENDPOINT || location.origin;
-  }
-  // return userContext.endpoint || userContext?.databaseAccount?.properties?.documentEndpoint;
-  return "https://test-craig-nosql-periodic-eastus.documents.azure.com:443/";
-};
-
-export async function getTokenFromAuthService2(
-  verb: string,
-  resourceType: string,
-  resourceId?: string,
-): Promise<AuthorizationToken> {
-  try {
-    const host = configContext.BACKEND_ENDPOINT;
-    const response = await _global.fetch(host + "/api/guest/runtimeproxy/authorizationTokens", {
-      method: "POST",
-      headers: {
-        "content-type": "application/json",
-        "x-ms-encrypted-auth-token": userContext.accessToken,
-      },
-      body: JSON.stringify({
-        verb,
-        resourceType,
-        resourceId,
-      }),
-    });
-    //TODO I am not sure why we have to parse the JSON again here. fetch should do it for us when we call .json()
-    const result = JSON.parse(await response.json());
-    return result;
-  } catch (error) {
-    logConsoleError(`Failed to get authorization headers for ${resourceType}: ${getErrorMessage(error)}`);
-    return Promise.reject(error);
-  }
-}
-
-// The Capability is a bitmap, which cosmosdb backend decodes as per the below enum
-enum SDKSupportedCapabilities {
-  None = 0,
-  PartitionMerge = 1 << 0,
-}
-
-let _client2: Cosmos.CosmosClient;
-
-export function client2(): Cosmos.CosmosClient {
-  console.log(`Called client2`);
-  if (_client2) return _client2;
-
-  let _defaultHeaders: Cosmos.CosmosHeaders = {};
-  _defaultHeaders["x-ms-cosmos-sdk-supportedcapabilities"] =
-    SDKSupportedCapabilities.None | SDKSupportedCapabilities.PartitionMerge;
-
-  if (
-    userContext.authType === AuthType.ConnectionString ||
-    userContext.authType === AuthType.EncryptedToken ||
-    userContext.authType === AuthType.ResourceToken
-  ) {
-    // Default to low priority. Needed for non-AAD-auth scenarios
-    // where we cannot use RP API, and thus, cannot detect whether priority
-    // based execution is enabled.
-    // The header will be ignored if priority based execution is disabled on the account.
-    _defaultHeaders["x-ms-cosmos-priority-level"] = PriorityLevel.Default;
-  }
-
-  const options: Cosmos.CosmosClientOptions = {
-    endpoint: endpoint2() || "https://cosmos.azure.com", // CosmosClient gets upset if we pass a bad URL. This should never actually get called
-    key: userContext.masterKey,
-    tokenProvider: tokenProvider2,
-    userAgentSuffix: "Azure Portal",
-    defaultHeaders: _defaultHeaders,
-    connectionPolicy: {
-      retryOptions: {
-        maxRetryAttemptCount: LocalStorageUtility.getEntryNumber(StorageKey.RetryAttempts),
-        fixedRetryIntervalInMilliseconds: LocalStorageUtility.getEntryNumber(StorageKey.RetryInterval),
-        maxWaitTimeInSeconds: LocalStorageUtility.getEntryNumber(StorageKey.MaxWaitTimeInSeconds),
-      },
-    },
-  };
-
-  // Account details from userContext.
-  console.log(`userContext details: ${JSON.stringify(userContext)}`);
-  console.log(`userContext.databaseaccount details: ${JSON.stringify(userContext.databaseAccount)}`);
-  console.log(
-    `userContext?.databaseAccount?.properties?.documentEndpoint details: ${JSON.stringify(
-      userContext?.databaseAccount?.properties?.documentEndpoint,
-    )}`,
-  );
-  console.log(`userContext?.endpoint details: ${JSON.stringify(userContext?.endpoint)}`);
-  console.log(
-    `userContext?.databaseAccount?.properties?.readLocations details: ${JSON.stringify(
-      userContext?.databaseAccount?.properties?.readLocations,
-    )}`,
-  );
-  console.log(
-    `userContext?.databaseAccount?.properties?.writeLocations details: ${JSON.stringify(
-      userContext?.databaseAccount?.properties?.writeLocations,
-    )}`,
-  );
-
-  if (configContext.PROXY_PATH !== undefined) {
-    (options as any).plugins = [{ on: "request", plugin: requestPlugin2 }];
-  }
-
-  // if (PriorityBasedExecutionUtils.isFeatureEnabled()) {
-  //   const plugins = (options as any).plugins || [];
-  //   plugins.push({ on: "request", plugin: PriorityBasedExecutionUtils.requestPlugin });
-  //   (options as any).plugins = plugins;
-  // }
-
-  _client2 = new Cosmos.CosmosClient(options);
-  return _client2;
-}

src/Common/dataAccess/bulkCreateDocument.ts

@@ -1,7 +1,7 @@
 import { JSONObject, OperationResponse } from "@azure/cosmos";
 import { CollectionBase } from "../../Contracts/ViewModels";
 import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-import { client } from "../CosmosClient";
+import { ClientOperationType, client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export const bulkCreateDocument = async (
@@ -13,7 +13,7 @@ export const bulkCreateDocument = async (
   );
 
   try {
-    const response = await client()
+    const response = await client(ClientOperationType.WRITE)
       .database(collection.databaseId)
       .container(collection.id())
       .items.bulk(

src/Common/dataAccess/createCollection.ts

@@ -6,14 +6,14 @@ import { Action, ActionModifiers } from "../../Shared/Telemetry/TelemetryConstan
 import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";
 import { userContext } from "../../UserContext";
 import { getCollectionName } from "../../Utils/APITypeUtils";
+import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { createUpdateCassandraTable } from "../../Utils/arm/generatedClients/cosmos/cassandraResources";
 import { createUpdateGremlinGraph } from "../../Utils/arm/generatedClients/cosmos/gremlinResources";
 import { createUpdateMongoDBCollection } from "../../Utils/arm/generatedClients/cosmos/mongoDBResources";
 import { createUpdateSqlContainer } from "../../Utils/arm/generatedClients/cosmos/sqlResources";
 import { createUpdateTable } from "../../Utils/arm/generatedClients/cosmos/tableResources";
 import * as ARMTypes from "../../Utils/arm/generatedClients/cosmos/types";
-import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
+import { ClientOperationType, client } from "../CosmosClient";
-import { client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 import { createMongoCollectionWithProxy } from "../MongoProxyClient";
 import { createDatabase } from "./createDatabase";
@@ -284,7 +284,9 @@ const createCollectionWithSDK = async (params: DataModels.CreateCollectionParams
     }
   }
 
-  const databaseResponse: DatabaseResponse = await client().databases.createIfNotExists(createDatabaseBody);
+  const databaseResponse: DatabaseResponse = await client(ClientOperationType.WRITE).databases.createIfNotExists(
+    createDatabaseBody,
+  );
   const collectionResponse: ContainerResponse = await databaseResponse?.database.containers.create(
     createCollectionBody,
     collectionOptions,

src/Common/dataAccess/createDatabase.ts

@@ -4,6 +4,7 @@ import * as DataModels from "../../Contracts/DataModels";
 import { useDatabases } from "../../Explorer/useDatabases";
 import { userContext } from "../../UserContext";
 import { getDatabaseName } from "../../Utils/APITypeUtils";
+import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { createUpdateCassandraKeyspace } from "../../Utils/arm/generatedClients/cosmos/cassandraResources";
 import { createUpdateGremlinDatabase } from "../../Utils/arm/generatedClients/cosmos/gremlinResources";
 import { createUpdateMongoDBDatabase } from "../../Utils/arm/generatedClients/cosmos/mongoDBResources";
@@ -15,8 +16,7 @@ import {
   MongoDBDatabaseCreateUpdateParameters,
   SqlDatabaseCreateUpdateParameters,
 } from "../../Utils/arm/generatedClients/cosmos/types";
-import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
+import { ClientOperationType, client } from "../CosmosClient";
-import { client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function createDatabase(params: DataModels.CreateDatabaseParams): Promise<DataModels.Database> {
@@ -153,7 +153,7 @@ async function createDatabaseWithSDK(params: DataModels.CreateDatabaseParams): P
     }
   }
 
-  const response: DatabaseResponse = await client().databases.create(createBody);
+  const response: DatabaseResponse = await client(ClientOperationType.WRITE).databases.create(createBody);
   return response.resource;
 }
 

src/Common/dataAccess/createDocument.ts

@@ -1,6 +1,6 @@
 import { CollectionBase } from "../../Contracts/ViewModels";
 import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-import { client } from "../CosmosClient";
+import { ClientOperationType, client } from "../CosmosClient";
 import { getEntityName } from "../DocumentUtility";
 import { handleError } from "../ErrorHandlingUtils";
 
@@ -9,7 +9,7 @@ export const createDocument = async (collection: CollectionBase, newDocument: un
   const clearMessage = logConsoleProgress(`Creating new ${entityName} for container ${collection.id()}`);
 
   try {
-    const response = await client()
+    const response = await client(ClientOperationType.WRITE)
       .database(collection.databaseId)
       .container(collection.id())
       .items.create(newDocument);

src/Common/dataAccess/createStoredProcedure.ts

@@ -1,6 +1,7 @@
 import { Resource, StoredProcedureDefinition } from "@azure/cosmos";
 import { AuthType } from "../../AuthType";
 import { userContext } from "../../UserContext";
+import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import {
   createUpdateSqlStoredProcedure,
   getSqlStoredProcedure,
@@ -9,8 +10,7 @@ import {
   SqlStoredProcedureCreateUpdateParameters,
   SqlStoredProcedureResource,
 } from "../../Utils/arm/generatedClients/cosmos/types";
-import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
+import { ClientOperationType, client } from "../CosmosClient";
-import { client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function createStoredProcedure(
@@ -63,7 +63,7 @@ export async function createStoredProcedure(
       return rpResponse && (rpResponse.properties?.resource as StoredProcedureDefinition & Resource);
     }
 
-    const response = await client()
+    const response = await client(ClientOperationType.WRITE)
       .database(databaseId)
       .container(collectionId)
       .scripts.storedProcedures.create(storedProcedure);

src/Common/dataAccess/createTrigger.ts

@@ -1,10 +1,10 @@
 import { TriggerDefinition } from "@azure/cosmos";
 import { AuthType } from "../../AuthType";
 import { userContext } from "../../UserContext";
+import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { createUpdateSqlTrigger, getSqlTrigger } from "../../Utils/arm/generatedClients/cosmos/sqlResources";
 import { SqlTriggerCreateUpdateParameters, SqlTriggerResource } from "../../Utils/arm/generatedClients/cosmos/types";
-import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
+import { ClientOperationType, client } from "../CosmosClient";
-import { client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function createTrigger(
@@ -55,7 +55,7 @@ export async function createTrigger(
       return rpResponse && rpResponse.properties?.resource;
     }
 
-    const response = await client()
+    const response = await client(ClientOperationType.WRITE)
       .database(databaseId)
       .container(collectionId)
       .scripts.triggers.create(trigger as unknown as TriggerDefinition); // TODO: TypeScript does not like the SQL SDK trigger type

src/Common/dataAccess/createUserDefinedFunction.ts

@@ -1,6 +1,7 @@
 import { Resource, UserDefinedFunctionDefinition } from "@azure/cosmos";
 import { AuthType } from "../../AuthType";
 import { userContext } from "../../UserContext";
+import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import {
   createUpdateSqlUserDefinedFunction,
   getSqlUserDefinedFunction,
@@ -9,8 +10,7 @@ import {
   SqlUserDefinedFunctionCreateUpdateParameters,
   SqlUserDefinedFunctionResource,
 } from "../../Utils/arm/generatedClients/cosmos/types";
-import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
+import { ClientOperationType, client } from "../CosmosClient";
-import { client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function createUserDefinedFunction(
@@ -63,7 +63,7 @@ export async function createUserDefinedFunction(
       return rpResponse && (rpResponse.properties?.resource as UserDefinedFunctionDefinition & Resource);
     }
 
-    const response = await client()
+    const response = await client(ClientOperationType.WRITE)
       .database(databaseId)
       .container(collectionId)
       .scripts.userDefinedFunctions.create(userDefinedFunction);

src/Common/dataAccess/deleteCollection.ts

@@ -6,7 +6,7 @@ import { deleteMongoDBCollection } from "../../Utils/arm/generatedClients/cosmos
 import { deleteSqlContainer } from "../../Utils/arm/generatedClients/cosmos/sqlResources";
 import { deleteTable } from "../../Utils/arm/generatedClients/cosmos/tableResources";
 import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-import { client } from "../CosmosClient";
+import { client, ClientOperationType } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function deleteCollection(databaseId: string, collectionId: string): Promise<void> {
@@ -15,7 +15,7 @@ export async function deleteCollection(databaseId: string, collectionId: string)
     if (userContext.authType === AuthType.AAD && !userContext.features.enableSDKoperations) {
       await deleteCollectionWithARM(databaseId, collectionId);
     } else {
-      await client().database(databaseId).container(collectionId).delete();
+      await client(ClientOperationType.WRITE).database(databaseId).container(collectionId).delete();
     }
     logConsoleInfo(`Successfully deleted container ${collectionId}`);
   } catch (error) {

src/Common/dataAccess/deleteConflict.ts

@@ -1,9 +1,9 @@
-import ConflictId from "../../Explorer/Tree/ConflictId";
-import { CollectionBase } from "../../Contracts/ViewModels";
 import { RequestOptions } from "@azure/cosmos";
-import { client } from "../CosmosClient";
+import { CollectionBase } from "../../Contracts/ViewModels";
+import ConflictId from "../../Explorer/Tree/ConflictId";
+import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
+import { ClientOperationType, client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
-import { logConsoleProgress, logConsoleInfo } from "../../Utils/NotificationConsoleUtils";
 
 export const deleteConflict = async (collection: CollectionBase, conflictId: ConflictId): Promise<void> => {
   const clearMessage = logConsoleProgress(`Deleting conflict ${conflictId.id()}`);
@@ -13,7 +13,7 @@ export const deleteConflict = async (collection: CollectionBase, conflictId: Con
       partitionKey: getPartitionKeyHeaderForConflict(conflictId),
     };
 
-    await client()
+    await client(ClientOperationType.WRITE)
       .database(collection.databaseId)
       .container(collection.id())
       .conflict(conflictId.id())

src/Common/dataAccess/deleteDatabase.ts

@@ -5,7 +5,7 @@ import { deleteGremlinDatabase } from "../../Utils/arm/generatedClients/cosmos/g
 import { deleteMongoDBDatabase } from "../../Utils/arm/generatedClients/cosmos/mongoDBResources";
 import { deleteSqlDatabase } from "../../Utils/arm/generatedClients/cosmos/sqlResources";
 import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-import { client } from "../CosmosClient";
+import { client, ClientOperationType } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function deleteDatabase(databaseId: string): Promise<void> {
@@ -15,7 +15,7 @@ export async function deleteDatabase(databaseId: string): Promise<void> {
     if (userContext.authType === AuthType.AAD && !userContext.features.enableSDKoperations) {
       await deleteDatabaseWithARM(databaseId);
     } else {
-      await client().database(databaseId).delete();
+      await client(ClientOperationType.WRITE).database(databaseId).delete();
     }
     logConsoleInfo(`Successfully deleted database ${databaseId}`);
   } catch (error) {

src/Common/dataAccess/deleteDocument.ts

@@ -1,7 +1,7 @@
 import { CollectionBase } from "../../Contracts/ViewModels";
 import DocumentId from "../../Explorer/Tree/DocumentId";
 import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-import { client } from "../CosmosClient";
+import { ClientOperationType, client } from "../CosmosClient";
 import { getEntityName } from "../DocumentUtility";
 import { handleError } from "../ErrorHandlingUtils";
 import { getPartitionKeyValue } from "./getPartitionKeyValue";
@@ -11,7 +11,7 @@ export const deleteDocument = async (collection: CollectionBase, documentId: Doc
   const clearMessage = logConsoleProgress(`Deleting ${entityName} ${documentId.id()}`);
 
   try {
-    await client()
+    await client(ClientOperationType.WRITE)
       .database(collection.databaseId)
       .container(collection.id())
       .item(documentId.id(), getPartitionKeyValue(documentId))

src/Common/dataAccess/deleteStoredProcedure.ts

@@ -2,7 +2,7 @@ import { AuthType } from "../../AuthType";
 import { userContext } from "../../UserContext";
 import { deleteSqlStoredProcedure } from "../../Utils/arm/generatedClients/cosmos/sqlResources";
 import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-import { client } from "../CosmosClient";
+import { client, ClientOperationType } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function deleteStoredProcedure(
@@ -26,7 +26,11 @@ export async function deleteStoredProcedure(
         storedProcedureId,
       );
     } else {
-      await client().database(databaseId).container(collectionId).scripts.storedProcedure(storedProcedureId).delete();
+      await client(ClientOperationType.WRITE)
+        .database(databaseId)
+        .container(collectionId)
+        .scripts.storedProcedure(storedProcedureId)
+        .delete();
     }
   } catch (error) {
     handleError(error, "DeleteStoredProcedure", `Error while deleting stored procedure ${storedProcedureId}`);

src/Common/dataAccess/deleteTrigger.ts

@@ -2,7 +2,7 @@ import { AuthType } from "../../AuthType";
 import { userContext } from "../../UserContext";
 import { deleteSqlTrigger } from "../../Utils/arm/generatedClients/cosmos/sqlResources";
 import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-import { client } from "../CosmosClient";
+import { client, ClientOperationType } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function deleteTrigger(databaseId: string, collectionId: string, triggerId: string): Promise<void> {
@@ -22,7 +22,11 @@ export async function deleteTrigger(databaseId: string, collectionId: string, tr
         triggerId,
       );
     } else {
-      await client().database(databaseId).container(collectionId).scripts.trigger(triggerId).delete();
+      await client(ClientOperationType.WRITE)
+        .database(databaseId)
+        .container(collectionId)
+        .scripts.trigger(triggerId)
+        .delete();
     }
   } catch (error) {
     handleError(error, "DeleteTrigger", `Error while deleting trigger ${triggerId}`);

src/Common/dataAccess/deleteUserDefinedFunction.ts

@@ -2,7 +2,7 @@ import { AuthType } from "../../AuthType";
 import { userContext } from "../../UserContext";
 import { deleteSqlUserDefinedFunction } from "../../Utils/arm/generatedClients/cosmos/sqlResources";
 import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-import { client } from "../CosmosClient";
+import { client, ClientOperationType } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function deleteUserDefinedFunction(databaseId: string, collectionId: string, id: string): Promise<void> {
@@ -22,7 +22,11 @@ export async function deleteUserDefinedFunction(databaseId: string, collectionId
         id,
       );
     } else {
-      await client().database(databaseId).container(collectionId).scripts.userDefinedFunction(id).delete();
+      await client(ClientOperationType.WRITE)
+        .database(databaseId)
+        .container(collectionId)
+        .scripts.userDefinedFunction(id)
+        .delete();
     }
   } catch (error) {
     handleError(error, "DeleteUserDefinedFunction", `Error while deleting user defined function ${id}`);

src/Common/dataAccess/executeStoredProcedure.ts

@@ -1,9 +1,9 @@
 import { Collection } from "../../Contracts/ViewModels";
-import { ClientDefaults, HttpHeaders } from "../Constants";
-import { client } from "../CosmosClient";
-import { handleError } from "../ErrorHandlingUtils";
-import { logConsoleProgress, logConsoleInfo } from "../../Utils/NotificationConsoleUtils";
 import StoredProcedure from "../../Explorer/Tree/StoredProcedure";
+import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
+import { ClientDefaults, HttpHeaders } from "../Constants";
+import { ClientOperationType, client } from "../CosmosClient";
+import { handleError } from "../ErrorHandlingUtils";
 
 export interface ExecuteSprocResult {
   result: StoredProcedure;
@@ -22,7 +22,7 @@ export const executeStoredProcedure = async (
   }, ClientDefaults.requestTimeoutMs);
 
   try {
-    const response = await client()
+    const response = await client(ClientOperationType.WRITE)
       .database(collection.databaseId)
       .container(collection.id())
       .scripts.storedProcedure(storedProcedure.id())

src/Common/dataAccess/getIndexTransformationProgress.ts

@@ -1,9 +1,9 @@
-import { client } from "../CosmosClient";
-import { handleError } from "../ErrorHandlingUtils";
-import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-import * as Constants from "../Constants";
 import { AuthType } from "../../AuthType";
 import { userContext } from "../../UserContext";
+import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
+import * as Constants from "../Constants";
+import { ClientOperationType, client } from "../CosmosClient";
+import { handleError } from "../ErrorHandlingUtils";
 
 export async function getIndexTransformationProgress(databaseId: string, collectionId: string): Promise<number> {
   if (userContext.authType !== AuthType.AAD) {
@@ -12,7 +12,10 @@ export async function getIndexTransformationProgress(databaseId: string, collect
   let indexTransformationPercentage: number;
   const clearMessage = logConsoleProgress(`Reading container ${collectionId}`);
   try {
-    const response = await client().database(databaseId).container(collectionId).read({ populateQuotaInfo: true });
+    const response = await client(ClientOperationType.READ)
+      .database(databaseId)
+      .container(collectionId)
+      .read({ populateQuotaInfo: true });
 
     indexTransformationPercentage = parseInt(
       response.headers[Constants.HttpHeaders.collectionIndexTransformationProgress] as string,

src/Common/dataAccess/queryConflicts.ts

@@ -1,5 +1,5 @@
 import { ConflictDefinition, FeedOptions, QueryIterator, Resource } from "@azure/cosmos";
-import { client } from "../CosmosClient";
+import { ClientOperationType, client } from "../CosmosClient";
 
 export const queryConflicts = (
   databaseId: string,
@@ -7,5 +7,5 @@ export const queryConflicts = (
   query: string,
   options: FeedOptions,
 ): QueryIterator<ConflictDefinition & Resource> => {
-  return client().database(databaseId).container(containerId).conflicts.query(query, options);
+  return client(ClientOperationType.READ).database(databaseId).container(containerId).conflicts.query(query, options);
 };

src/Common/dataAccess/queryDocuments.ts

@@ -1,8 +1,7 @@
 import { FeedOptions, ItemDefinition, QueryIterator, Resource } from "@azure/cosmos";
 import { LocalStorageUtility, StorageKey } from "../../Shared/StorageUtility";
 import { Queries } from "../Constants";
-// import { client2 } from "../ReadRegionCosmosClient";
+import { ClientOperationType, client } from "../CosmosClient";
-import { client } from "../CosmosClient";
 
 export const queryDocuments = (
   databaseId: string,
@@ -11,9 +10,7 @@ export const queryDocuments = (
   options: FeedOptions,
 ): QueryIterator<ItemDefinition & Resource> => {
   options = getCommonQueryOptions(options);
-  // console.log(`${JSON.stringify(client2().getReadEndpoint())}`);
+  return client(ClientOperationType.READ).database(databaseId).container(containerId).items.query(query, options);
-  // return client2().database(databaseId).container(containerId).items.query(query, options);
-  return client().database(databaseId).container(containerId).items.query(query, options);
 };
 
 export const getCommonQueryOptions = (options: FeedOptions): FeedOptions => {

src/Common/dataAccess/readCollection.ts

@@ -3,11 +3,11 @@ import { sampleDataClient } from "Common/SampleDataClient";
 import { userContext } from "UserContext";
 import * as DataModels from "../../Contracts/DataModels";
 import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-import { client } from "../CosmosClient";
+import { ClientOperationType, client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function readCollection(databaseId: string, collectionId: string): Promise<DataModels.Collection> {
-  const cosmosClient = client();
+  const cosmosClient = client(ClientOperationType.READ);
   return await readCollectionInternal(cosmosClient, databaseId, collectionId);
 }
 

src/Common/dataAccess/readCollections.ts

@@ -10,7 +10,7 @@ import { listGremlinGraphs } from "../../Utils/arm/generatedClients/cosmos/greml
 import { listMongoDBCollections } from "../../Utils/arm/generatedClients/cosmos/mongoDBResources";
 import { listSqlContainers } from "../../Utils/arm/generatedClients/cosmos/sqlResources";
 import { listTables } from "../../Utils/arm/generatedClients/cosmos/tableResources";
-import { client } from "../CosmosClient";
+import { ClientOperationType, client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function readCollections(databaseId: string): Promise<DataModels.Collection[]> {
@@ -31,7 +31,7 @@ export async function readCollections(databaseId: string): Promise<DataModels.Co
       const tokenCollectionId = resourceIdObj[3];
 
       if (tokenDatabaseId === databaseId) {
-        promises.push(client().database(databaseId).container(tokenCollectionId).read());
+        promises.push(client(ClientOperationType.READ).database(databaseId).container(tokenCollectionId).read());
       }
     }
 
@@ -61,7 +61,7 @@ export async function readCollections(databaseId: string): Promise<DataModels.Co
       return await readCollectionsWithARM(databaseId);
     }
 
-    const sdkResponse = await client().database(databaseId).containers.readAll().fetchAll();
+    const sdkResponse = await client(ClientOperationType.READ).database(databaseId).containers.readAll().fetchAll();
     return sdkResponse.resources as DataModels.Collection[];
   } catch (error) {
     handleError(error, "ReadCollections", `Error while querying containers for database ${databaseId}`);
@@ -77,7 +77,7 @@ export async function readCollectionsWithPagination(
 ): Promise<DataModels.CollectionsWithPagination> {
   const clearMessage = logConsoleProgress(`Querying containers for database ${databaseId}`);
   try {
-    const sdkResponse = await client()
+    const sdkResponse = await client(ClientOperationType.READ)
       .database(databaseId)
       .containers.query(
         { query: "SELECT * FROM c" },

src/Common/dataAccess/readDatabases.ts

@@ -7,7 +7,7 @@ import { listCassandraKeyspaces } from "../../Utils/arm/generatedClients/cosmos/
 import { listGremlinDatabases } from "../../Utils/arm/generatedClients/cosmos/gremlinResources";
 import { listMongoDBDatabases } from "../../Utils/arm/generatedClients/cosmos/mongoDBResources";
 import { listSqlDatabases } from "../../Utils/arm/generatedClients/cosmos/sqlResources";
-import { client } from "../CosmosClient";
+import { ClientOperationType, client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function readDatabases(): Promise<DataModels.Database[]> {
@@ -56,7 +56,7 @@ export async function readDatabases(): Promise<DataModels.Database[]> {
     ) {
       databases = await readDatabasesWithARM();
     } else {
-      const sdkResponse = await client().databases.readAll().fetchAll();
+      const sdkResponse = await client(ClientOperationType.READ).databases.readAll().fetchAll();
       databases = sdkResponse.resources as DataModels.Database[];
     }
   } catch (error) {

src/Common/dataAccess/readDocument.ts

@@ -3,10 +3,9 @@ import { CollectionBase } from "../../Contracts/ViewModels";
 import DocumentId from "../../Explorer/Tree/DocumentId";
 import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { HttpHeaders } from "../Constants";
-import { client } from "../CosmosClient";
+import { ClientOperationType, client } from "../CosmosClient";
 import { getEntityName } from "../DocumentUtility";
 import { handleError } from "../ErrorHandlingUtils";
-// import { client2 } from "../ReadRegionCosmosClient";
 import { getPartitionKeyValue } from "./getPartitionKeyValue";
 
 export const readDocument = async (collection: CollectionBase, documentId: DocumentId): Promise<Item> => {
@@ -20,18 +19,12 @@ export const readDocument = async (collection: CollectionBase, documentId: Docum
             [HttpHeaders.partitionKey]: documentId.partitionKeyValue,
           }
         : {};
-    const response = await client()
+    const response = await client(ClientOperationType.READ)
       .database(collection.databaseId)
       .container(collection.id())
       .item(documentId.id(), getPartitionKeyValue(documentId))
       .read(options);
 
-    // const response = await client2()
-    //   .database(collection.databaseId)
-    //   .container(collection.id())
-    //   .item(documentId.id(), getPartitionKeyValue(documentId))
-    //   .read(options);
-
     return response?.resource;
   } catch (error) {
     handleError(error, "ReadDocument", `Failed to read ${entityName} ${documentId.id()}`);

src/Common/dataAccess/readOfferWithSDK.ts

@@ -1,7 +1,7 @@
 import { RequestOptions } from "@azure/cosmos";
 import { Offer } from "../../Contracts/DataModels";
 import { HttpHeaders } from "../Constants";
-import { client } from "../CosmosClient";
+import { ClientOperationType, client } from "../CosmosClient";
 import { parseSDKOfferResponse } from "../OfferUtility";
 import { readOffers } from "./readOffers";
 
@@ -21,7 +21,7 @@ export const readOfferWithSDK = async (offerId: string, resourceId: string): Pro
       [HttpHeaders.populateCollectionThroughputInfo]: true,
     },
   };
-  const response = await client().offer(offerId).read(options);
+  const response = await client(ClientOperationType.READ).offer(offerId).read(options);
 
   return parseSDKOfferResponse(response);
 };

src/Common/dataAccess/readOffers.ts

@@ -1,13 +1,13 @@
 import { SDKOfferDefinition } from "../../Contracts/DataModels";
 import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-import { client } from "../CosmosClient";
+import { ClientOperationType, client } from "../CosmosClient";
-import { handleError, getErrorMessage } from "../ErrorHandlingUtils";
+import { getErrorMessage, handleError } from "../ErrorHandlingUtils";
 
 export const readOffers = async (): Promise<SDKOfferDefinition[]> => {
   const clearMessage = logConsoleProgress(`Querying offers`);
 
   try {
-    const response = await client().offers.readAll().fetchAll();
+    const response = await client(ClientOperationType.READ).offers.readAll().fetchAll();
     return response?.resources;
   } catch (error) {
     // This should be removed when we can correctly identify if an account is serverless when connected using connection string too.

src/Common/dataAccess/readStoredProcedures.ts

@@ -4,7 +4,7 @@ import { AuthType } from "../../AuthType";
 import { userContext } from "../../UserContext";
 import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { listSqlStoredProcedures } from "../../Utils/arm/generatedClients/cosmos/sqlResources";
-import { client } from "../CosmosClient";
+import { ClientOperationType, client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function readStoredProcedures(
@@ -34,7 +34,7 @@ export async function readStoredProcedures(
       throw new Error(cloudError?.error?.message);
     }
 
-    const response = await client()
+    const response = await client(ClientOperationType.READ)
       .database(databaseId)
       .container(collectionId)
       .scripts.storedProcedures.readAll()

src/Common/dataAccess/readTriggers.ts

@@ -1,10 +1,10 @@
 import { TriggerDefinition } from "@azure/cosmos";
 import { AuthType } from "../../AuthType";
 import { userContext } from "../../UserContext";
+import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { listSqlTriggers } from "../../Utils/arm/generatedClients/cosmos/sqlResources";
 import { SqlTriggerResource } from "../../Utils/arm/generatedClients/cosmos/types";
-import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
+import { ClientOperationType, client } from "../CosmosClient";
-import { client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function readTriggers(
@@ -28,7 +28,11 @@ export async function readTriggers(
       return rpResponse?.value?.map((trigger) => trigger.properties?.resource);
     }
 
-    const response = await client().database(databaseId).container(collectionId).scripts.triggers.readAll().fetchAll();
+    const response = await client(ClientOperationType.READ)
+      .database(databaseId)
+      .container(collectionId)
+      .scripts.triggers.readAll()
+      .fetchAll();
     return response?.resources;
   } catch (error) {
     handleError(error, "ReadTriggers", `Failed to query triggers for container ${collectionId}`);

src/Common/dataAccess/readUserDefinedFunctions.ts

@@ -1,9 +1,9 @@
 import { Resource, UserDefinedFunctionDefinition } from "@azure/cosmos";
 import { AuthType } from "../../AuthType";
 import { userContext } from "../../UserContext";
-import { listSqlUserDefinedFunctions } from "../../Utils/arm/generatedClients/cosmos/sqlResources";
 import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-import { client } from "../CosmosClient";
+import { listSqlUserDefinedFunctions } from "../../Utils/arm/generatedClients/cosmos/sqlResources";
+import { ClientOperationType, client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function readUserDefinedFunctions(
@@ -24,7 +24,7 @@ export async function readUserDefinedFunctions(
       return rpResponse?.value?.map((udf) => udf.properties?.resource as UserDefinedFunctionDefinition & Resource);
     }
 
-    const response = await client()
+    const response = await client(ClientOperationType.READ)
       .database(databaseId)
       .container(collectionId)
       .scripts.userDefinedFunctions.readAll()

src/Common/dataAccess/updateCollection.ts

@@ -2,6 +2,7 @@ import { ContainerDefinition, RequestOptions } from "@azure/cosmos";
 import { AuthType } from "../../AuthType";
 import { Collection } from "../../Contracts/DataModels";
 import { userContext } from "../../UserContext";
+import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import {
   createUpdateCassandraTable,
   getCassandraTable,
@@ -19,8 +20,7 @@ import {
   SqlContainerCreateUpdateParameters,
   SqlContainerResource,
 } from "../../Utils/arm/generatedClients/cosmos/types";
-import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
+import { ClientOperationType, client } from "../CosmosClient";
-import { client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function updateCollection(
@@ -40,7 +40,7 @@ export async function updateCollection(
     ) {
       collection = await updateCollectionWithARM(databaseId, collectionId, newCollection);
     } else {
-      const sdkResponse = await client()
+      const sdkResponse = await client(ClientOperationType.WRITE)
         .database(databaseId)
         .container(collectionId)
         .replace(newCollection as ContainerDefinition, options);

src/Common/dataAccess/updateDocument.ts

@@ -3,7 +3,7 @@ import { HttpHeaders } from "Common/Constants";
 import { CollectionBase } from "../../Contracts/ViewModels";
 import DocumentId from "../../Explorer/Tree/DocumentId";
 import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-import { client } from "../CosmosClient";
+import { ClientOperationType, client } from "../CosmosClient";
 import { getEntityName } from "../DocumentUtility";
 import { handleError } from "../ErrorHandlingUtils";
 import { getPartitionKeyValue } from "./getPartitionKeyValue";
@@ -23,7 +23,7 @@ export const updateDocument = async (
             [HttpHeaders.partitionKey]: documentId.partitionKeyValue,
           }
         : {};
-    const response = await client()
+    const response = await client(ClientOperationType.WRITE)
       .database(collection.databaseId)
       .container(collection.id())
       .item(documentId.id(), getPartitionKeyValue(documentId))

src/Common/dataAccess/updateOffer.ts

@@ -2,6 +2,7 @@ import { OfferDefinition, RequestOptions } from "@azure/cosmos";
 import { AuthType } from "../../AuthType";
 import { Offer, SDKOfferDefinition, UpdateOfferParams } from "../../Contracts/DataModels";
 import { userContext } from "../../UserContext";
+import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import {
   migrateCassandraKeyspaceToAutoscale,
   migrateCassandraKeyspaceToManualThroughput,
@@ -40,9 +41,8 @@ import {
   updateTableThroughput,
 } from "../../Utils/arm/generatedClients/cosmos/tableResources";
 import { ThroughputSettingsUpdateParameters } from "../../Utils/arm/generatedClients/cosmos/types";
-import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { HttpHeaders } from "../Constants";
-import { client } from "../CosmosClient";
+import { ClientOperationType, client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 import { parseSDKOfferResponse } from "../OfferUtility";
 import { readCollectionOffer } from "./readCollectionOffer";
@@ -401,7 +401,7 @@ const updateOfferWithSDK = async (params: UpdateOfferParams): Promise<Offer> =>
     newOffer.content.offerAutopilotSettings = { maxThroughput: 0 };
   }
 
-  const sdkResponse = await client()
+  const sdkResponse = await client(ClientOperationType.WRITE)
     .offer(params.currentOffer.id)
     // TODO Remove casting when SDK types are fixed (https://github.com/Azure/azure-sdk-for-js/issues/10660)
     .replace(newOffer as unknown as OfferDefinition, options);

src/Common/dataAccess/updateStoredProcedure.ts

@@ -1,6 +1,7 @@
 import { Resource, StoredProcedureDefinition } from "@azure/cosmos";
 import { AuthType } from "../../AuthType";
 import { userContext } from "../../UserContext";
+import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import {
   createUpdateSqlStoredProcedure,
   getSqlStoredProcedure,
@@ -9,8 +10,7 @@ import {
   SqlStoredProcedureCreateUpdateParameters,
   SqlStoredProcedureResource,
 } from "../../Utils/arm/generatedClients/cosmos/types";
-import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
+import { ClientOperationType, client } from "../CosmosClient";
-import { client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function updateStoredProcedure(
@@ -54,7 +54,7 @@ export async function updateStoredProcedure(
       throw new Error(`Failed to update stored procedure: ${storedProcedure.id} does not exist.`);
     }
 
-    const response = await client()
+    const response = await client(ClientOperationType.WRITE)
       .database(databaseId)
       .container(collectionId)
       .scripts.storedProcedure(storedProcedure.id)

src/Common/dataAccess/updateTrigger.ts

@@ -1,10 +1,10 @@
 import { TriggerDefinition } from "@azure/cosmos";
 import { AuthType } from "../../AuthType";
 import { userContext } from "../../UserContext";
+import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { createUpdateSqlTrigger, getSqlTrigger } from "../../Utils/arm/generatedClients/cosmos/sqlResources";
 import { SqlTriggerCreateUpdateParameters, SqlTriggerResource } from "../../Utils/arm/generatedClients/cosmos/types";
-import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
+import { ClientOperationType, client } from "../CosmosClient";
-import { client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function updateTrigger(
@@ -47,7 +47,7 @@ export async function updateTrigger(
       throw new Error(`Failed to update trigger: ${trigger.id} does not exist.`);
     }
 
-    const response = await client()
+    const response = await client(ClientOperationType.WRITE)
       .database(databaseId)
       .container(collectionId)
       .scripts.trigger(trigger.id)

src/Common/dataAccess/updateUserDefinedFunction.ts

@@ -1,6 +1,7 @@
 import { Resource, UserDefinedFunctionDefinition } from "@azure/cosmos";
 import { AuthType } from "../../AuthType";
 import { userContext } from "../../UserContext";
+import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import {
   createUpdateSqlUserDefinedFunction,
   getSqlUserDefinedFunction,
@@ -9,8 +10,7 @@ import {
   SqlUserDefinedFunctionCreateUpdateParameters,
   SqlUserDefinedFunctionResource,
 } from "../../Utils/arm/generatedClients/cosmos/types";
-import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
+import { ClientOperationType, client } from "../CosmosClient";
-import { client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
 
 export async function updateUserDefinedFunction(
@@ -53,7 +53,7 @@ export async function updateUserDefinedFunction(
       throw new Error(`Failed to update user defined function: ${userDefinedFunction.id} does not exist.`);
     }
 
-    const response = await client()
+    const response = await client(ClientOperationType.WRITE)
       .database(databaseId)
       .container(collectionId)
       .scripts.userDefinedFunction(userDefinedFunction.id)