Lazy loading containers (#1411)

Co-authored-by: Asier Isayas <aisayas@microsoft.com>

less/resourceTree.less

@@ -11,6 +11,10 @@
   .collectionHeader {
     font-size: 12px;
   }
+
+  .loadMoreHeader {
+    color: RGB(5, 99, 193);
+  }
 }
 
 .notebookResourceTree {
@@ -23,6 +27,4 @@
   .clickDisabled {
     pointer-events: none;
   }
-}
+}
-
-

src/Common/Constants.ts

@@ -141,7 +141,7 @@ export class Queries {
   public static UnlimitedPageOption: string = "unlimited";
   public static itemsPerPage: number = 100;
   public static unlimitedItemsPerPage: number = 100; // TODO: Figure out appropriate value so it works for accounts with a large number of partitions
-
+  public static containersPerPage: number = 50;
   public static QueryEditorMinHeightRatio: number = 0.1;
   public static QueryEditorMaxHeightRatio: number = 0.4;
   public static readonly DefaultMaxDegreeOfParallelism = 6;

src/Common/dataAccess/readCollections.ts

@@ -1,3 +1,4 @@
+import { Queries } from "Common/Constants";
 import { AuthType } from "../../AuthType";
 import * as DataModels from "../../Contracts/DataModels";
 import { userContext } from "../../UserContext";
@@ -31,6 +32,35 @@ export async function readCollections(databaseId: string): Promise<DataModels.Co
   }
 }
 
+export async function readCollectionsWithPagination(
+  databaseId: string,
+  continuationToken?: string
+): Promise<DataModels.CollectionsWithPagination> {
+  const clearMessage = logConsoleProgress(`Querying containers for database ${databaseId}`);
+  try {
+    const sdkResponse = await client()
+      .database(databaseId)
+      .containers.query(
+        { query: "SELECT * FROM c" },
+        {
+          continuationToken,
+          maxItemCount: Queries.containersPerPage,
+        }
+      )
+      .fetchNext();
+    const collectionsWithPagination: DataModels.CollectionsWithPagination = {
+      collections: sdkResponse.resources as DataModels.Collection[],
+      continuationToken: sdkResponse.continuationToken,
+    };
+    return collectionsWithPagination;
+  } catch (error) {
+    handleError(error, "ReadCollections", `Error while querying containers for database ${databaseId}`);
+    throw error;
+  } finally {
+    clearMessage();
+  }
+}
+
 async function readCollectionsWithARM(databaseId: string): Promise<DataModels.Collection[]> {
   let rpResponse;
 

src/Contracts/DataModels.ts

@@ -156,6 +156,11 @@ export interface Collection extends Resource {
   requestSchema?: () => void;
 }
 
+export interface CollectionsWithPagination {
+  continuationToken?: string;
+  collections?: Collection[];
+}
+
 export interface Database extends Resource {
   collections?: Collection[];
 }

src/Contracts/ViewModels.ts

@@ -87,13 +87,13 @@ export interface Database extends TreeNode {
   isDatabaseExpanded: ko.Observable<boolean>;
   isDatabaseShared: ko.Computed<boolean>;
   isSampleDB?: boolean;
-
+  collectionsContinuationToken?: string;
   selectedSubnodeKind: ko.Observable<CollectionTabKind>;
 
   expandDatabase(): Promise<void>;
   collapseDatabase(): void;
 
-  loadCollections(): Promise<void>;
+  loadCollections(restart?: boolean): Promise<void>;
   findCollectionWithId(collectionId: string): Collection;
   openAddCollection(database: Database, event: MouseEvent): void;
   onSettingsClick: () => void;

src/Explorer/Explorer.tsx

@@ -577,7 +577,7 @@ export default class Explorer {
     try {
       await Promise.all(
         databasesToLoad.map(async (database: ViewModels.Database) => {
-          await database.loadCollections();
+          await database.loadCollections(true);
           const isNewDatabase: boolean = _.some(newDatabases, (db: ViewModels.Database) => db.id() === database.id());
           if (isNewDatabase) {
             database.expandDatabase();

src/Explorer/Panes/SettingsPane/SettingsPane.tsx

@@ -21,6 +21,11 @@ export const SettingsPane: FunctionComponent = () => {
   const [customItemPerPage, setCustomItemPerPage] = useState<number>(
     LocalStorageUtility.getEntryNumber(StorageKey.CustomItemPerPage) || 0
   );
+  const [containerPaginationEnabled, setContainerPaginationEnabled] = useState<boolean>(
+    LocalStorageUtility.hasItem(StorageKey.ContainerPaginationEnabled)
+      ? LocalStorageUtility.getEntryString(StorageKey.ContainerPaginationEnabled) === "true"
+      : false
+  );
   const [crossPartitionQueryEnabled, setCrossPartitionQueryEnabled] = useState<boolean>(
     LocalStorageUtility.hasItem(StorageKey.IsCrossPartitionQueryEnabled)
       ? LocalStorageUtility.getEntryString(StorageKey.IsCrossPartitionQueryEnabled) === "true"
@@ -50,6 +55,7 @@ export const SettingsPane: FunctionComponent = () => {
       isCustomPageOptionSelected() ? customItemPerPage : Constants.Queries.unlimitedItemsPerPage
     );
     LocalStorageUtility.setEntryNumber(StorageKey.CustomItemPerPage, customItemPerPage);
+    LocalStorageUtility.setEntryString(StorageKey.ContainerPaginationEnabled, containerPaginationEnabled.toString());
     LocalStorageUtility.setEntryString(StorageKey.IsCrossPartitionQueryEnabled, crossPartitionQueryEnabled.toString());
     LocalStorageUtility.setEntryNumber(StorageKey.MaxDegreeOfParellism, maxDegreeOfParallelism);
 
@@ -185,6 +191,25 @@ export const SettingsPane: FunctionComponent = () => {
             </div>
           </div>
         )}
+        <div className="settingsSection">
+          <div className="settingsSectionPart">
+            <div className="settingsSectionLabel">
+              Enable container pagination
+              <InfoTooltip>
+                Load 50 containers at a time. Currently, containers are not pulled in alphanumeric order.
+              </InfoTooltip>
+            </div>
+            <Checkbox
+              styles={{
+                label: { padding: 0 },
+              }}
+              className="padding"
+              ariaLabel="Enable container pagination"
+              checked={containerPaginationEnabled}
+              onChange={() => setContainerPaginationEnabled(!containerPaginationEnabled)}
+            />
+          </div>
+        </div>
         {shouldShowCrossPartitionOption && (
           <div className="settingsSection">
             <div className="settingsSectionPart">

src/Explorer/Panes/SettingsPane/__snapshots__/SettingsPane.test.tsx.snap

@@ -97,6 +97,35 @@ exports[`Settings Pane should render Default properly 1`] = `
         </div>
       </div>
     </div>
+    <div
+      className="settingsSection"
+    >
+      <div
+        className="settingsSectionPart"
+      >
+        <div
+          className="settingsSectionLabel"
+        >
+          Enable container pagination
+          <InfoTooltip>
+            Load 50 containers at a time. Currently, containers are not pulled in alphanumeric order.
+          </InfoTooltip>
+        </div>
+        <StyledCheckboxBase
+          ariaLabel="Enable container pagination"
+          checked={false}
+          className="padding"
+          onChange={[Function]}
+          styles={
+            Object {
+              "label": Object {
+                "padding": 0,
+              },
+            }
+          }
+        />
+      </div>
+    </div>
     <div
       className="settingsSection"
     >
@@ -182,6 +211,35 @@ exports[`Settings Pane should render Gremlin properly 1`] = `
   <div
     className="paneMainContent"
   >
+    <div
+      className="settingsSection"
+    >
+      <div
+        className="settingsSectionPart"
+      >
+        <div
+          className="settingsSectionLabel"
+        >
+          Enable container pagination
+          <InfoTooltip>
+            Load 50 containers at a time. Currently, containers are not pulled in alphanumeric order.
+          </InfoTooltip>
+        </div>
+        <StyledCheckboxBase
+          ariaLabel="Enable container pagination"
+          checked={false}
+          className="padding"
+          onChange={[Function]}
+          styles={
+            Object {
+              "label": Object {
+                "padding": 0,
+              },
+            }
+          }
+        />
+      </div>
+    </div>
     <div
       className="settingsSection"
     >

src/Explorer/Tree/Database.tsx

@@ -3,7 +3,7 @@ import React from "react";
 import * as _ from "underscore";
 import { AuthType } from "../../AuthType";
 import * as Constants from "../../Common/Constants";
-import { readCollections } from "../../Common/dataAccess/readCollections";
+import { readCollections, readCollectionsWithPagination } from "../../Common/dataAccess/readCollections";
 import { readDatabaseOffer } from "../../Common/dataAccess/readDatabaseOffer";
 import { getErrorMessage, getErrorStack } from "../../Common/ErrorHandlingUtils";
 import * as Logger from "../../Common/Logger";
@@ -13,6 +13,7 @@ import * as ViewModels from "../../Contracts/ViewModels";
 import { useSidePanel } from "../../hooks/useSidePanel";
 import { useTabs } from "../../hooks/useTabs";
 import { IJunoResponse, JunoClient } from "../../Juno/JunoClient";
+import * as StorageUtility from "../../Shared/StorageUtility";
 import { Action, ActionModifiers } from "../../Shared/Telemetry/TelemetryConstants";
 import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";
 import { userContext } from "../../UserContext";
@@ -38,6 +39,7 @@ export default class Database implements ViewModels.Database {
   public selectedSubnodeKind: ko.Observable<ViewModels.CollectionTabKind>;
   public junoClient: JunoClient;
   public isSampleDB: boolean;
+  public collectionsContinuationToken?: string;
   private isOfferRead: boolean;
 
   constructor(container: Explorer, data: DataModels.Database) {
@@ -140,7 +142,11 @@ export default class Database implements ViewModels.Database {
     }
 
     await this.loadOffer();
-    await this.loadCollections();
+
+    if (this.collections()?.length === 0) {
+      await this.loadCollections(true);
+    }
+
     this.isDatabaseExpanded(true);
     TelemetryProcessor.trace(Action.ExpandTreeNode, ActionModifiers.Mark, {
       description: "Database node",
@@ -162,9 +168,31 @@ export default class Database implements ViewModels.Database {
     });
   }
 
-  public async loadCollections(): Promise<void> {
+  public async loadCollections(restart = false) {
     const collectionVMs: Collection[] = [];
-    const collections: DataModels.Collection[] = await readCollections(this.id());
+    let collections: DataModels.Collection[] = [];
+    if (restart) {
+      this.collectionsContinuationToken = undefined;
+    }
+    const containerPaginationEnabled =
+      StorageUtility.LocalStorageUtility.getEntryString(StorageUtility.StorageKey.ContainerPaginationEnabled) ===
+      "true";
+    if (containerPaginationEnabled) {
+      const collectionsWithPagination: DataModels.CollectionsWithPagination = await readCollectionsWithPagination(
+        this.id(),
+        this.collectionsContinuationToken
+      );
+
+      if (collectionsWithPagination.collections?.length === Constants.Queries.containersPerPage) {
+        this.collectionsContinuationToken = collectionsWithPagination.continuationToken;
+      } else {
+        this.collectionsContinuationToken = undefined;
+      }
+      collections = collectionsWithPagination.collections;
+    } else {
+      collections = await readCollections(this.id());
+    }
+
     // TODO Remove
     // This is a hack to make Mongo collections read via ARM have a SQL-ish partitionKey property
     if (userContext.apiType === "Mongo" && userContext.authType === AuthType.AAD) {
@@ -199,7 +227,9 @@ export default class Database implements ViewModels.Database {
 
     //merge collections
     this.addCollectionsToList(collectionVMs);
-    this.deleteCollectionsFromList(deltaCollections.toDelete);
+    if (!containerPaginationEnabled || restart) {
+      this.deleteCollectionsFromList(deltaCollections.toDelete);
+    }
 
     useDatabases.getState().updateDatabase(this);
   }

src/Explorer/Tree/ResourceTree.tsx

@@ -479,6 +479,18 @@ export const ResourceTree: React.FC<ResourceTreeProps> = ({ container }: Resourc
           databaseNode.children.push(buildCollectionNode(database, collection))
         );
 
+      if (database.collectionsContinuationToken) {
+        const loadMoreNode: TreeNode = {
+          label: "load more",
+          className: "loadMoreHeader",
+          onClick: async () => {
+            await database.loadCollections();
+            useDatabases.getState().updateDatabase(database);
+          },
+        };
+        databaseNode.children.push(loadMoreNode);
+      }
+
       database.collections.subscribe((collections: ViewModels.Collection[]) => {
         collections.forEach((collection: ViewModels.Collection) =>
           databaseNode.children.push(buildCollectionNode(database, collection))

src/Shared/StorageUtility.ts

@@ -4,6 +4,7 @@ import * as SessionStorageUtility from "./SessionStorageUtility";
 export { LocalStorageUtility, SessionStorageUtility };
 export enum StorageKey {
   ActualItemPerPage,
+  ContainerPaginationEnabled,
   CustomItemPerPage,
   DatabaseAccountId,
   EncryptedKeyToken,