Move readCollection, readCollections, and readDatabases to ARM (#134)

src/Common/DataAccessUtilityBase.ts

@@ -418,26 +418,6 @@ export function deleteTrigger(
   );
 }
 
-export function readCollections(database: ViewModels.Database, options: any): Q.Promise<DataModels.Collection[]> {
-  return Q(
-    client()
-      .database(database.id())
-      .containers.readAll()
-      .fetchAll()
-      .then(response => response.resources as DataModels.Collection[])
-  );
-}
-
-export function readCollection(databaseId: string, collectionId: string): Q.Promise<DataModels.Collection> {
-  return Q(
-    client()
-      .database(databaseId)
-      .container(collectionId)
-      .read()
-      .then(response => response.resource as DataModels.Collection)
-  );
-}
-
 export function readCollectionQuotaInfo(
   collection: ViewModels.Collection,
   options: any
@@ -507,26 +487,6 @@ export function readOffer(requestedResource: DataModels.Offer, options: any): Q.
   );
 }
 
-export function readDatabases(options: any): Q.Promise<DataModels.Database[]> {
-  try {
-    if (configContext.platform === Platform.Portal) {
-      return sendCachedDataMessage<DataModels.Database[]>(MessageTypes.AllDatabases, [
-        (<any>window).dataExplorer.databaseAccount().id,
-        Constants.ClientDefaults.portalCacheTimeoutMs
-      ]);
-    }
-  } catch (error) {
-    // If error getting cached DBs, continue on and read via SDK
-  }
-
-  return Q(
-    client()
-      .databases.readAll()
-      .fetchAll()
-      .then(response => response.resources as DataModels.Database[])
-  );
-}
-
 export function getOrCreateDatabaseAndCollection(
   request: DataModels.CreateDatabaseAndCollectionRequest,
   options: any

src/Common/DocumentClientUtilityBase.ts

@@ -806,63 +806,6 @@ export function refreshCachedOffers(): Q.Promise<void> {
   return DataAccessUtilityBase.refreshCachedOffers();
 }
 
-export function readCollections(database: ViewModels.Database, options: any = {}): Q.Promise<DataModels.Collection[]> {
-  var deferred = Q.defer<DataModels.Collection[]>();
-  const id = NotificationConsoleUtils.logConsoleMessage(
-    ConsoleDataType.InProgress,
-    `Querying containers for database ${database.id()}`
-  );
-  DataAccessUtilityBase.readCollections(database, options)
-    .then(
-      (collections: DataModels.Collection[]) => {
-        deferred.resolve(collections);
-      },
-      (error: any) => {
-        NotificationConsoleUtils.logConsoleMessage(
-          ConsoleDataType.Error,
-          `Error while querying containers for database ${database.id()}:\n ${JSON.stringify(error)}`
-        );
-        Logger.logError(JSON.stringify(error), "ReadCollections", error.code);
-        sendNotificationForError(error);
-        deferred.reject(error);
-      }
-    )
-    .finally(() => {
-      NotificationConsoleUtils.clearInProgressMessageWithId(id);
-    });
-
-  return deferred.promise;
-}
-
-export function readCollection(databaseId: string, collectionId: string): Q.Promise<DataModels.Collection> {
-  const deferred = Q.defer<DataModels.Collection>();
-  const id = NotificationConsoleUtils.logConsoleMessage(
-    ConsoleDataType.InProgress,
-    `Querying container ${collectionId}`
-  );
-
-  DataAccessUtilityBase.readCollection(databaseId, collectionId)
-    .then(
-      (collection: DataModels.Collection) => {
-        deferred.resolve(collection);
-      },
-      (error: any) => {
-        NotificationConsoleUtils.logConsoleMessage(
-          ConsoleDataType.Error,
-          `Error while querying containers for database ${databaseId}:\n ${JSON.stringify(error)}`
-        );
-        Logger.logError(JSON.stringify(error), "ReadCollections", error.code);
-        sendNotificationForError(error);
-        deferred.reject(error);
-      }
-    )
-    .finally(() => {
-      NotificationConsoleUtils.clearInProgressMessageWithId(id);
-    });
-
-  return deferred.promise;
-}
-
 export function readCollectionQuotaInfo(
   collection: ViewModels.Collection,
   options?: any
@@ -950,31 +893,6 @@ export function readOffer(
   return deferred.promise;
 }
 
-export function readDatabases(options: any): Q.Promise<DataModels.Database[]> {
-  var deferred = Q.defer<any>();
-  const id = NotificationConsoleUtils.logConsoleMessage(ConsoleDataType.InProgress, "Querying databases");
-  DataAccessUtilityBase.readDatabases(options)
-    .then(
-      (databases: DataModels.Database[]) => {
-        deferred.resolve(databases);
-      },
-      (error: any) => {
-        NotificationConsoleUtils.logConsoleMessage(
-          ConsoleDataType.Error,
-          `Error while querying databases:\n ${JSON.stringify(error)}`
-        );
-        Logger.logError(JSON.stringify(error), "ReadDatabases", error.code);
-        sendNotificationForError(error);
-        deferred.reject(error);
-      }
-    )
-    .finally(() => {
-      NotificationConsoleUtils.clearInProgressMessageWithId(id);
-    });
-
-  return deferred.promise;
-}
-
 export function getOrCreateDatabaseAndCollection(
   request: DataModels.CreateDatabaseAndCollectionRequest,
   options: any = {}

src/Common/dataAccess/deleteCollection.test.ts

@@ -1,16 +1,17 @@
 jest.mock("../../Utils/arm/request");
 jest.mock("../MessageHandler");
+jest.mock("../CosmosClient");
 import { deleteCollection } from "./deleteCollection";
 import { armRequest } from "../../Utils/arm/request";
 import { AuthType } from "../../AuthType";
+import { client } from "../CosmosClient";
 import { updateUserContext } from "../../UserContext";
 import { DatabaseAccount } from "../../Contracts/DataModels";
 import { sendCachedDataMessage } from "../MessageHandler";
 import { DefaultAccountExperienceType } from "../../DefaultAccountExperienceType";
 
 describe("deleteCollection", () => {
-  it("should call ARM if logged in with AAD", async () => {
-    window.authType = AuthType.AAD;
+  beforeAll(() => {
     updateUserContext({
       databaseAccount: {
         name: "test"
@@ -18,8 +19,28 @@ describe("deleteCollection", () => {
       defaultExperience: DefaultAccountExperienceType.DocumentDB
     });
     (sendCachedDataMessage as jest.Mock).mockResolvedValue(undefined);
+  });
+
+  it("should call ARM if logged in with AAD", async () => {
+    window.authType = AuthType.AAD;
     await deleteCollection("database", "collection");
     expect(armRequest).toHaveBeenCalled();
   });
-  // TODO: Test non-AAD case
+
+  it("should call SDK if not logged in with non-AAD method", async () => {
+    window.authType = AuthType.MasterKey;
+    (client as jest.Mock).mockReturnValue({
+      database: () => {
+        return {
+          container: () => {
+            return {
+              delete: (): unknown => undefined
+            };
+          }
+        };
+      }
+    });
+    await deleteCollection("database", "collection");
+    expect(client).toHaveBeenCalled();
+  });
 });

src/Common/dataAccess/deleteCollection.ts

@@ -6,6 +6,8 @@ import { deleteMongoDBCollection } from "../../Utils/arm/generatedClients/2020-0
 import { deleteGremlinGraph } from "../../Utils/arm/generatedClients/2020-04-01/gremlinResources";
 import { deleteTable } from "../../Utils/arm/generatedClients/2020-04-01/tableResources";
 import { logConsoleError, logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
+import { logError } from "../Logger";
+import { sendNotificationForError } from "./sendNotificationForError";
 import { userContext } from "../../UserContext";
 import { client } from "../CosmosClient";
 import { refreshCachedResources } from "../DataAccessUtilityBase";
@@ -23,6 +25,8 @@ export async function deleteCollection(databaseId: string, collectionId: string)
     }
   } catch (error) {
     logConsoleError(`Error while deleting container ${collectionId}:\n ${JSON.stringify(error)}`);
+    logError(JSON.stringify(error), "DeleteCollection", error.code);
+    sendNotificationForError(error);
     throw error;
   }
   logConsoleInfo(`Successfully deleted container ${collectionId}`);

src/Common/dataAccess/deleteDatabase.test.ts

@@ -1,16 +1,17 @@
 jest.mock("../../Utils/arm/request");
 jest.mock("../MessageHandler");
+jest.mock("../CosmosClient");
 import { deleteDatabase } from "./deleteDatabase";
 import { armRequest } from "../../Utils/arm/request";
 import { AuthType } from "../../AuthType";
+import { client } from "../CosmosClient";
 import { updateUserContext } from "../../UserContext";
 import { DatabaseAccount } from "../../Contracts/DataModels";
 import { sendCachedDataMessage } from "../MessageHandler";
 import { DefaultAccountExperienceType } from "../../DefaultAccountExperienceType";
 
 describe("deleteDatabase", () => {
-  it("should call ARM if logged in with AAD", async () => {
-    window.authType = AuthType.AAD;
+  beforeAll(() => {
     updateUserContext({
       databaseAccount: {
         name: "test"
@@ -18,8 +19,24 @@ describe("deleteDatabase", () => {
       defaultExperience: DefaultAccountExperienceType.DocumentDB
     });
     (sendCachedDataMessage as jest.Mock).mockResolvedValue(undefined);
+  });
+
+  it("should call ARM if logged in with AAD", async () => {
+    window.authType = AuthType.AAD;
     await deleteDatabase("database");
     expect(armRequest).toHaveBeenCalled();
   });
-  // TODO: Test non-AAD case
+
+  it("should call SDK if not logged in with non-AAD method", async () => {
+    window.authType = AuthType.MasterKey;
+    (client as jest.Mock).mockReturnValue({
+      database: () => {
+        return {
+          delete: (): unknown => undefined
+        };
+      }
+    });
+    await deleteDatabase("database");
+    expect(client).toHaveBeenCalled();
+  });
 });

src/Common/dataAccess/readCollection.test.ts

@@ -0,0 +1,35 @@
+jest.mock("../CosmosClient");
+import { AuthType } from "../../AuthType";
+import { DatabaseAccount } from "../../Contracts/DataModels";
+import { DefaultAccountExperienceType } from "../../DefaultAccountExperienceType";
+import { client } from "../CosmosClient";
+import { readCollection } from "./readCollection";
+import { updateUserContext } from "../../UserContext";
+
+describe("readCollection", () => {
+  beforeAll(() => {
+    updateUserContext({
+      databaseAccount: {
+        name: "test"
+      } as DatabaseAccount,
+      defaultExperience: DefaultAccountExperienceType.DocumentDB
+    });
+  });
+
+  it("should call SDK if logged in with resource token", async () => {
+    window.authType = AuthType.ResourceToken;
+    (client as jest.Mock).mockReturnValue({
+      database: () => {
+        return {
+          container: () => {
+            return {
+              read: (): unknown => ({})
+            };
+          }
+        };
+      }
+    });
+    await readCollection("database", "collection");
+    expect(client).toHaveBeenCalled();
+  });
+});

src/Common/dataAccess/readCollection.ts

@@ -0,0 +1,24 @@
+import * as DataModels from "../../Contracts/DataModels";
+import { client } from "../CosmosClient";
+import { logConsoleProgress, logConsoleError } from "../../Utils/NotificationConsoleUtils";
+import { logError } from "../Logger";
+import { sendNotificationForError } from "./sendNotificationForError";
+
+export async function readCollection(databaseId: string, collectionId: string): Promise<DataModels.Collection> {
+  let collection: DataModels.Collection;
+  const clearMessage = logConsoleProgress(`Querying container ${collectionId}`);
+  try {
+    const response = await client()
+      .database(databaseId)
+      .container(collectionId)
+      .read();
+    collection = response.resource as DataModels.Collection;
+  } catch (error) {
+    logConsoleError(`Error while querying container ${collectionId}:\n ${JSON.stringify(error)}`);
+    logError(JSON.stringify(error), "ReadCollection", error.code);
+    sendNotificationForError(error);
+    throw error;
+  }
+  clearMessage();
+  return collection;
+}

src/Common/dataAccess/readCollections.test.ts

@@ -0,0 +1,45 @@
+jest.mock("../../Utils/arm/request");
+jest.mock("../CosmosClient");
+import { AuthType } from "../../AuthType";
+import { DatabaseAccount } from "../../Contracts/DataModels";
+import { DefaultAccountExperienceType } from "../../DefaultAccountExperienceType";
+import { armRequest } from "../../Utils/arm/request";
+import { client } from "../CosmosClient";
+import { readCollections } from "./readCollections";
+import { updateUserContext } from "../../UserContext";
+
+describe("readCollections", () => {
+  beforeAll(() => {
+    updateUserContext({
+      databaseAccount: {
+        name: "test"
+      } as DatabaseAccount,
+      defaultExperience: DefaultAccountExperienceType.DocumentDB
+    });
+  });
+
+  it("should call ARM if logged in with AAD", async () => {
+    window.authType = AuthType.AAD;
+    await readCollections("database");
+    expect(armRequest).toHaveBeenCalled();
+  });
+
+  it("should call SDK if not logged in with non-AAD method", async () => {
+    window.authType = AuthType.MasterKey;
+    (client as jest.Mock).mockReturnValue({
+      database: () => {
+        return {
+          containers: {
+            readAll: () => {
+              return {
+                fetchAll: (): unknown => []
+              };
+            }
+          }
+        };
+      }
+    });
+    await readCollections("database");
+    expect(client).toHaveBeenCalled();
+  });
+});

src/Common/dataAccess/readCollections.ts

@@ -0,0 +1,66 @@
+import * as DataModels from "../../Contracts/DataModels";
+import { AuthType } from "../../AuthType";
+import { DefaultAccountExperienceType } from "../../DefaultAccountExperienceType";
+import { client } from "../CosmosClient";
+import { listSqlContainers } from "../../Utils/arm/generatedClients/2020-04-01/sqlResources";
+import { listCassandraTables } from "../../Utils/arm/generatedClients/2020-04-01/cassandraResources";
+import { listMongoDBCollections } from "../../Utils/arm/generatedClients/2020-04-01/mongoDBResources";
+import { listGremlinGraphs } from "../../Utils/arm/generatedClients/2020-04-01/gremlinResources";
+import { listTables } from "../../Utils/arm/generatedClients/2020-04-01/tableResources";
+import { logConsoleProgress, logConsoleError } from "../../Utils/NotificationConsoleUtils";
+import { logError } from "../Logger";
+import { sendNotificationForError } from "./sendNotificationForError";
+import { userContext } from "../../UserContext";
+
+export async function readCollections(databaseId: string): Promise<DataModels.Collection[]> {
+  let collections: DataModels.Collection[];
+  const clearMessage = logConsoleProgress(`Querying containers for database ${databaseId}`);
+  try {
+    if (window.authType === AuthType.AAD) {
+      collections = await readCollectionsWithARM(databaseId);
+    } else {
+      const sdkResponse = await client()
+        .database(databaseId)
+        .containers.readAll()
+        .fetchAll();
+      collections = sdkResponse.resources as DataModels.Collection[];
+    }
+  } catch (error) {
+    logConsoleError(`Error while querying containers for database ${databaseId}:\n ${JSON.stringify(error)}`);
+    logError(JSON.stringify(error), "ReadCollections", error.code);
+    sendNotificationForError(error);
+    throw error;
+  }
+  clearMessage();
+  return collections;
+}
+
+async function readCollectionsWithARM(databaseId: string): Promise<DataModels.Collection[]> {
+  let rpResponse;
+  const subscriptionId = userContext.subscriptionId;
+  const resourceGroup = userContext.resourceGroup;
+  const accountName = userContext.databaseAccount.name;
+  const defaultExperience = userContext.defaultExperience;
+
+  switch (defaultExperience) {
+    case DefaultAccountExperienceType.DocumentDB:
+      rpResponse = await listSqlContainers(subscriptionId, resourceGroup, accountName, databaseId);
+      break;
+    case DefaultAccountExperienceType.MongoDB:
+      rpResponse = await listMongoDBCollections(subscriptionId, resourceGroup, accountName, databaseId);
+      break;
+    case DefaultAccountExperienceType.Cassandra:
+      rpResponse = await listCassandraTables(subscriptionId, resourceGroup, accountName, databaseId);
+      break;
+    case DefaultAccountExperienceType.Graph:
+      rpResponse = await listGremlinGraphs(subscriptionId, resourceGroup, accountName, databaseId);
+      break;
+    case DefaultAccountExperienceType.Table:
+      rpResponse = await listTables(subscriptionId, resourceGroup, accountName);
+      break;
+    default:
+      throw new Error(`Unsupported default experience type: ${defaultExperience}`);
+  }
+
+  return rpResponse?.value?.map(collection => collection.properties?.resource as DataModels.Collection);
+}

src/Common/dataAccess/readDatabases.test.ts

@@ -0,0 +1,41 @@
+jest.mock("../../Utils/arm/request");
+jest.mock("../CosmosClient");
+import { AuthType } from "../../AuthType";
+import { DatabaseAccount } from "../../Contracts/DataModels";
+import { DefaultAccountExperienceType } from "../../DefaultAccountExperienceType";
+import { armRequest } from "../../Utils/arm/request";
+import { client } from "../CosmosClient";
+import { readDatabases } from "./readDatabases";
+import { updateUserContext } from "../../UserContext";
+
+describe("readDatabases", () => {
+  beforeAll(() => {
+    updateUserContext({
+      databaseAccount: {
+        name: "test"
+      } as DatabaseAccount,
+      defaultExperience: DefaultAccountExperienceType.DocumentDB
+    });
+  });
+
+  it("should call ARM if logged in with AAD", async () => {
+    window.authType = AuthType.AAD;
+    await readDatabases();
+    expect(armRequest).toHaveBeenCalled();
+  });
+
+  it("should call SDK if not logged in with non-AAD method", async () => {
+    window.authType = AuthType.MasterKey;
+    (client as jest.Mock).mockReturnValue({
+      databases: {
+        readAll: () => {
+          return {
+            fetchAll: (): unknown => []
+          };
+        }
+      }
+    });
+    await readDatabases();
+    expect(client).toHaveBeenCalled();
+  });
+});

src/Common/dataAccess/readDatabases.ts

@@ -0,0 +1,61 @@
+import * as DataModels from "../../Contracts/DataModels";
+import { AuthType } from "../../AuthType";
+import { DefaultAccountExperienceType } from "../../DefaultAccountExperienceType";
+import { client } from "../CosmosClient";
+import { listSqlDatabases } from "../../Utils/arm/generatedClients/2020-04-01/sqlResources";
+import { listCassandraKeyspaces } from "../../Utils/arm/generatedClients/2020-04-01/cassandraResources";
+import { listMongoDBDatabases } from "../../Utils/arm/generatedClients/2020-04-01/mongoDBResources";
+import { listGremlinDatabases } from "../../Utils/arm/generatedClients/2020-04-01/gremlinResources";
+import { logConsoleProgress, logConsoleError } from "../../Utils/NotificationConsoleUtils";
+import { logError } from "../Logger";
+import { sendNotificationForError } from "./sendNotificationForError";
+import { userContext } from "../../UserContext";
+
+export async function readDatabases(): Promise<DataModels.Database[]> {
+  let databases: DataModels.Database[];
+  const clearMessage = logConsoleProgress(`Querying databases`);
+  try {
+    if (window.authType === AuthType.AAD) {
+      databases = await readDatabasesWithARM();
+    } else {
+      const sdkResponse = await client()
+        .databases.readAll()
+        .fetchAll();
+      databases = sdkResponse.resources as DataModels.Database[];
+    }
+  } catch (error) {
+    logConsoleError(`Error while querying databases:\n ${JSON.stringify(error)}`);
+    logError(JSON.stringify(error), "ReadDatabases", error.code);
+    sendNotificationForError(error);
+    throw error;
+  }
+  clearMessage();
+  return databases;
+}
+
+async function readDatabasesWithARM(): Promise<DataModels.Database[]> {
+  let rpResponse;
+  const subscriptionId = userContext.subscriptionId;
+  const resourceGroup = userContext.resourceGroup;
+  const accountName = userContext.databaseAccount.name;
+  const defaultExperience = userContext.defaultExperience;
+
+  switch (defaultExperience) {
+    case DefaultAccountExperienceType.DocumentDB:
+      rpResponse = await listSqlDatabases(subscriptionId, resourceGroup, accountName);
+      break;
+    case DefaultAccountExperienceType.MongoDB:
+      rpResponse = await listMongoDBDatabases(subscriptionId, resourceGroup, accountName);
+      break;
+    case DefaultAccountExperienceType.Cassandra:
+      rpResponse = await listCassandraKeyspaces(subscriptionId, resourceGroup, accountName);
+      break;
+    case DefaultAccountExperienceType.Graph:
+      rpResponse = await listGremlinDatabases(subscriptionId, resourceGroup, accountName);
+      break;
+    default:
+      throw new Error(`Unsupported default experience type: ${defaultExperience}`);
+  }
+
+  return rpResponse?.value?.map(database => database.properties?.resource as DataModels.Database);
+}

src/Explorer/Explorer.ts

@@ -15,7 +15,9 @@ import CassandraAddCollectionPane from "./Panes/CassandraAddCollectionPane";
 import Database from "./Tree/Database";
 import DeleteCollectionConfirmationPane from "./Panes/DeleteCollectionConfirmationPane";
 import DeleteDatabaseConfirmationPane from "./Panes/DeleteDatabaseConfirmationPane";
-import { readDatabases, readCollection, readOffers, refreshCachedResources } from "../Common/DocumentClientUtilityBase";
+import { readOffers, refreshCachedResources } from "../Common/DocumentClientUtilityBase";
+import { readCollection } from "../Common/dataAccess/readCollection";
+import { readDatabases } from "../Common/dataAccess/readDatabases";
 import EditTableEntityPane from "./Panes/Tables/EditTableEntityPane";
 import EnvironmentUtility from "../Common/EnvironmentUtility";
 import GraphStylingPane from "./Panes/GraphStylingPane";
@@ -1420,7 +1422,7 @@ export default class Explorer {
 
     const refreshDatabases = (offers?: DataModels.Offer[]) => {
       this._setLoadingStatusText("Fetching databases...");
-      readDatabases(null /*options*/).then(
+      readDatabases().then(
         (databases: DataModels.Database[]) => {
           this._setLoadingStatusText("Successfully fetched databases.");
           TelemetryProcessor.traceSuccess(

src/Explorer/Tree/Database.ts

@@ -12,7 +12,8 @@ import * as NotificationConsoleUtils from "../../Utils/NotificationConsoleUtils"
 import { ConsoleDataType } from "../Menus/NotificationConsole/NotificationConsoleComponent";
 import * as Logger from "../../Common/Logger";
 import Explorer from "../Explorer";
-import { readCollections, readOffers, readOffer } from "../../Common/DocumentClientUtilityBase";
+import { readOffers, readOffer } from "../../Common/DocumentClientUtilityBase";
+import { readCollections } from "../../Common/dataAccess/readCollections";
 
 export default class Database implements ViewModels.Database {
   public nodeKind: string;
@@ -259,7 +260,7 @@ export default class Database implements ViewModels.Database {
     let collectionVMs: Collection[] = [];
     let deferred: Q.Deferred<void> = Q.defer<void>();
 
-    readCollections(this).then(
+    readCollections(this.id()).then(
       (collections: DataModels.Collection[]) => {
         let collectionsToAddVMPromises: Q.Promise<any>[] = [];
         let deltaCollections = this.getDeltaCollections(collections);

src/Utils/arm/generatedClients/2020-04-01/types.ts

@@ -162,6 +162,7 @@ export type DatabaseAccountGetResults = ARMResourceProperties & {
 };
 
 /* The system generated resource properties associated with SQL databases, SQL containers, Gremlin databases and Gremlin graphs. */
+// TODO: ExtendedResourceProperties was missing some properties such as _self which was manually added. Need to fix this in the RP spec.
 export interface ExtendedResourceProperties {
   /* A system generated property. A unique identifier. */
   readonly _rid: string;
@@ -169,6 +170,8 @@ export interface ExtendedResourceProperties {
   readonly _ts: unknown;
   /* A system generated property representing the resource etag required for optimistic concurrency control. */
   readonly _etag: string;
+  // TODO: This property was manually added. It should be auto-generated like the other properties.
+  readonly _self: string;
 }
 
 /* An Azure Cosmos DB resource throughput. */