Get collection usage size with ARM metrics call (#327)

- Removed `readCollectionQuotaInfo` call. The only data we need is the usage size which we can get via the ARM metrics call instead.
- Added `getCollectionUsageSize` which fetches the `DataUsage` and `IndexUsage` metrics, converts them to KB, and returns the sum as the total usage size

src/Common/dataAccess/getCollectionDataUsageSize.ts

@@ -0,0 +1,87 @@
+import { armRequest } from "../../Utils/arm/request";
+import { configContext } from "../../ConfigContext";
+import { handleError } from "../ErrorHandlingUtils";
+import { userContext } from "../../UserContext";
+
+interface TimeSeriesData {
+  data: {
+    timeStamp: string;
+    total: number;
+  }[];
+  metadatavalues: {
+    name: {
+      localizedValue: string;
+      value: string;
+    };
+    value: string;
+  };
+}
+
+interface MetricsData {
+  displayDescription: string;
+  errorCode: string;
+  id: string;
+  name: {
+    value: string;
+    localizedValue: string;
+  };
+  timeseries: TimeSeriesData[];
+  type: string;
+  unit: string;
+}
+
+interface MetricsResponse {
+  cost: number;
+  interval: string;
+  namespace: string;
+  resourceregion: string;
+  timespan: string;
+  value: MetricsData[];
+}
+
+export const getCollectionUsageSizeInKB = async (databaseName: string, containerName: string): Promise<number> => {
+  const subscriptionId = userContext.subscriptionId;
+  const resourceGroup = userContext.resourceGroup;
+  const accountName = userContext.databaseAccount.name;
+  const filter = `DatabaseName eq '${databaseName}' and CollectionName eq '${containerName}'`;
+  const metricNames = "DataUsage,IndexUsage";
+  const path = `/subscriptions/${subscriptionId}/resourceGroups/${resourceGroup}/providers/Microsoft.DocumentDB/databaseAccounts/${accountName}/providers/microsoft.insights/metrics`;
+
+  try {
+    const metricsResponse: MetricsResponse = await armRequest({
+      host: configContext.ARM_ENDPOINT,
+      path,
+      method: "GET",
+      apiVersion: "2018-01-01",
+      queryParams: {
+        filter,
+        metricNames
+      }
+    });
+
+    if (metricsResponse?.value?.length !== 2) {
+      return undefined;
+    }
+
+    const dataUsageData: MetricsData = metricsResponse.value[0];
+    const indexUsagedata: MetricsData = metricsResponse.value[1];
+    const dataUsageSizeInKb: number = getUsageSizeInKb(dataUsageData);
+    const indexUsageSizeInKb: number = getUsageSizeInKb(indexUsagedata);
+
+    return dataUsageSizeInKb + indexUsageSizeInKb;
+  } catch (error) {
+    handleError(error, "getCollectionUsageSize");
+    throw error;
+  }
+};
+
+const getUsageSizeInKb = (metricsData: MetricsData): number => {
+  if (metricsData?.errorCode !== "Success") {
+    throw Error(`Get collection usage size failed: ${metricsData.errorCode}`);
+  }
+
+  const timeSeriesData: TimeSeriesData = metricsData?.timeseries?.[0];
+  const usageSizeInBytes: number = timeSeriesData?.data?.[0]?.total;
+
+  return usageSizeInBytes ? usageSizeInBytes / 1024 : 0;
+};

src/Common/dataAccess/getIndexTransformationProgress.ts

@@ -14,7 +14,7 @@ export async function getIndexTransformationProgress(databaseId: string, collect
     const response = await client()
       .database(databaseId)
       .container(collectionId)
-      .read({ populateQuotaInfo: true });
+      .read();
 
     indexTransformationPercentage = parseInt(
       response.headers[Constants.HttpHeaders.collectionIndexTransformationProgress] as string

src/Common/dataAccess/readCollectionQuotaInfo.ts

@@ -1,45 +0,0 @@
-import * as DataModels from "../../Contracts/DataModels";
-import * as HeadersUtility from "../HeadersUtility";
-import * as ViewModels from "../../Contracts/ViewModels";
-import { ContainerDefinition, Resource } from "@azure/cosmos";
-import { HttpHeaders } from "../Constants";
-import { RequestOptions } from "@azure/cosmos/dist-esm";
-import { client } from "../CosmosClient";
-import { handleError } from "../ErrorHandlingUtils";
-import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-
-interface ResourceWithStatistics {
-  statistics: DataModels.Statistic[];
-}
-
-export const readCollectionQuotaInfo = async (
-  collection: ViewModels.Collection
-): Promise<DataModels.CollectionQuotaInfo> => {
-  const clearMessage = logConsoleProgress(`Querying containers for database ${collection.id}`);
-  const options: RequestOptions = {};
-  options.populateQuotaInfo = true;
-  options.initialHeaders = options.initialHeaders || {};
-  options.initialHeaders[HttpHeaders.populatePartitionStatistics] = true;
-
-  try {
-    const response = await client()
-      .database(collection.databaseId)
-      .container(collection.id())
-      .read(options);
-    const quota: DataModels.CollectionQuotaInfo = HeadersUtility.getQuota(response.headers);
-    const resource = response.resource as ContainerDefinition & Resource & ResourceWithStatistics;
-    quota["usageSizeInKB"] = resource.statistics.reduce(
-      (previousValue: number, currentValue: DataModels.Statistic) => previousValue + currentValue.sizeInKB,
-      0
-    );
-    quota["numPartitions"] = resource.statistics.length;
-    quota["uniqueKeyPolicy"] = collection.uniqueKeyPolicy; // TODO: Remove after refactoring (#119617)
-
-    return quota;
-  } catch (error) {
-    handleError(error, "ReadCollectionQuotaInfo", `Error while querying quota info for container ${collection.id}`);
-    throw error;
-  } finally {
-    clearMessage();
-  }
-};

src/Contracts/DataModels.ts

@@ -230,18 +230,6 @@ export interface SDKOfferDefinition extends Resource {
   offerResourceId?: string;
 }
 
-export interface CollectionQuotaInfo {
-  storedProcedures: number;
-  triggers: number;
-  functions: number;
-  documentsSize: number;
-  collectionSize: number;
-  documentsCount: number;
-  usageSizeInKB: number;
-  numPartitions: number;
-  uniqueKeyPolicy?: UniqueKeyPolicy; // TODO: This should ideally not be a part of the collection quota. Remove after refactoring. (#119617)
-}
-
 export interface OfferThroughputInfo {
   minimumRUForCollection: number;
   numPhysicalPartitions: number;

src/Contracts/ViewModels.ts

@@ -120,7 +120,7 @@ export interface Collection extends CollectionBase {
   requestSchema?: () => void;
   indexingPolicy: ko.Observable<DataModels.IndexingPolicy>;
   uniqueKeyPolicy: DataModels.UniqueKeyPolicy;
-  quotaInfo: ko.Observable<DataModels.CollectionQuotaInfo>;
+  usageSizeInKB: ko.Observable<number>;
   offer: ko.Observable<DataModels.Offer>;
   conflictResolutionPolicy: ko.Observable<DataModels.ConflictResolutionPolicy>;
   changeFeedPolicy: ko.Observable<DataModels.ChangeFeedPolicy>;

src/Explorer/Controls/Settings/SettingsSubComponents/ScaleComponent.tsx

@@ -186,7 +186,7 @@ export class ScaleComponent extends React.Component<ScaleComponentProps> {
       onScaleSaveableChange={this.props.onScaleSaveableChange}
       onScaleDiscardableChange={this.props.onScaleDiscardableChange}
       getThroughputWarningMessage={this.getThroughputWarningMessage}
-      usageSizeInKB={this.props.collection.quotaInfo().usageSizeInKB}
+      usageSizeInKB={this.props.collection.usageSizeInKB()}
     />
   );
 

src/Explorer/Controls/Settings/SettingsSubComponents/__snapshots__/ScaleComponent.test.tsx.snap

@@ -58,6 +58,7 @@ exports[`ScaleComponent renders with correct initial notification 1`] = `
       spendAckChecked={false}
       throughput={1000}
       throughputBaseline={1000}
+      usageSizeInKB={100}
       wasAutopilotOriginallySet={true}
     />
     <Stack

src/Explorer/Controls/Settings/TestUtils.tsx

@@ -18,7 +18,7 @@ export const collection = ({
     excludedPaths: []
   }),
   uniqueKeyPolicy: {} as DataModels.UniqueKeyPolicy,
-  quotaInfo: ko.observable<DataModels.CollectionQuotaInfo>({} as DataModels.CollectionQuotaInfo),
+  usageSizeInKB: ko.observable(100),
   offer: ko.observable<DataModels.Offer>({
     autoscaleMaxThroughput: undefined,
     manualThroughput: 10000,

src/Explorer/Controls/Settings/__snapshots__/SettingsComponent.test.tsx.snap

@@ -1300,9 +1300,9 @@ exports[`SettingsComponent renders 1`] = `
                 "version": 2,
               },
               "partitionKeyProperty": "partitionKey",
-              "quotaInfo": [Function],
               "readSettings": [Function],
               "uniqueKeyPolicy": Object {},
+              "usageSizeInKB": [Function],
             }
           }
           container={
@@ -3871,9 +3871,9 @@ exports[`SettingsComponent renders 1`] = `
                 "version": 2,
               },
               "partitionKeyProperty": "partitionKey",
-              "quotaInfo": [Function],
               "readSettings": [Function],
               "uniqueKeyPolicy": Object {},
+              "usageSizeInKB": [Function],
             }
           }
           container={

src/Explorer/Tree/Collection.test.ts

@@ -10,10 +10,9 @@ describe("Collection", () => {
     container: Explorer,
     databaseId: string,
     data: DataModels.Collection,
-    quotaInfo: DataModels.CollectionQuotaInfo,
     offer: DataModels.Offer
   ): Collection {
-    return new Collection(container, databaseId, data, quotaInfo, offer);
+    return new Collection(container, databaseId, data);
   }
 
   function generateMockCollectionsDataModelWithPartitionKey(
@@ -50,7 +49,7 @@ describe("Collection", () => {
     });
     mockContainer.deleteCollectionText = ko.observable<string>("delete collection");
 
-    return generateCollection(mockContainer, "abc", data, {} as DataModels.CollectionQuotaInfo, {} as DataModels.Offer);
+    return generateCollection(mockContainer, "abc", data, {} as DataModels.Offer);
   }
 
   describe("Partition key path parsing", () => {

src/Explorer/Tree/Collection.ts

@@ -10,7 +10,7 @@ import { readTriggers } from "../../Common/dataAccess/readTriggers";
 import { readUserDefinedFunctions } from "../../Common/dataAccess/readUserDefinedFunctions";
 import { createDocument } from "../../Common/DocumentClientUtilityBase";
 import { readCollectionOffer } from "../../Common/dataAccess/readCollectionOffer";
-import { readCollectionQuotaInfo } from "../../Common/dataAccess/readCollectionQuotaInfo";
+import { getCollectionUsageSizeInKB } from "../../Common/dataAccess/getCollectionDataUsageSize";
 import * as Logger from "../../Common/Logger";
 import * as DataModels from "../../Contracts/DataModels";
 import * as ViewModels from "../../Contracts/ViewModels";
@@ -37,7 +37,6 @@ import UserDefinedFunction from "./UserDefinedFunction";
 import { configContext, Platform } from "../../ConfigContext";
 import Explorer from "../Explorer";
 import { userContext } from "../../UserContext";
-import TabsBase from "../Tabs/TabsBase";
 import { fetchPortalNotifications } from "../../Common/PortalNotifications";
 import { getErrorMessage, getErrorStack } from "../../Common/ErrorHandlingUtils";
 
@@ -54,7 +53,8 @@ export default class Collection implements ViewModels.Collection {
   public defaultTtl: ko.Observable<number>;
   public indexingPolicy: ko.Observable<DataModels.IndexingPolicy>;
   public uniqueKeyPolicy: DataModels.UniqueKeyPolicy;
-  public quotaInfo: ko.Observable<DataModels.CollectionQuotaInfo>;
+  public usageSizeInKB: ko.Observable<number>;
+
   public offer: ko.Observable<DataModels.Offer>;
   public conflictResolutionPolicy: ko.Observable<DataModels.ConflictResolutionPolicy>;
   public changeFeedPolicy: ko.Observable<DataModels.ChangeFeedPolicy>;
@@ -95,13 +95,7 @@ export default class Collection implements ViewModels.Collection {
   public userDefinedFunctionsFocused: ko.Observable<boolean>;
   public triggersFocused: ko.Observable<boolean>;
 
-  constructor(
-    container: Explorer,
-    databaseId: string,
-    data: DataModels.Collection,
-    quotaInfo: DataModels.CollectionQuotaInfo,
-    offer: DataModels.Offer
-  ) {
+  constructor(container: Explorer, databaseId: string, data: DataModels.Collection) {
     this.nodeKind = "Collection";
     this.container = container;
     this.self = data._self;
@@ -113,8 +107,8 @@ export default class Collection implements ViewModels.Collection {
     this.id = ko.observable(data.id);
     this.defaultTtl = ko.observable(data.defaultTtl);
     this.indexingPolicy = ko.observable(data.indexingPolicy);
-    this.quotaInfo = ko.observable(quotaInfo);
-    this.offer = ko.observable(offer);
+    this.usageSizeInKB = ko.observable();
+    this.offer = ko.observable();
     this.conflictResolutionPolicy = ko.observable(data.conflictResolutionPolicy);
     this.changeFeedPolicy = ko.observable<DataModels.ChangeFeedPolicy>(data.changeFeedPolicy);
     this.analyticalStorageTtl = ko.observable(data.analyticalStorageTtl);
@@ -607,14 +601,6 @@ export default class Collection implements ViewModels.Collection {
     }
   };
 
-  private async loadCollectionQuotaInfo(): Promise<void> {
-    // TODO: Use the collection entity cache to get quota info
-    const quotaInfoWithUniqueKeyPolicy = await readCollectionQuotaInfo(this);
-    this.uniqueKeyPolicy = quotaInfoWithUniqueKeyPolicy.uniqueKeyPolicy;
-    const quotaInfo = _.omit(quotaInfoWithUniqueKeyPolicy, "uniqueKeyPolicy");
-    this.quotaInfo(quotaInfo);
-  }
-
   public onNewQueryClick(source: any, event: MouseEvent, queryText?: string) {
     const collection: ViewModels.Collection = source.collection || source;
     const id = this.container.tabsManager.getTabs(ViewModels.CollectionTabKind.Query).length + 1;
@@ -1287,7 +1273,7 @@ export default class Collection implements ViewModels.Collection {
 
       try {
         this.offer(await readCollectionOffer(params));
-        await this.loadCollectionQuotaInfo();
+        this.usageSizeInKB(await getCollectionUsageSizeInKB(this.databaseId, this.id()));
 
         TelemetryProcessor.traceSuccess(
           Action.LoadOffers,

src/Explorer/Tree/Database.ts

@@ -193,7 +193,7 @@ export default class Database implements ViewModels.Database {
     });
 
     deltaCollections.toAdd.forEach((collection: DataModels.Collection) => {
-      const collectionVM: Collection = new Collection(this.container, this.id(), collection, null, null);
+      const collectionVM: Collection = new Collection(this.container, this.id(), collection);
       collectionVMs.push(collectionVM);
     });
 

src/Explorer/Tree/ResourceTreeAdapter.test.tsx

@@ -229,9 +229,7 @@ const createMockCollection = (): ViewModels.Collection => {
   const mockCollectionVM: ViewModels.Collection = new Collection(
     createMockContainer(),
     "fakeDatabaseId",
-    mockCollection,
-    undefined,
-    undefined
+    mockCollection
   );
 
   return mockCollectionVM;

src/Utils/arm/request.ts

@@ -33,18 +33,36 @@ export class ARMError extends Error {
   public code: string | number;
 }
 
+interface ARMQueryParams {
+  filter?: string;
+  metricNames?: string;
+}
+
 interface Options {
   host: string;
   path: string;
   apiVersion: string;
   method: "GET" | "POST" | "PUT" | "PATCH" | "DELETE" | "HEAD";
   body?: unknown;
+  queryParams?: ARMQueryParams;
 }
 
 // TODO: This is very similar to what is happening in ResourceProviderClient.ts. Should probably merge them.
-export async function armRequest<T>({ host, path, apiVersion, method, body: requestBody }: Options): Promise<T> {
+export async function armRequest<T>({
+  host,
+  path,
+  apiVersion,
+  method,
+  body: requestBody,
+  queryParams
+}: Options): Promise<T> {
   const url = new URL(path, host);
   url.searchParams.append("api-version", configContext.armAPIVersion || apiVersion);
+  if (queryParams) {
+    queryParams.filter && url.searchParams.append("$filter", queryParams.filter);
+    queryParams.metricNames && url.searchParams.append("metricnames", queryParams.metricNames);
+  }
+
   const response = await window.fetch(url.href, {
     method,
     headers: {