import { Constants as CosmosSDKConstants } from "@azure/cosmos"; import { getMongoGuidRepresentation } from "Shared/StorageUtility"; import { AuthType } from "../AuthType"; import { configContext } from "../ConfigContext"; import * as DataModels from "../Contracts/DataModels"; import { MessageTypes } from "../Contracts/ExplorerContracts"; import { Collection } from "../Contracts/ViewModels"; import DocumentId from "../Explorer/Tree/DocumentId"; import { userContext } from "../UserContext"; import { logConsoleError } from "../Utils/NotificationConsoleUtils"; import { ApiType, ContentType, HttpHeaders, HttpStatusCodes } from "./Constants"; import { MinimalQueryIterator } from "./IteratorUtilities"; import { sendMessage } from "./MessageHandler"; const defaultHeaders = { [HttpHeaders.apiType]: ApiType.MongoDB.toString(), [CosmosSDKConstants.HttpHeaders.MaxEntityCount]: "100", [CosmosSDKConstants.HttpHeaders.Version]: "2017-11-15", }; function authHeaders() { if (userContext.authType === AuthType.EncryptedToken) { return { [HttpHeaders.guestAccessToken]: userContext.accessToken }; } else { return { [HttpHeaders.authorization]: userContext.authorizationToken }; } } export function queryIterator(databaseId: string, collection: Collection, query: string): MinimalQueryIterator { let continuationToken: string; return { fetchNext: () => { return queryDocuments(databaseId, collection, false, query).then((response) => { continuationToken = response.continuationToken; const headers: { [key: string]: string | number } = {}; response.headers.forEach((value, key) => { headers[key] = value; }); return { resources: response.documents, headers, requestCharge: Number(headers[CosmosSDKConstants.HttpHeaders.RequestCharge]), activityId: String(headers[CosmosSDKConstants.HttpHeaders.ActivityId]), hasMoreResults: !!continuationToken, }; }); }, }; } interface QueryResponse { continuationToken: string; documents: DataModels.DocumentId[]; headers: Headers; } export function queryDocuments( databaseId: string, collection: Collection, isResourceList: boolean, query: string, continuationToken?: string, ): Promise { const { databaseAccount } = userContext; const resourceEndpoint = databaseAccount.properties.mongoEndpoint || databaseAccount.properties.documentEndpoint; const params = { databaseID: databaseId, collectionID: collection.id(), resourceUrl: `${resourceEndpoint}dbs/${databaseId}/colls/${collection.id()}/docs/`, resourceID: collection.rid, resourceType: "docs", subscriptionID: userContext.subscriptionId, resourceGroup: userContext.resourceGroup, databaseAccountName: databaseAccount.name, partitionKey: collection && collection.partitionKey && !collection.partitionKey.systemKey ? collection.partitionKeyProperties?.[0] : "", query, }; const endpoint = getEndpoint(configContext.MONGO_PROXY_ENDPOINT) || ""; const headers = { ...defaultHeaders, ...authHeaders(), [CosmosSDKConstants.HttpHeaders.IsQuery]: "true", [CosmosSDKConstants.HttpHeaders.PopulateQueryMetrics]: "true", [CosmosSDKConstants.HttpHeaders.EnableScanInQuery]: "true", [CosmosSDKConstants.HttpHeaders.EnableCrossPartitionQuery]: "true", [CosmosSDKConstants.HttpHeaders.ParallelizeCrossPartitionQuery]: "true", [HttpHeaders.contentType]: "application/query+json", }; if (continuationToken) { headers[CosmosSDKConstants.HttpHeaders.Continuation] = continuationToken; } const path = isResourceList ? "/resourcelist" : "/queryDocuments"; return window .fetch(`${endpoint}${path}`, { method: "POST", body: JSON.stringify(params), headers, }) .then(async (response) => { if (response.ok) { return { continuationToken: response.headers.get(CosmosSDKConstants.HttpHeaders.Continuation), documents: (await response.json()).Documents as DataModels.DocumentId[], headers: response.headers, }; } await errorHandling(response, "querying documents", params); return undefined; }); } export function readDocument( databaseId: string, collection: Collection, documentId: DocumentId, ): Promise { const { databaseAccount } = userContext; const resourceEndpoint = databaseAccount.properties.mongoEndpoint || databaseAccount.properties.documentEndpoint; const idComponents = documentId.self.split("/"); const path = idComponents.slice(0, 4).join("/"); const rid = encodeURIComponent(idComponents[5]); const params = { databaseID: databaseId, collectionID: collection.id(), resourceUrl: `${resourceEndpoint}${path}/${rid}`, resourceID: rid, resourceType: "docs", subscriptionID: userContext.subscriptionId, resourceGroup: userContext.resourceGroup, databaseAccountName: databaseAccount.name, partitionKey: documentId && documentId.partitionKey && !documentId.partitionKey.systemKey ? documentId.partitionKeyProperties?.[0] : "", clientSettings: { guidRepresentation: getMongoGuidRepresentation(), }, }; const endpoint = getEndpoint(configContext.MONGO_PROXY_ENDPOINT); return window .fetch(endpoint, { method: "POST", body: JSON.stringify(params), headers: { ...defaultHeaders, ...authHeaders(), [HttpHeaders.contentType]: ContentType.applicationJson, }, }) .then(async (response) => { if (response.ok) { return response.json(); } return await errorHandling(response, "reading document", params); }); } export function createDocument( databaseId: string, collection: Collection, partitionKeyProperty: string, documentContent: unknown, ): Promise { const { databaseAccount } = userContext; const resourceEndpoint = databaseAccount.properties.mongoEndpoint || databaseAccount.properties.documentEndpoint; const params = { databaseID: databaseId, collectionID: collection.id(), resourceUrl: `${resourceEndpoint}dbs/${databaseId}/colls/${collection.id()}/docs/`, resourceID: collection.rid, resourceType: "docs", subscriptionID: userContext.subscriptionId, resourceGroup: userContext.resourceGroup, databaseAccountName: databaseAccount.name, partitionKey: collection && collection.partitionKey && !collection.partitionKey.systemKey ? partitionKeyProperty : "", documentContent: JSON.stringify(documentContent), clientSettings: { guidRepresentation: getMongoGuidRepresentation(), }, }; const endpoint = getEndpoint(configContext.MONGO_PROXY_ENDPOINT); return window .fetch(`${endpoint}/createDocument`, { method: "POST", body: JSON.stringify(params), headers: { ...defaultHeaders, ...authHeaders(), [HttpHeaders.contentType]: ContentType.applicationJson, }, }) .then(async (response) => { if (response.ok) { return response.json(); } return await errorHandling(response, "creating document", params); }); } export function updateDocument( databaseId: string, collection: Collection, documentId: DocumentId, documentContent: string, ): Promise { const { databaseAccount } = userContext; const resourceEndpoint = databaseAccount.properties.mongoEndpoint || databaseAccount.properties.documentEndpoint; const idComponents = documentId.self.split("/"); const path = idComponents.slice(0, 5).join("/"); const rid = encodeURIComponent(idComponents[5]); const params = { databaseID: databaseId, collectionID: collection.id(), resourceUrl: `${resourceEndpoint}${path}/${rid}`, resourceID: rid, resourceType: "docs", subscriptionID: userContext.subscriptionId, resourceGroup: userContext.resourceGroup, databaseAccountName: databaseAccount.name, partitionKey: documentId && documentId.partitionKey && !documentId.partitionKey.systemKey ? documentId.partitionKeyProperties?.[0] : "", documentContent, clientSettings: { guidRepresentation: getMongoGuidRepresentation(), }, }; const endpoint = getEndpoint(configContext.MONGO_PROXY_ENDPOINT); return window .fetch(endpoint, { method: "PUT", body: JSON.stringify(params), headers: { ...defaultHeaders, ...authHeaders(), [HttpHeaders.contentType]: ContentType.applicationJson, [CosmosSDKConstants.HttpHeaders.PartitionKey]: JSON.stringify(documentId.partitionKeyHeader()), }, }) .then(async (response) => { if (response.ok) { return response.json(); } return await errorHandling(response, "updating document", params); }); } export function deleteDocuments( databaseId: string, collection: Collection, documentIds: DocumentId[], ): Promise<{ deletedCount: number; isAcknowledged: boolean; }> { const { databaseAccount } = userContext; const resourceEndpoint = databaseAccount.properties.mongoEndpoint || databaseAccount.properties.documentEndpoint; const rids: string[] = documentIds.map((documentId) => { const idComponents = documentId.self.split("/"); return idComponents[5]; }); const params = { databaseID: databaseId, collectionID: collection.id(), resourceUrl: `${resourceEndpoint}`, resourceIDs: rids, subscriptionID: userContext.subscriptionId, resourceGroup: userContext.resourceGroup, databaseAccountName: databaseAccount.name, clientSettings: { guidRepresentation: getMongoGuidRepresentation(), }, }; const endpoint = getEndpoint(configContext.MONGO_PROXY_ENDPOINT); return window .fetch(`${endpoint}/bulkdelete`, { method: "DELETE", body: JSON.stringify(params), headers: { ...defaultHeaders, ...authHeaders(), [HttpHeaders.contentType]: ContentType.applicationJson, }, }) .then(async (response) => { if (response.ok) { const result = await response.json(); return result; } return await errorHandling(response, "deleting documents", params); }); } export function createMongoCollectionWithProxy( params: DataModels.CreateCollectionParams, ): Promise { const { databaseAccount } = userContext; const shardKey: string = params.partitionKey?.paths[0]; const createCollectionParams = { databaseID: params.databaseId, collectionID: params.collectionId, resourceUrl: databaseAccount.properties.mongoEndpoint || databaseAccount.properties.documentEndpoint, resourceID: "", resourceType: "colls", subscriptionID: userContext.subscriptionId, resourceGroup: userContext.resourceGroup, databaseAccountName: databaseAccount.name, partitionKey: shardKey, isAutoscale: !!params.autoPilotMaxThroughput, hasSharedThroughput: params.databaseLevelThroughput, offerThroughput: params.autoPilotMaxThroughput || params.offerThroughput, createDatabase: params.createNewDatabase, isSharded: !!shardKey, }; const endpoint = getEndpoint(configContext.MONGO_PROXY_ENDPOINT); return window .fetch(`${endpoint}/createCollection`, { method: "POST", body: JSON.stringify(createCollectionParams), headers: { ...defaultHeaders, ...authHeaders(), [HttpHeaders.contentType]: ContentType.applicationJson, }, }) .then(async (response) => { if (response.ok) { return response.json(); } return await errorHandling(response, "creating collection", createCollectionParams); }); } export function getEndpoint(endpoint: string): string { let url = endpoint + "/api/mongo/explorer"; if (userContext.authType === AuthType.EncryptedToken) { if (endpoint === configContext.MONGO_PROXY_ENDPOINT) { url = url.replace("api/mongo", "api/connectionstring/mongo"); } else { url = url.replace("api/mongo", "api/guest/mongo"); } } return url; } export class ThrottlingError extends Error { constructor(message: string) { super(message); } } // TODO: This function throws most of the time except on Forbidden which is a bit strange // It causes problems for TypeScript understanding the types async function errorHandling(response: Response, action: string, params: unknown): Promise { const errorMessage = await response.text(); // Log the error where the user can see it logConsoleError(`Error ${action}: ${errorMessage}, Payload: ${JSON.stringify(params)}`); if (response.status === HttpStatusCodes.Forbidden) { sendMessage({ type: MessageTypes.ForbiddenError, reason: errorMessage }); return; } else if ( response.status === HttpStatusCodes.BadRequest && errorMessage.includes("Error=16500") && errorMessage.includes("RetryAfterMs=") ) { // If throttling is happening, Cosmos DB will return a 400 with a body of: // A write operation resulted in an error. Error=16500, RetryAfterMs=4, Details='Batch write error. throw new ThrottlingError(errorMessage); } throw new Error(errorMessage); } export function getARMCreateCollectionEndpoint(params: DataModels.MongoParameters): string { return `subscriptions/${params.sid}/resourceGroups/${params.rg}/providers/Microsoft.DocumentDB/databaseAccounts/${userContext.databaseAccount.name}/mongodbDatabases/${params.db}/collections/${params.coll}`; }