From 56f50e13209356a8eb855966aa43088703b0cced Mon Sep 17 00:00:00 2001 From: "Craig Boger (from Dev Box)" Date: Tue, 20 Feb 2024 16:31:42 -0800 Subject: [PATCH] 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 | 132 +++++++----- src/Common/ReadRegionCosmosClient.ts | 200 ------------------ src/Common/dataAccess/bulkCreateDocument.ts | 4 +- src/Common/dataAccess/createCollection.ts | 8 +- src/Common/dataAccess/createDatabase.ts | 6 +- src/Common/dataAccess/createDocument.ts | 4 +- .../dataAccess/createStoredProcedure.ts | 6 +- src/Common/dataAccess/createTrigger.ts | 6 +- .../dataAccess/createUserDefinedFunction.ts | 6 +- src/Common/dataAccess/deleteCollection.ts | 4 +- src/Common/dataAccess/deleteConflict.ts | 10 +- src/Common/dataAccess/deleteDatabase.ts | 4 +- src/Common/dataAccess/deleteDocument.ts | 4 +- .../dataAccess/deleteStoredProcedure.ts | 8 +- src/Common/dataAccess/deleteTrigger.ts | 8 +- .../dataAccess/deleteUserDefinedFunction.ts | 8 +- .../dataAccess/executeStoredProcedure.ts | 10 +- .../getIndexTransformationProgress.ts | 13 +- src/Common/dataAccess/queryConflicts.ts | 4 +- src/Common/dataAccess/queryDocuments.ts | 7 +- src/Common/dataAccess/readCollection.ts | 4 +- src/Common/dataAccess/readCollections.ts | 8 +- src/Common/dataAccess/readDatabases.ts | 4 +- src/Common/dataAccess/readDocument.ts | 11 +- src/Common/dataAccess/readOfferWithSDK.ts | 4 +- src/Common/dataAccess/readOffers.ts | 6 +- src/Common/dataAccess/readStoredProcedures.ts | 4 +- src/Common/dataAccess/readTriggers.ts | 10 +- .../dataAccess/readUserDefinedFunctions.ts | 6 +- src/Common/dataAccess/updateCollection.ts | 6 +- src/Common/dataAccess/updateDocument.ts | 4 +- src/Common/dataAccess/updateOffer.ts | 6 +- .../dataAccess/updateStoredProcedure.ts | 6 +- src/Common/dataAccess/updateTrigger.ts | 6 +- .../dataAccess/updateUserDefinedFunction.ts | 6 +- 35 files changed, 193 insertions(+), 350 deletions(-) delete mode 100644 src/Common/ReadRegionCosmosClient.ts diff --git a/src/Common/CosmosClient.ts b/src/Common/CosmosClient.ts index e4f857b1d..777a41b44 100644 --- a/src/Common/CosmosClient.ts +++ b/src/Common/CosmosClient.ts @@ -85,10 +85,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 = 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); }; @@ -135,36 +136,89 @@ enum SDKSupportedCapabilities { PartitionMerge = 1 << 0, } -// Need to put in some kind of function here to recreate the CosmosClient with a new endpoint. -// changeClientEndpoint....... +// Client Management let _client: Cosmos.CosmosClient; -let _currentClientEndpoint: string; +let _readClients: Map = 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) { - return _client; + 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 = {}; @@ -183,11 +237,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", @@ -201,26 +252,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 }]; } @@ -231,6 +262,7 @@ export function client(): Cosmos.CosmosClient { (options as any).plugins = plugins; } - _client = new Cosmos.CosmosClient(options); - return _client; + _readClients.set(selectedEndpoint, new Cosmos.CosmosClient(options)); + + return _readClients.get(selectedEndpoint); } diff --git a/src/Common/ReadRegionCosmosClient.ts b/src/Common/ReadRegionCosmosClient.ts deleted file mode 100644 index 235fd6042..000000000 --- a/src/Common/ReadRegionCosmosClient.ts +++ /dev/null @@ -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( - 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 = 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 { - 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; -} diff --git a/src/Common/dataAccess/bulkCreateDocument.ts b/src/Common/dataAccess/bulkCreateDocument.ts index c39b37821..e041d222a 100644 --- a/src/Common/dataAccess/bulkCreateDocument.ts +++ b/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( diff --git a/src/Common/dataAccess/createCollection.ts b/src/Common/dataAccess/createCollection.ts index 73ccff3f4..2d6fb8fec 100644 --- a/src/Common/dataAccess/createCollection.ts +++ b/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 { client } from "../CosmosClient"; +import { ClientOperationType, 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, diff --git a/src/Common/dataAccess/createDatabase.ts b/src/Common/dataAccess/createDatabase.ts index 21e54b95c..e52a2b58e 100644 --- a/src/Common/dataAccess/createDatabase.ts +++ b/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 { client } from "../CosmosClient"; +import { ClientOperationType, client } from "../CosmosClient"; import { handleError } from "../ErrorHandlingUtils"; export async function createDatabase(params: DataModels.CreateDatabaseParams): Promise { @@ -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; } diff --git a/src/Common/dataAccess/createDocument.ts b/src/Common/dataAccess/createDocument.ts index 94dde951d..2763c886c 100644 --- a/src/Common/dataAccess/createDocument.ts +++ b/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); diff --git a/src/Common/dataAccess/createStoredProcedure.ts b/src/Common/dataAccess/createStoredProcedure.ts index ab7839f3e..b1951f296 100644 --- a/src/Common/dataAccess/createStoredProcedure.ts +++ b/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 { client } from "../CosmosClient"; +import { ClientOperationType, 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); diff --git a/src/Common/dataAccess/createTrigger.ts b/src/Common/dataAccess/createTrigger.ts index dd4ec1af5..4faa8ee5a 100644 --- a/src/Common/dataAccess/createTrigger.ts +++ b/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 { client } from "../CosmosClient"; +import { ClientOperationType, 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 diff --git a/src/Common/dataAccess/createUserDefinedFunction.ts b/src/Common/dataAccess/createUserDefinedFunction.ts index 3d1bca86e..b858ca09e 100644 --- a/src/Common/dataAccess/createUserDefinedFunction.ts +++ b/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 { client } from "../CosmosClient"; +import { ClientOperationType, 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); diff --git a/src/Common/dataAccess/deleteCollection.ts b/src/Common/dataAccess/deleteCollection.ts index f83126dd1..54184dbc6 100644 --- a/src/Common/dataAccess/deleteCollection.ts +++ b/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 { @@ -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) { diff --git a/src/Common/dataAccess/deleteConflict.ts b/src/Common/dataAccess/deleteConflict.ts index 5c778f474..01131e5e5 100644 --- a/src/Common/dataAccess/deleteConflict.ts +++ b/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 => { 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()) diff --git a/src/Common/dataAccess/deleteDatabase.ts b/src/Common/dataAccess/deleteDatabase.ts index fc16fe205..ba224f9c5 100644 --- a/src/Common/dataAccess/deleteDatabase.ts +++ b/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 { @@ -15,7 +15,7 @@ export async function deleteDatabase(databaseId: string): Promise { 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) { diff --git a/src/Common/dataAccess/deleteDocument.ts b/src/Common/dataAccess/deleteDocument.ts index 5caef9e0e..e5ce9b8f8 100644 --- a/src/Common/dataAccess/deleteDocument.ts +++ b/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)) diff --git a/src/Common/dataAccess/deleteStoredProcedure.ts b/src/Common/dataAccess/deleteStoredProcedure.ts index 403b707ff..adf59dad7 100644 --- a/src/Common/dataAccess/deleteStoredProcedure.ts +++ b/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}`); diff --git a/src/Common/dataAccess/deleteTrigger.ts b/src/Common/dataAccess/deleteTrigger.ts index 22b77f009..96a443b3a 100644 --- a/src/Common/dataAccess/deleteTrigger.ts +++ b/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 { @@ -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}`); diff --git a/src/Common/dataAccess/deleteUserDefinedFunction.ts b/src/Common/dataAccess/deleteUserDefinedFunction.ts index ee70b803c..7227b0bac 100644 --- a/src/Common/dataAccess/deleteUserDefinedFunction.ts +++ b/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 { @@ -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}`); diff --git a/src/Common/dataAccess/executeStoredProcedure.ts b/src/Common/dataAccess/executeStoredProcedure.ts index 7dbbcdd20..d7b30fc39 100644 --- a/src/Common/dataAccess/executeStoredProcedure.ts +++ b/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()) diff --git a/src/Common/dataAccess/getIndexTransformationProgress.ts b/src/Common/dataAccess/getIndexTransformationProgress.ts index 674a13a11..1244812d8 100644 --- a/src/Common/dataAccess/getIndexTransformationProgress.ts +++ b/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 { 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, diff --git a/src/Common/dataAccess/queryConflicts.ts b/src/Common/dataAccess/queryConflicts.ts index 5e0fdc103..bb42f9151 100644 --- a/src/Common/dataAccess/queryConflicts.ts +++ b/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 => { - return client().database(databaseId).container(containerId).conflicts.query(query, options); + return client(ClientOperationType.READ).database(databaseId).container(containerId).conflicts.query(query, options); }; diff --git a/src/Common/dataAccess/queryDocuments.ts b/src/Common/dataAccess/queryDocuments.ts index 582fa4c6d..1e577cb39 100644 --- a/src/Common/dataAccess/queryDocuments.ts +++ b/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 { client } from "../CosmosClient"; +import { ClientOperationType, client } from "../CosmosClient"; export const queryDocuments = ( databaseId: string, @@ -11,9 +10,7 @@ export const queryDocuments = ( options: FeedOptions, ): QueryIterator => { options = getCommonQueryOptions(options); - // console.log(`${JSON.stringify(client2().getReadEndpoint())}`); - // return client2().database(databaseId).container(containerId).items.query(query, options); - return client().database(databaseId).container(containerId).items.query(query, options); + return client(ClientOperationType.READ).database(databaseId).container(containerId).items.query(query, options); }; export const getCommonQueryOptions = (options: FeedOptions): FeedOptions => { diff --git a/src/Common/dataAccess/readCollection.ts b/src/Common/dataAccess/readCollection.ts index 3943b762d..8325ded25 100644 --- a/src/Common/dataAccess/readCollection.ts +++ b/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 { - const cosmosClient = client(); + const cosmosClient = client(ClientOperationType.READ); return await readCollectionInternal(cosmosClient, databaseId, collectionId); } diff --git a/src/Common/dataAccess/readCollections.ts b/src/Common/dataAccess/readCollections.ts index 4098b3fe8..18bed4409 100644 --- a/src/Common/dataAccess/readCollections.ts +++ b/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 { @@ -31,7 +31,7 @@ export async function readCollections(databaseId: string): Promise { 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" }, diff --git a/src/Common/dataAccess/readDatabases.ts b/src/Common/dataAccess/readDatabases.ts index 9cc0b0641..d281b369b 100644 --- a/src/Common/dataAccess/readDatabases.ts +++ b/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 { @@ -56,7 +56,7 @@ export async function readDatabases(): Promise { ) { 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) { diff --git a/src/Common/dataAccess/readDocument.ts b/src/Common/dataAccess/readDocument.ts index 12d3f0fcb..8cd6b47a9 100644 --- a/src/Common/dataAccess/readDocument.ts +++ b/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 => { @@ -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()}`); diff --git a/src/Common/dataAccess/readOfferWithSDK.ts b/src/Common/dataAccess/readOfferWithSDK.ts index 7332a4e4f..071f9b090 100644 --- a/src/Common/dataAccess/readOfferWithSDK.ts +++ b/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); }; diff --git a/src/Common/dataAccess/readOffers.ts b/src/Common/dataAccess/readOffers.ts index 7205e4ada..a4a328a30 100644 --- a/src/Common/dataAccess/readOffers.ts +++ b/src/Common/dataAccess/readOffers.ts @@ -1,13 +1,13 @@ import { SDKOfferDefinition } from "../../Contracts/DataModels"; import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils"; -import { client } from "../CosmosClient"; -import { handleError, getErrorMessage } from "../ErrorHandlingUtils"; +import { ClientOperationType, client } from "../CosmosClient"; +import { getErrorMessage, handleError } from "../ErrorHandlingUtils"; export const readOffers = async (): Promise => { 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. diff --git a/src/Common/dataAccess/readStoredProcedures.ts b/src/Common/dataAccess/readStoredProcedures.ts index cdb60160e..46d418c68 100644 --- a/src/Common/dataAccess/readStoredProcedures.ts +++ b/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() diff --git a/src/Common/dataAccess/readTriggers.ts b/src/Common/dataAccess/readTriggers.ts index 61b03b2e5..59f831855 100644 --- a/src/Common/dataAccess/readTriggers.ts +++ b/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 { client } from "../CosmosClient"; +import { ClientOperationType, 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}`); diff --git a/src/Common/dataAccess/readUserDefinedFunctions.ts b/src/Common/dataAccess/readUserDefinedFunctions.ts index e98ec0b50..eeaa3561d 100644 --- a/src/Common/dataAccess/readUserDefinedFunctions.ts +++ b/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() diff --git a/src/Common/dataAccess/updateCollection.ts b/src/Common/dataAccess/updateCollection.ts index 15515f5b5..4ff351091 100644 --- a/src/Common/dataAccess/updateCollection.ts +++ b/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 { client } from "../CosmosClient"; +import { ClientOperationType, 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); diff --git a/src/Common/dataAccess/updateDocument.ts b/src/Common/dataAccess/updateDocument.ts index dd97b46de..b491fedd6 100644 --- a/src/Common/dataAccess/updateDocument.ts +++ b/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)) diff --git a/src/Common/dataAccess/updateOffer.ts b/src/Common/dataAccess/updateOffer.ts index 0e507ce37..bf999198e 100644 --- a/src/Common/dataAccess/updateOffer.ts +++ b/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 => 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); diff --git a/src/Common/dataAccess/updateStoredProcedure.ts b/src/Common/dataAccess/updateStoredProcedure.ts index e641081f2..c267e5b62 100644 --- a/src/Common/dataAccess/updateStoredProcedure.ts +++ b/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 { client } from "../CosmosClient"; +import { ClientOperationType, 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) diff --git a/src/Common/dataAccess/updateTrigger.ts b/src/Common/dataAccess/updateTrigger.ts index b6681894d..8770430a6 100644 --- a/src/Common/dataAccess/updateTrigger.ts +++ b/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 { client } from "../CosmosClient"; +import { ClientOperationType, 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) diff --git a/src/Common/dataAccess/updateUserDefinedFunction.ts b/src/Common/dataAccess/updateUserDefinedFunction.ts index c8bc71b52..b75f14bc0 100644 --- a/src/Common/dataAccess/updateUserDefinedFunction.ts +++ b/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 { client } from "../CosmosClient"; +import { ClientOperationType, 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)