Implement Sample data import for Fabric Home (#2101)

* Implement dialog to import sample data

* Fix format

* Cosmetic fixes

* fix: update help link to point to the new documentation URL

---------

Co-authored-by: Sevo Kukol <sevoku@microsoft.com>

src/Contracts/ViewModels.ts

@@ -1,4 +1,5 @@
 import {
+  JSONObject,
   QueryMetrics,
   Resource,
   StoredProcedureDefinition,
@@ -206,6 +207,12 @@ export interface Collection extends CollectionBase {
   onDragOver(source: Collection, event: { originalEvent: DragEvent }): void;
   onDrop(source: Collection, event: { originalEvent: DragEvent }): void;
   uploadFiles(fileList: FileList): Promise<{ data: UploadDetailsRecord[] }>;
+  bulkInsertDocuments(documents: JSONObject[]): Promise<{
+    numSucceeded: number;
+    numFailed: number;
+    numThrottled: number;
+    errors: string[];
+  }>;
 }
 
 /**

src/Explorer/DataSamples/DataSamplesUtil.ts

@@ -6,6 +6,7 @@ import Explorer from "../Explorer";
 import { useDatabases } from "../useDatabases";
 import { ContainerSampleGenerator } from "./ContainerSampleGenerator";
 
+// TODO: this does not seem to be used. Remove?
 export class DataSamplesUtil {
   private static readonly DialogTitle = "Create Sample Container";
   constructor(private container: Explorer) {}

src/Explorer/SplashScreen/FabricHome.tsx

@@ -3,6 +3,8 @@
  */
 import { Link, makeStyles, tokens } from "@fluentui/react-components";
 import { DocumentAddRegular, LinkMultipleRegular } from "@fluentui/react-icons";
+import { SampleDataImportDialog } from "Explorer/SplashScreen/SampleDataImportDialog";
+import { CosmosFluentProvider } from "Explorer/Theme/ThemeUtil";
 import { isFabricNative } from "Platform/Fabric/FabricUtil";
 import * as React from "react";
 import { userContext } from "UserContext";
@@ -108,12 +110,10 @@ const FabricHomeScreenButton: React.FC<FabricHomeScreenButtonProps & { className
   onClick,
 }) => {
   const styles = useStyles();
-
-  // TODO Make this a11y copmliant: aria-label for icon
   return (
     <div role="button" className={`${styles.buttonContainer} ${className}`} onClick={onClick}>
       <div className={styles.buttonUpperPart}>{icon}</div>
-      <div className={styles.buttonLowerPart}>
+      <div aria-label={title} className={styles.buttonLowerPart}>
         <div>{title}</div>
         <div>{description}</div>
       </div>
@@ -123,6 +123,8 @@ const FabricHomeScreenButton: React.FC<FabricHomeScreenButtonProps & { className
 
 export const FabricHomeScreen: React.FC<SplashScreenProps> = (props: SplashScreenProps) => {
   const styles = useStyles();
+  const [openSampleDataImportDialog, setOpenSampleDataImportDialog] = React.useState(false);
+
   const getSplashScreenButtons = (): JSX.Element => {
     const buttons: FabricHomeScreenButtonProps[] = [
       {
@@ -138,11 +140,13 @@ export const FabricHomeScreen: React.FC<SplashScreenProps> = (props: SplashScree
         title: "Sample data",
         description: "Automatically load sample data in your database",
         icon: <img src={CosmosDbBlackIcon} />,
+        onClick: () => setOpenSampleDataImportDialog(true),
       },
       {
         title: "App development",
         description: "Start here to use an SDK to build your apps",
         icon: <LinkMultipleRegular />,
+        onClick: () => window.open("https://aka.ms/cosmosdbfabricsdk", "_blank"),
       },
     ];
 
@@ -157,17 +161,25 @@ export const FabricHomeScreen: React.FC<SplashScreenProps> = (props: SplashScree
 
   const title = "Build your database";
   return (
-    <div className={styles.homeContainer}>
-      <div className={styles.title} role="heading" aria-label={title}>
-        {title}
-      </div>
-      {getSplashScreenButtons()}
-      <div className={styles.footer}>
-        Need help?{" "}
-        <Link href="https://cosmos.azure.com/docs" target="_blank">
-          Learn more <img src={LinkIcon} alt="Learn more" />
-        </Link>
-      </div>
-    </div>
+    <>
+      <CosmosFluentProvider className={styles.homeContainer}>
+        <SampleDataImportDialog
+          open={openSampleDataImportDialog}
+          setOpen={setOpenSampleDataImportDialog}
+          explorer={props.explorer}
+          databaseName={userContext.fabricContext?.databaseName}
+        />
+        <div className={styles.title} role="heading" aria-label={title}>
+          {title}
+        </div>
+        {getSplashScreenButtons()}
+        <div className={styles.footer}>
+          Need help?{" "}
+          <Link href="https://aka.ms/cosmosdbfabricdocs" target="_blank">
+            Learn more <img src={LinkIcon} alt="Learn more" />
+          </Link>
+        </div>
+      </CosmosFluentProvider>
+    </>
   );
 };

src/Explorer/SplashScreen/SampleDataImportDialog.tsx

@@ -0,0 +1,158 @@
+import {
+  Button,
+  Dialog,
+  DialogActions,
+  DialogBody,
+  DialogContent,
+  DialogSurface,
+  DialogTitle,
+  makeStyles,
+  Spinner,
+  tokens,
+} from "@fluentui/react-components";
+import Explorer from "Explorer/Explorer";
+import { checkContainerExists, createContainer, importData } from "Explorer/SplashScreen/SampleUtil";
+import React, { useEffect, useState } from "react";
+import * as ViewModels from "../../Contracts/ViewModels";
+
+const SAMPLE_DATA_CONTAINER_NAME = "SampleData";
+
+const useStyles = makeStyles({
+  dialogContent: {
+    alignItems: "center",
+    marginBottom: tokens.spacingVerticalL,
+  },
+});
+
+/**
+ * This dialog:
+ * - creates a container
+ * - imports data into the container
+ * @param props
+ * @returns
+ */
+export const SampleDataImportDialog: React.FC<{
+  open: boolean;
+  setOpen: (open: boolean) => void;
+  explorer: Explorer;
+  databaseName: string;
+}> = (props) => {
+  const [status, setStatus] = useState<"idle" | "creating" | "importing" | "completed" | "error">("idle");
+  const [errorMessage, setErrorMessage] = useState<string | null>(null);
+  const containerName = SAMPLE_DATA_CONTAINER_NAME;
+  const [collection, setCollection] = useState<ViewModels.Collection>(undefined);
+  const styles = useStyles();
+
+  useEffect(() => {
+    // Reset state when dialog opens
+    if (props.open) {
+      setStatus("idle");
+      setErrorMessage(undefined);
+    }
+  }, [props.open]);
+
+  const handleStartImport = async (): Promise<void> => {
+    setStatus("creating");
+    const databaseName = props.databaseName;
+    if (checkContainerExists(databaseName, containerName)) {
+      const msg = `The container "${containerName}" in database "${databaseName}" already exists. Please delete it and retry.`;
+      setStatus("error");
+      setErrorMessage(msg);
+      return;
+    }
+
+    let collection;
+    try {
+      collection = await createContainer(databaseName, containerName, props.explorer);
+    } catch (error) {
+      setStatus("error");
+      setErrorMessage(`Failed to create container: ${error instanceof Error ? error.message : String(error)}`);
+      return;
+    }
+
+    try {
+      setStatus("importing");
+      await importData(collection);
+      setCollection(collection);
+      setStatus("completed");
+    } catch (error) {
+      setStatus("error");
+      setErrorMessage(`Failed to import data: ${error instanceof Error ? error.message : String(error)}`);
+    }
+  };
+
+  const handleActionOnClick = () => {
+    switch (status) {
+      case "idle":
+        handleStartImport();
+        break;
+      case "error":
+        props.setOpen(false);
+        break;
+      case "creating":
+      case "importing":
+        props.setOpen(false);
+        break;
+      case "completed":
+        props.setOpen(false);
+        collection.openTab();
+        break;
+    }
+  };
+
+  const renderContent = () => {
+    switch (status) {
+      case "idle":
+        return `Create a container "${containerName}" and import sample data into it. This may take a few minutes.`;
+
+      case "creating":
+        return <Spinner size="small" labelPosition="above" label={`Creating container "${containerName}"...`} />;
+      case "importing":
+        return <Spinner size="small" labelPosition="above" label={`Importing data into "${containerName}"...`} />;
+      case "completed":
+        return `Successfully created "${containerName}" with sample data.`;
+      case "error":
+        return (
+          <div style={{ color: "red" }}>
+            <div>Error: {errorMessage}</div>
+          </div>
+        );
+    }
+  };
+
+  const getButtonLabel = () => {
+    switch (status) {
+      case "idle":
+        return "Start";
+      case "creating":
+      case "importing":
+        return "Close";
+      case "completed":
+        return "Close";
+      case "error":
+        return "Close";
+    }
+  };
+
+  return (
+    <Dialog open={props.open} onOpenChange={(event, data) => props.setOpen(data.open)}>
+      <DialogSurface>
+        <DialogBody>
+          <DialogTitle>Sample Data</DialogTitle>
+          <DialogContent>
+            <div className={styles.dialogContent}>{renderContent()}</div>
+          </DialogContent>
+          <DialogActions>
+            <Button
+              appearance="primary"
+              onClick={handleActionOnClick}
+              disabled={status === "creating" || status === "importing"}
+            >
+              {getButtonLabel()}
+            </Button>
+          </DialogActions>
+        </DialogBody>
+      </DialogSurface>
+    </Dialog>
+  );
+};

src/Explorer/SplashScreen/SampleUtil.ts

@@ -0,0 +1,56 @@
+import { createCollection } from "Common/dataAccess/createCollection";
+import Explorer from "Explorer/Explorer";
+import { useDatabases } from "Explorer/useDatabases";
+import * as DataModels from "../../Contracts/DataModels";
+import * as ViewModels from "../../Contracts/ViewModels";
+
+/**
+ * Public for unit tests
+ * @param databaseName
+ * @param containerName
+ * @param containerDatabases
+ */
+const hasContainer = (
+  databaseName: string,
+  containerName: string,
+  containerDatabases: ViewModels.Database[],
+): boolean => {
+  const filteredDatabases = containerDatabases.filter((database) => database.id() === databaseName);
+  return (
+    filteredDatabases.length > 0 &&
+    filteredDatabases[0].collections().filter((collection) => collection.id() === containerName).length > 0
+  );
+};
+
+export const checkContainerExists = (databaseName: string, containerName: string) =>
+  hasContainer(databaseName, containerName, useDatabases.getState().databases);
+
+export const createContainer = async (
+  databaseName: string,
+  containerName: string,
+  explorer: Explorer,
+): Promise<ViewModels.Collection> => {
+  const createRequest: DataModels.CreateCollectionParams = {
+    createNewDatabase: false,
+    collectionId: containerName,
+    databaseId: databaseName,
+    databaseLevelThroughput: false,
+  };
+  await createCollection(createRequest);
+  await explorer.refreshAllDatabases();
+  const database = useDatabases.getState().findDatabaseWithId(databaseName);
+  if (!database) {
+    return undefined;
+  }
+  await database.loadCollections();
+  const newCollection = database.findCollectionWithId(containerName);
+  return newCollection;
+};
+
+export const importData = async (collection: ViewModels.Collection): Promise<void> => {
+  // TODO: keep same chunk as ContainerSampleGenerator
+  const dataFileContent = await import(
+    /* webpackChunkName: "queryCopilotSampleData" */ "../../../sampleData/queryCopilotSampleData.json"
+  );
+  await collection.bulkInsertDocuments(dataFileContent.data);
+};

src/Explorer/Tree/Collection.ts

@@ -1,4 +1,10 @@
-import { Resource, StoredProcedureDefinition, TriggerDefinition, UserDefinedFunctionDefinition } from "@azure/cosmos";
+import {
+  JSONObject,
+  Resource,
+  StoredProcedureDefinition,
+  TriggerDefinition,
+  UserDefinedFunctionDefinition,
+} from "@azure/cosmos";
 import { useNotebook } from "Explorer/Notebook/useNotebook";
 import { DocumentsTabV2 } from "Explorer/Tabs/DocumentsTabV2/DocumentsTabV2";
 import { isFabricMirrored } from "Platform/Fabric/FabricUtil";
@@ -1086,6 +1092,56 @@ export default class Collection implements ViewModels.Collection {
     });
   }
 
+  public async bulkInsertDocuments(documents: JSONObject[]): Promise<{
+    numSucceeded: number;
+    numFailed: number;
+    numThrottled: number;
+    errors: string[];
+  }> {
+    const stats = {
+      numSucceeded: 0,
+      numFailed: 0,
+      numThrottled: 0,
+      errors: [] as string[],
+    };
+
+    const chunkSize = 100; // 100 is the max # of bulk operations the SDK currently accepts
+    const chunkedContent = Array.from({ length: Math.ceil(documents.length / chunkSize) }, (_, index) =>
+      documents.slice(index * chunkSize, index * chunkSize + chunkSize),
+    );
+    for (const chunk of chunkedContent) {
+      let retryAttempts = 0;
+      let chunkComplete = false;
+      let documentsToAttempt = chunk;
+      while (retryAttempts < 10 && !chunkComplete) {
+        const responses = await bulkCreateDocument(this, documentsToAttempt);
+        const attemptedDocuments = [...documentsToAttempt];
+        documentsToAttempt = [];
+        responses.forEach((response, index) => {
+          if (response.statusCode === 201) {
+            stats.numSucceeded++;
+          } else if (response.statusCode === 429) {
+            documentsToAttempt.push(attemptedDocuments[index]);
+          } else {
+            stats.numFailed++;
+            stats.errors.push(JSON.stringify(response.resourceBody));
+          }
+        });
+        if (documentsToAttempt.length === 0) {
+          chunkComplete = true;
+          break;
+        }
+        logConsoleInfo(
+          `${documentsToAttempt.length} document creations were throttled. Waiting ${retryAttempts} seconds and retrying throttled documents`,
+        );
+        retryAttempts++;
+        await sleep(retryAttempts);
+      }
+    }
+
+    return stats;
+  }
+
   private async _createDocumentsFromFile(fileName: string, documentContent: string): Promise<UploadDetailsRecord> {
     const record: UploadDetailsRecord = {
       fileName: fileName,
@@ -1098,38 +1154,11 @@ export default class Collection implements ViewModels.Collection {
     try {
       const parsedContent = JSON.parse(documentContent);
       if (Array.isArray(parsedContent)) {
-        const chunkSize = 100; // 100 is the max # of bulk operations the SDK currently accepts
-        const chunkedContent = Array.from({ length: Math.ceil(parsedContent.length / chunkSize) }, (_, index) =>
-          parsedContent.slice(index * chunkSize, index * chunkSize + chunkSize),
-        );
-        for (const chunk of chunkedContent) {
-          let retryAttempts = 0;
-          let chunkComplete = false;
-          let documentsToAttempt = chunk;
-          while (retryAttempts < 10 && !chunkComplete) {
-            const responses = await bulkCreateDocument(this, documentsToAttempt);
-            const attemptedDocuments = [...documentsToAttempt];
-            documentsToAttempt = [];
-            responses.forEach((response, index) => {
-              if (response.statusCode === 201) {
-                record.numSucceeded++;
-              } else if (response.statusCode === 429) {
-                documentsToAttempt.push(attemptedDocuments[index]);
-              } else {
-                record.numFailed++;
-              }
-            });
-            if (documentsToAttempt.length === 0) {
-              chunkComplete = true;
-              break;
-            }
-            logConsoleInfo(
-              `${documentsToAttempt.length} document creations were throttled. Waiting ${retryAttempts} seconds and retrying throttled documents`,
-            );
-            retryAttempts++;
-            await sleep(retryAttempts);
-          }
-        }
+        const { numSucceeded, numFailed, numThrottled, errors } = await this.bulkInsertDocuments(parsedContent);
+        record.numSucceeded = numSucceeded;
+        record.numFailed = numFailed;
+        record.numThrottled = numThrottled;
+        record.errors = errors;
       } else {
         await createDocument(this, parsedContent);
         record.numSucceeded++;