import { OfferDefinition, RequestOptions } from "@azure/cosmos"; import { AuthType } from "../../AuthType"; import { Offer, SDKOfferDefinition, UpdateOfferParams } from "../../Contracts/DataModels"; import { userContext } from "../../UserContext"; import { migrateCassandraKeyspaceToAutoscale, migrateCassandraKeyspaceToManualThroughput, migrateCassandraTableToAutoscale, migrateCassandraTableToManualThroughput, updateCassandraKeyspaceThroughput, updateCassandraTableThroughput, } from "../../Utils/arm/generatedClients/cosmos/cassandraResources"; import { migrateGremlinDatabaseToAutoscale, migrateGremlinDatabaseToManualThroughput, migrateGremlinGraphToAutoscale, migrateGremlinGraphToManualThroughput, updateGremlinDatabaseThroughput, updateGremlinGraphThroughput, } from "../../Utils/arm/generatedClients/cosmos/gremlinResources"; import { migrateMongoDBCollectionToAutoscale, migrateMongoDBCollectionToManualThroughput, migrateMongoDBDatabaseToAutoscale, migrateMongoDBDatabaseToManualThroughput, updateMongoDBCollectionThroughput, updateMongoDBDatabaseThroughput, } from "../../Utils/arm/generatedClients/cosmos/mongoDBResources"; import { migrateSqlContainerToAutoscale, migrateSqlContainerToManualThroughput, migrateSqlDatabaseToAutoscale, migrateSqlDatabaseToManualThroughput, updateSqlContainerThroughput, updateSqlDatabaseThroughput, } from "../../Utils/arm/generatedClients/cosmos/sqlResources"; import { migrateTableToAutoscale, migrateTableToManualThroughput, 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 { handleError } from "../ErrorHandlingUtils"; import { parseSDKOfferResponse } from "../OfferUtility"; import { readCollectionOffer } from "./readCollectionOffer"; import { readDatabaseOffer } from "./readDatabaseOffer"; export const updateOffer = async (params: UpdateOfferParams): Promise => { let updatedOffer: Offer; const offerResourceText: string = params.collectionId ? `collection ${params.collectionId}` : `database ${params.databaseId}`; const clearMessage = logConsoleProgress(`Updating offer for ${offerResourceText}`); try { if (userContext.authType === AuthType.AAD && !userContext.features.enableSDKoperations) { if (params.collectionId) { updatedOffer = await updateCollectionOfferWithARM(params); } else if (userContext.apiType === "Tables") { // update table's database offer with SDK since RP doesn't support it updatedOffer = await updateOfferWithSDK(params); } else { updatedOffer = await updateDatabaseOfferWithARM(params); } } else { updatedOffer = await updateOfferWithSDK(params); } logConsoleInfo(`Successfully updated offer for ${offerResourceText}`); return updatedOffer; } catch (error) { handleError(error, "UpdateCollection", `Error updating offer for ${offerResourceText}`); throw error; } finally { clearMessage(); } }; const updateCollectionOfferWithARM = async (params: UpdateOfferParams): Promise => { try { switch (userContext.apiType) { case "SQL": await updateSqlContainerOffer(params); break; case "Mongo": await updateMongoCollectionOffer(params); break; case "Cassandra": await updateCassandraTableOffer(params); break; case "Gremlin": await updateGremlinGraphOffer(params); break; case "Tables": await updateTableOffer(params); break; default: throw new Error(`Unsupported default experience type: ${userContext.apiType}`); } } catch (error) { if (error.code !== "MethodNotAllowed") { throw error; } } return await readCollectionOffer({ collectionId: params.collectionId, databaseId: params.databaseId, offerId: params.currentOffer.id, }); }; const updateDatabaseOfferWithARM = async (params: UpdateOfferParams): Promise => { try { switch (userContext.apiType) { case "SQL": await updateSqlDatabaseOffer(params); break; case "Mongo": await updateMongoDatabaseOffer(params); break; case "Cassandra": await updateCassandraKeyspaceOffer(params); break; case "Gremlin": await updateGremlinDatabaseOffer(params); break; default: throw new Error(`Unsupported default experience type: ${userContext.apiType}`); } } catch (error) { if (error.code !== "MethodNotAllowed") { throw error; } } return await readDatabaseOffer({ databaseId: params.databaseId, offerId: params.currentOffer.id, }); }; const updateSqlContainerOffer = async (params: UpdateOfferParams): Promise => { const { subscriptionId, resourceGroup, databaseAccount } = userContext; const accountName = databaseAccount.name; if (params.migrateToAutoPilot) { await migrateSqlContainerToAutoscale( subscriptionId, resourceGroup, accountName, params.databaseId, params.collectionId, ); } else if (params.migrateToManual) { await migrateSqlContainerToManualThroughput( subscriptionId, resourceGroup, accountName, params.databaseId, params.collectionId, ); } else { const body: ThroughputSettingsUpdateParameters = createUpdateOfferBody(params); await updateSqlContainerThroughput( subscriptionId, resourceGroup, accountName, params.databaseId, params.collectionId, body, ); } }; const updateMongoCollectionOffer = async (params: UpdateOfferParams): Promise => { const { subscriptionId, resourceGroup, databaseAccount } = userContext; const accountName = databaseAccount.name; if (params.migrateToAutoPilot) { await migrateMongoDBCollectionToAutoscale( subscriptionId, resourceGroup, accountName, params.databaseId, params.collectionId, ); } else if (params.migrateToManual) { await migrateMongoDBCollectionToManualThroughput( subscriptionId, resourceGroup, accountName, params.databaseId, params.collectionId, ); } else { const body: ThroughputSettingsUpdateParameters = createUpdateOfferBody(params); await updateMongoDBCollectionThroughput( subscriptionId, resourceGroup, accountName, params.databaseId, params.collectionId, body, ); } }; const updateCassandraTableOffer = async (params: UpdateOfferParams): Promise => { const { subscriptionId, resourceGroup, databaseAccount } = userContext; const accountName = databaseAccount.name; if (params.migrateToAutoPilot) { await migrateCassandraTableToAutoscale( subscriptionId, resourceGroup, accountName, params.databaseId, params.collectionId, ); } else if (params.migrateToManual) { await migrateCassandraTableToManualThroughput( subscriptionId, resourceGroup, accountName, params.databaseId, params.collectionId, ); } else { const body: ThroughputSettingsUpdateParameters = createUpdateOfferBody(params); await updateCassandraTableThroughput( subscriptionId, resourceGroup, accountName, params.databaseId, params.collectionId, body, ); } }; const updateGremlinGraphOffer = async (params: UpdateOfferParams): Promise => { const { subscriptionId, resourceGroup, databaseAccount } = userContext; const accountName = databaseAccount.name; if (params.migrateToAutoPilot) { await migrateGremlinGraphToAutoscale( subscriptionId, resourceGroup, accountName, params.databaseId, params.collectionId, ); } else if (params.migrateToManual) { await migrateGremlinGraphToManualThroughput( subscriptionId, resourceGroup, accountName, params.databaseId, params.collectionId, ); } else { const body: ThroughputSettingsUpdateParameters = createUpdateOfferBody(params); await updateGremlinGraphThroughput( subscriptionId, resourceGroup, accountName, params.databaseId, params.collectionId, body, ); } }; const updateTableOffer = async (params: UpdateOfferParams): Promise => { const { subscriptionId, resourceGroup, databaseAccount } = userContext; const accountName = databaseAccount.name; if (params.migrateToAutoPilot) { await migrateTableToAutoscale(subscriptionId, resourceGroup, accountName, params.collectionId); } else if (params.migrateToManual) { await migrateTableToManualThroughput(subscriptionId, resourceGroup, accountName, params.collectionId); } else { const body: ThroughputSettingsUpdateParameters = createUpdateOfferBody(params); await updateTableThroughput(subscriptionId, resourceGroup, accountName, params.collectionId, body); } }; const updateSqlDatabaseOffer = async (params: UpdateOfferParams): Promise => { const { subscriptionId, resourceGroup, databaseAccount } = userContext; const accountName = databaseAccount.name; if (params.migrateToAutoPilot) { await migrateSqlDatabaseToAutoscale(subscriptionId, resourceGroup, accountName, params.databaseId); } else if (params.migrateToManual) { await migrateSqlDatabaseToManualThroughput(subscriptionId, resourceGroup, accountName, params.databaseId); } else { const body: ThroughputSettingsUpdateParameters = createUpdateOfferBody(params); await updateSqlDatabaseThroughput(subscriptionId, resourceGroup, accountName, params.databaseId, body); } }; const updateMongoDatabaseOffer = async (params: UpdateOfferParams): Promise => { const { subscriptionId, resourceGroup, databaseAccount } = userContext; const accountName = databaseAccount.name; if (params.migrateToAutoPilot) { await migrateMongoDBDatabaseToAutoscale(subscriptionId, resourceGroup, accountName, params.databaseId); } else if (params.migrateToManual) { await migrateMongoDBDatabaseToManualThroughput(subscriptionId, resourceGroup, accountName, params.databaseId); } else { const body: ThroughputSettingsUpdateParameters = createUpdateOfferBody(params); await updateMongoDBDatabaseThroughput(subscriptionId, resourceGroup, accountName, params.databaseId, body); } }; const updateCassandraKeyspaceOffer = async (params: UpdateOfferParams): Promise => { const { subscriptionId, resourceGroup, databaseAccount } = userContext; const accountName = databaseAccount.name; if (params.migrateToAutoPilot) { await migrateCassandraKeyspaceToAutoscale(subscriptionId, resourceGroup, accountName, params.databaseId); } else if (params.migrateToManual) { await migrateCassandraKeyspaceToManualThroughput(subscriptionId, resourceGroup, accountName, params.databaseId); } else { const body: ThroughputSettingsUpdateParameters = createUpdateOfferBody(params); await updateCassandraKeyspaceThroughput(subscriptionId, resourceGroup, accountName, params.databaseId, body); } }; const updateGremlinDatabaseOffer = async (params: UpdateOfferParams): Promise => { const { subscriptionId, resourceGroup, databaseAccount } = userContext; const accountName = databaseAccount.name; if (params.migrateToAutoPilot) { await migrateGremlinDatabaseToAutoscale(subscriptionId, resourceGroup, accountName, params.databaseId); } else if (params.migrateToManual) { await migrateGremlinDatabaseToManualThroughput(subscriptionId, resourceGroup, accountName, params.databaseId); } else { const body: ThroughputSettingsUpdateParameters = createUpdateOfferBody(params); await updateGremlinDatabaseThroughput(subscriptionId, resourceGroup, accountName, params.databaseId, body); } }; const createUpdateOfferBody = (params: UpdateOfferParams): ThroughputSettingsUpdateParameters => { const body: ThroughputSettingsUpdateParameters = { properties: { resource: {}, }, }; if (params.autopilotThroughput) { body.properties.resource.autoscaleSettings = { maxThroughput: params.autopilotThroughput, }; } else { body.properties.resource.throughput = params.manualThroughput; } return body; }; const updateOfferWithSDK = async (params: UpdateOfferParams): Promise => { const sdkOfferDefinition = params.currentOffer.offerDefinition; const newOffer: SDKOfferDefinition = { content: { offerThroughput: undefined, offerIsRUPerMinuteThroughputEnabled: false, }, _etag: undefined, _ts: undefined, _rid: sdkOfferDefinition._rid, _self: sdkOfferDefinition._self, id: sdkOfferDefinition.id, offerResourceId: sdkOfferDefinition.offerResourceId, offerVersion: sdkOfferDefinition.offerVersion, offerType: sdkOfferDefinition.offerType, resource: sdkOfferDefinition.resource, }; if (params.autopilotThroughput) { newOffer.content.offerAutopilotSettings = { maxThroughput: params.autopilotThroughput, }; } else { newOffer.content.offerThroughput = params.manualThroughput; } const options: RequestOptions = {}; if (params.migrateToAutoPilot) { options.initialHeaders = { [HttpHeaders.migrateOfferToAutopilot]: "true", }; delete newOffer.content.offerAutopilotSettings; } else if (params.migrateToManual) { options.initialHeaders = { [HttpHeaders.migrateOfferToManualThroughput]: "true", }; newOffer.content.offerAutopilotSettings = { maxThroughput: 0 }; } const sdkResponse = await client() .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); return parseSDKOfferResponse(sdkResponse); };