import * as _ from "underscore"; import * as Constants from "./Constants"; import * as DataModels from "../Contracts/DataModels"; import * as HeadersUtility from "./HeadersUtility"; import * as ViewModels from "../Contracts/ViewModels"; import Q from "q"; import { ConflictDefinition, ContainerDefinition, ContainerResponse, DatabaseResponse, FeedOptions, ItemDefinition, PartitionKeyDefinition, QueryIterator, Resource, TriggerDefinition } from "@azure/cosmos"; import { ContainerRequest } from "@azure/cosmos/dist-esm/client/Container/ContainerRequest"; import { client } from "./CosmosClient"; import { DatabaseRequest } from "@azure/cosmos/dist-esm/client/Database/DatabaseRequest"; import { LocalStorageUtility, StorageKey } from "../Shared/StorageUtility"; import { sendCachedDataMessage } from "./MessageHandler"; import { MessageTypes } from "../Contracts/ExplorerContracts"; import { OfferUtils } from "../Utils/OfferUtils"; import { RequestOptions } from "@azure/cosmos/dist-esm"; import StoredProcedure from "../Explorer/Tree/StoredProcedure"; import { Platform, configContext } from "../ConfigContext"; import DocumentId from "../Explorer/Tree/DocumentId"; import ConflictId from "../Explorer/Tree/ConflictId"; export function getCommonQueryOptions(options: FeedOptions): any { const storedItemPerPageSetting: number = LocalStorageUtility.getEntryNumber(StorageKey.ActualItemPerPage); options = options || {}; options.populateQueryMetrics = true; options.enableScanInQuery = options.enableScanInQuery || true; if (!options.partitionKey) { options.forceQueryPlan = true; } options.maxItemCount = options.maxItemCount || (storedItemPerPageSetting !== undefined && storedItemPerPageSetting) || Constants.Queries.itemsPerPage; options.maxDegreeOfParallelism = LocalStorageUtility.getEntryNumber(StorageKey.MaxDegreeOfParellism); return options; } export function queryDocuments( databaseId: string, containerId: string, query: string, options: any ): Q.Promise> { options = getCommonQueryOptions(options); const documentsIterator = client() .database(databaseId) .container(containerId) .items.query(query, options); return Q(documentsIterator); } export function readStoredProcedures( collection: ViewModels.Collection, options?: any ): Q.Promise { return Q( client() .database(collection.databaseId) .container(collection.id()) .scripts.storedProcedures.readAll(options) .fetchAll() .then(response => response.resources as DataModels.StoredProcedure[]) ); } export function readStoredProcedure( collection: ViewModels.Collection, requestedResource: DataModels.Resource, options?: any ): Q.Promise { return Q( client() .database(collection.databaseId) .container(collection.id()) .scripts.storedProcedure(requestedResource.id) .read(options) .then(response => response.resource as DataModels.StoredProcedure) ); } export function readUserDefinedFunctions( collection: ViewModels.Collection, options: any ): Q.Promise { return Q( client() .database(collection.databaseId) .container(collection.id()) .scripts.userDefinedFunctions.readAll(options) .fetchAll() .then(response => response.resources as DataModels.UserDefinedFunction[]) ); } export function readUserDefinedFunction( collection: ViewModels.Collection, requestedResource: DataModels.Resource, options?: any ): Q.Promise { return Q( client() .database(collection.databaseId) .container(collection.id()) .scripts.userDefinedFunction(requestedResource.id) .read(options) .then(response => response.resource as DataModels.UserDefinedFunction) ); } export function readTriggers(collection: ViewModels.Collection, options: any): Q.Promise { return Q( client() .database(collection.databaseId) .container(collection.id()) .scripts.triggers.readAll(options) .fetchAll() .then(response => response.resources as DataModels.Trigger[]) ); } export function readTrigger( collection: ViewModels.Collection, requestedResource: DataModels.Resource, options?: any ): Q.Promise { return Q( client() .database(collection.databaseId) .container(collection.id()) .scripts.trigger(requestedResource.id) .read(options) .then(response => response.resource as DataModels.Trigger) ); } export function executeStoredProcedure( collection: ViewModels.Collection, storedProcedure: StoredProcedure, partitionKeyValue: any, params: any[] ): Q.Promise { // TODO remove this deferred. Kept it because of timeout code at bottom of function const deferred = Q.defer(); client() .database(collection.databaseId) .container(collection.id()) .scripts.storedProcedure(storedProcedure.id()) .execute(partitionKeyValue, params, { enableScriptLogging: true }) .then(response => deferred.resolve({ result: response.resource, scriptLogs: response.headers[Constants.HttpHeaders.scriptLogResults] }) ) .catch(error => deferred.reject(error)); return deferred.promise.timeout( Constants.ClientDefaults.requestTimeoutMs, `Request timed out while executing stored procedure ${storedProcedure.id()}` ); } export function readDocument(collection: ViewModels.CollectionBase, documentId: DocumentId): Q.Promise { const partitionKey = documentId.partitionKeyValue; return Q( client() .database(collection.databaseId) .container(collection.id()) .item(documentId.id(), partitionKey) .read() .then(response => response.resource) ); } export function getPartitionKeyHeaderForConflict(conflictId: ConflictId): Object { const partitionKeyDefinition: DataModels.PartitionKey = conflictId.partitionKey; const partitionKeyValue: any = conflictId.partitionKeyValue; return getPartitionKeyHeader(partitionKeyDefinition, partitionKeyValue); } export function getPartitionKeyHeader(partitionKeyDefinition: DataModels.PartitionKey, partitionKeyValue: any): Object { if (!partitionKeyDefinition) { return undefined; } if (partitionKeyValue === undefined) { return [{}]; } return [partitionKeyValue]; } export function updateCollection( databaseId: string, collectionId: string, newCollection: DataModels.Collection, options: any = {} ): Q.Promise { return Q( client() .database(databaseId) .container(collectionId) .replace(newCollection as ContainerDefinition, options) .then(async (response: ContainerResponse) => { return refreshCachedResources().then(() => response.resource as DataModels.Collection); }) ); } export function updateDocument( collection: ViewModels.CollectionBase, documentId: DocumentId, newDocument: any ): Q.Promise { const partitionKey = documentId.partitionKeyValue; return Q( client() .database(collection.databaseId) .container(collection.id()) .item(documentId.id(), partitionKey) .replace(newDocument) .then(response => response.resource) ); } export function updateOffer( offer: DataModels.Offer, newOffer: DataModels.Offer, options?: RequestOptions ): Q.Promise { return Q( client() .offer(offer.id) .replace(newOffer, options) .then(response => { return Promise.all([refreshCachedOffers(), refreshCachedResources()]).then(() => response.resource); }) ); } export function updateStoredProcedure( collection: ViewModels.Collection, storedProcedure: DataModels.StoredProcedure, options: any ): Q.Promise { return Q( client() .database(collection.databaseId) .container(collection.id()) .scripts.storedProcedure(storedProcedure.id) .replace(storedProcedure, options) .then(response => response.resource as DataModels.StoredProcedure) ); } export function updateUserDefinedFunction( collection: ViewModels.Collection, userDefinedFunction: DataModels.UserDefinedFunction, options?: any ): Q.Promise { return Q( client() .database(collection.databaseId) .container(collection.id()) .scripts.userDefinedFunction(userDefinedFunction.id) .replace(userDefinedFunction, options) .then(response => response.resource as DataModels.StoredProcedure) ); } export function updateTrigger( collection: ViewModels.Collection, trigger: DataModels.Trigger, options?: any ): Q.Promise { return Q( client() .database(collection.databaseId) .container(collection.id()) .scripts.trigger(trigger.id) .replace(trigger as TriggerDefinition, options) .then(response => response.resource as DataModels.Trigger) ); } export function createDocument(collection: ViewModels.CollectionBase, newDocument: any): Q.Promise { return Q( client() .database(collection.databaseId) .container(collection.id()) .items.create(newDocument) .then(response => response.resource as DataModels.StoredProcedure) ); } export function createStoredProcedure( collection: ViewModels.Collection, newStoredProcedure: DataModels.StoredProcedure, options?: any ): Q.Promise { return Q( client() .database(collection.databaseId) .container(collection.id()) .scripts.storedProcedures.create(newStoredProcedure, options) .then(response => response.resource as DataModels.StoredProcedure) ); } export function createUserDefinedFunction( collection: ViewModels.Collection, newUserDefinedFunction: DataModels.UserDefinedFunction, options: any ): Q.Promise { return Q( client() .database(collection.databaseId) .container(collection.id()) .scripts.userDefinedFunctions.create(newUserDefinedFunction, options) .then(response => response.resource as DataModels.UserDefinedFunction) ); } export function createTrigger( collection: ViewModels.Collection, newTrigger: DataModels.Trigger, options?: any ): Q.Promise { return Q( client() .database(collection.databaseId) .container(collection.id()) .scripts.triggers.create(newTrigger as TriggerDefinition, options) .then(response => response.resource as DataModels.Trigger) ); } export function deleteDocument(collection: ViewModels.CollectionBase, documentId: DocumentId): Q.Promise { const partitionKey = documentId.partitionKeyValue; return Q( client() .database(collection.databaseId) .container(collection.id()) .item(documentId.id(), partitionKey) .delete() ); } export function deleteConflict( collection: ViewModels.CollectionBase, conflictId: ConflictId, options: any = {} ): Q.Promise { options.partitionKey = options.partitionKey || getPartitionKeyHeaderForConflict(conflictId); return Q( client() .database(collection.databaseId) .container(collection.id()) .conflict(conflictId.id()) .delete(options) ); } export function deleteStoredProcedure( collection: ViewModels.Collection, storedProcedure: DataModels.StoredProcedure, options: any ): Q.Promise { return Q( client() .database(collection.databaseId) .container(collection.id()) .scripts.storedProcedure(storedProcedure.id) .delete() ); } export function deleteUserDefinedFunction( collection: ViewModels.Collection, userDefinedFunction: DataModels.UserDefinedFunction, options: any ): Q.Promise { return Q( client() .database(collection.databaseId) .container(collection.id()) .scripts.userDefinedFunction(userDefinedFunction.id) .delete() ); } export function deleteTrigger( collection: ViewModels.Collection, trigger: DataModels.Trigger, options: any ): Q.Promise { return Q( client() .database(collection.databaseId) .container(collection.id()) .scripts.trigger(trigger.id) .delete() ); } export function readCollectionQuotaInfo( collection: ViewModels.Collection, options: any ): Q.Promise { options = options || {}; options.populateQuotaInfo = true; options.initialHeaders = options.initialHeaders || {}; options.initialHeaders[Constants.HttpHeaders.populatePartitionStatistics] = true; return Q( client() .database(collection.databaseId) .container(collection.id()) .read(options) // TODO any needed because SDK does not properly type response.resource.statistics .then((response: any) => { let quota: DataModels.CollectionQuotaInfo = HeadersUtility.getQuota(response.headers); quota["usageSizeInKB"] = response.resource.statistics.reduce( ( previousValue: number, currentValue: DataModels.Statistic, currentIndex: number, array: DataModels.Statistic[] ) => { return previousValue + currentValue.sizeInKB; }, 0 ); quota["numPartitions"] = response.resource.statistics.length; quota["uniqueKeyPolicy"] = collection.uniqueKeyPolicy; // TODO: Remove after refactoring (#119617) return quota; }) ); } export function readOffers(options: any): Q.Promise { try { if (configContext.platform === Platform.Portal) { return sendCachedDataMessage(MessageTypes.AllOffers, [ (window).dataExplorer.databaseAccount().id, Constants.ClientDefaults.portalCacheTimeoutMs ]); } } catch (error) { // If error getting cached Offers, continue on and read via SDK } return Q( client() .offers.readAll() .fetchAll() .then(response => response.resources) ); } export function readOffer(requestedResource: DataModels.Offer, options: any): Q.Promise { options = options || {}; options.initialHeaders = options.initialHeaders || {}; if (!OfferUtils.isOfferV1(requestedResource)) { options.initialHeaders[Constants.HttpHeaders.populateCollectionThroughputInfo] = true; } return Q( client() .offer(requestedResource.id) .read(options) .then(response => ({ ...response.resource, headers: response.headers })) ); } export function getOrCreateDatabaseAndCollection( request: DataModels.CreateDatabaseAndCollectionRequest, options: any ): Q.Promise { const databaseOptions: any = options && _.omit(options, "sharedOfferThroughput"); const { databaseId, databaseLevelThroughput, collectionId, partitionKey, indexingPolicy, uniqueKeyPolicy, offerThroughput, analyticalStorageTtl, hasAutoPilotV2FeatureFlag } = request; const createBody: DatabaseRequest = { id: databaseId }; // TODO: replace when SDK support autopilot const initialHeaders = request.autoPilot ? !hasAutoPilotV2FeatureFlag ? { [Constants.HttpHeaders.autoPilotThroughputSDK]: JSON.stringify({ maxThroughput: request.autoPilot.maxThroughput }) } : { [Constants.HttpHeaders.autoPilotTier]: request.autoPilot.autopilotTier } : undefined; if (databaseLevelThroughput) { if (request.autoPilot) { databaseOptions.initialHeaders = initialHeaders; } createBody.throughput = offerThroughput; } return Q( client() .databases.createIfNotExists(createBody, databaseOptions) .then(response => { return response.database.containers.create( { id: collectionId, partitionKey: (partitionKey || undefined) as PartitionKeyDefinition, indexingPolicy: indexingPolicy ? indexingPolicy : undefined, uniqueKeyPolicy: uniqueKeyPolicy ? uniqueKeyPolicy : undefined, analyticalStorageTtl: analyticalStorageTtl, throughput: databaseLevelThroughput || request.autoPilot ? undefined : offerThroughput } as ContainerRequest, // TODO: remove cast when https://github.com/Azure/azure-cosmos-js/issues/423 is fixed { initialHeaders: databaseLevelThroughput ? undefined : initialHeaders } ); }) .then(containerResponse => containerResponse.resource as DataModels.Collection) .finally(() => refreshCachedResources(options)) ); } export function createDatabase( request: DataModels.CreateDatabaseRequest, options: any ): Q.Promise { var deferred = Q.defer(); _createDatabase(request, options).then( (createdDatabase: DataModels.Database) => { refreshCachedOffers().then(() => { deferred.resolve(createdDatabase); }); }, _createDatabaseError => { deferred.reject(_createDatabaseError); } ); return deferred.promise; } export function refreshCachedOffers(): Q.Promise { if (configContext.platform === Platform.Portal) { return sendCachedDataMessage(MessageTypes.RefreshOffers, []); } else { return Q(); } } export function refreshCachedResources(options?: any): Q.Promise { if (configContext.platform === Platform.Portal) { return sendCachedDataMessage(MessageTypes.RefreshResources, []); } else { return Q(); } } export function queryConflicts( databaseId: string, containerId: string, query: string, options: any ): Q.Promise> { const documentsIterator = client() .database(databaseId) .container(containerId) .conflicts.query(query, options); return Q(documentsIterator); } function _createDatabase(request: DataModels.CreateDatabaseRequest, options: any = {}): Q.Promise { const { databaseId, databaseLevelThroughput, offerThroughput, autoPilot, hasAutoPilotV2FeatureFlag } = request; const createBody: DatabaseRequest = { id: databaseId }; const databaseOptions: any = options && _.omit(options, "sharedOfferThroughput"); // TODO: replace when SDK support autopilot const initialHeaders = autoPilot ? !hasAutoPilotV2FeatureFlag ? { [Constants.HttpHeaders.autoPilotThroughputSDK]: JSON.stringify({ maxThroughput: autoPilot.maxThroughput }) } : { [Constants.HttpHeaders.autoPilotTier]: autoPilot.autopilotTier } : undefined; if (!!databaseLevelThroughput) { if (autoPilot) { databaseOptions.initialHeaders = initialHeaders; } createBody.throughput = offerThroughput; } return Q( client() .databases.create(createBody, databaseOptions) .then((response: DatabaseResponse) => { return refreshCachedResources(databaseOptions).then(() => response.resource); }) ); }