Copilot settings (#1719)

* gating copilot with AFEC

* Add enable sample DB in settings

* Add enable sample DB in settings

* Add enable sample DB in settings

* PR comments

less/documentDB.less

@@ -2897,9 +2897,21 @@ a:link {
     padding-left: 8px;
   }
 
+  .settingsSectionInlineCheckbox {
+    display: flex;
+    flex-direction: row-reverse;
+    justify-content: flex-end;
+    align-items: center;
+    gap: 5px;
+  }
+
   .settingsSectionLabel {
     margin-bottom: @DefaultSpace;
     margin-right: 5px;
+
+    .panelInfoIcon {
+      margin-left: 5px;
+    }
   }
 
   .pageOptionsPart {

src/Explorer/Explorer.tsx

@@ -3,9 +3,10 @@ import { isPublicInternetAccessAllowed } from "Common/DatabaseAccountUtility";
 import { sendMessage } from "Common/MessageHandler";
 import { Platform, configContext } from "ConfigContext";
 import { MessageTypes } from "Contracts/ExplorerContracts";
-import { getCopilotEnabled } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
+import { getCopilotEnabled, isCopilotFeatureRegistered } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
 import { IGalleryItem } from "Juno/JunoClient";
 import { requestDatabaseResourceTokens } from "Platform/Fabric/FabricUtil";
+import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
 import { allowedNotebookServerUrls, validateEndpoint } from "Utils/EndpointValidation";
 import { useQueryCopilot } from "hooks/useQueryCopilot";
 import * as ko from "knockout";
@@ -1389,9 +1390,18 @@ export default class Explorer {
     if (userContext.apiType !== "SQL" || !userContext.subscriptionId) {
       return;
     }
-    const copilotEnabled = await getCopilotEnabled();
+    const copilotEnabledPromise = getCopilotEnabled();
-    useQueryCopilot.getState().setCopilotEnabled(copilotEnabled);
+    const copilotUserDBEnabledPromise = isCopilotFeatureRegistered(userContext.subscriptionId);
-    useQueryCopilot.getState().setCopilotUserDBEnabled(copilotEnabled);
+    const [copilotEnabled, copilotUserDBEnabled] = await Promise.all([
+      copilotEnabledPromise,
+      copilotUserDBEnabledPromise,
+    ]);
+    const copilotSampleDBEnabled = LocalStorageUtility.getEntryString(StorageKey.CopilotSampleDBEnabled) === "true";
+    useQueryCopilot.getState().setCopilotEnabled(copilotEnabled && copilotUserDBEnabled);
+    useQueryCopilot.getState().setCopilotUserDBEnabled(copilotUserDBEnabled);
+    useQueryCopilot
+      .getState()
+      .setCopilotSampleDBEnabled(copilotEnabled && copilotUserDBEnabled && copilotSampleDBEnabled);
   }
 
   public async refreshSampleData(): Promise<void> {

src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx

@@ -200,7 +200,7 @@ export function createControlCommandBarButtons(container: Explorer): CommandButt
     {
       iconSrc: SettingsIcon,
       iconAlt: "Settings",
-      onCommandClick: () => useSidePanel.getState().openSidePanel("Settings", <SettingsPane />),
+      onCommandClick: () => useSidePanel.getState().openSidePanel("Settings", <SettingsPane explorer={container} />),
       commandButtonLabel: undefined,
       ariaLabel: "Settings",
       tooltipText: "Settings",

src/Explorer/OpenActions/OpenActions.tsx

@@ -154,7 +154,7 @@ function openPane(action: ActionContracts.OpenPane, explorer: Explorer) {
     action.paneKind === ActionContracts.PaneKind.GlobalSettings ||
     action.paneKind === ActionContracts.PaneKind[ActionContracts.PaneKind.GlobalSettings]
   ) {
-    useSidePanel.getState().openSidePanel("Settings", <SettingsPane />);
+    useSidePanel.getState().openSidePanel("Settings", <SettingsPane explorer={explorer} />);
   }
 }
 

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

@@ -6,7 +6,7 @@ import { SettingsPane } from "./SettingsPane";
 
 describe("Settings Pane", () => {
   it("should render Default properly", () => {
-    const wrapper = shallow(<SettingsPane />);
+    const wrapper = shallow(<SettingsPane explorer={null} />);
     expect(wrapper).toMatchSnapshot();
   });
 
@@ -18,7 +18,7 @@ describe("Settings Pane", () => {
         },
       } as DatabaseAccount,
     });
-    const wrapper = shallow(<SettingsPane />);
+    const wrapper = shallow(<SettingsPane explorer={null} />);
     expect(wrapper).toMatchSnapshot();
   });
 });

src/Explorer/Panes/SettingsPane/SettingsPane.tsx

@@ -16,13 +16,20 @@ import * as StringUtility from "Shared/StringUtility";
 import { userContext } from "UserContext";
 import { logConsoleInfo } from "Utils/NotificationConsoleUtils";
 import * as PriorityBasedExecutionUtils from "Utils/PriorityBasedExecutionUtils";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
 import { useSidePanel } from "hooks/useSidePanel";
 import React, { FunctionComponent, useState } from "react";
+import Explorer from "../../Explorer";
 import { RightPaneForm, RightPaneFormProps } from "../RightPaneForm/RightPaneForm";
 
-export const SettingsPane: FunctionComponent = () => {
+export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
+  explorer,
+}: {
+  explorer: Explorer;
+}): JSX.Element => {
   const closeSidePanel = useSidePanel((state) => state.closeSidePanel);
   const [isExecuting, setIsExecuting] = useState<boolean>(false);
+  const [refreshExplorer, setRefreshExplorer] = useState<boolean>(false);
   const [pageOption, setPageOption] = useState<string>(
     LocalStorageUtility.getEntryNumber(StorageKey.ActualItemPerPage) === Constants.Queries.unlimitedItemsPerPage
       ? Constants.Queries.UnlimitedPageOption
@@ -78,13 +85,17 @@ export const SettingsPane: FunctionComponent = () => {
       ? LocalStorageUtility.getEntryString(StorageKey.PriorityLevel)
       : Constants.PriorityLevel.Default,
   );
+  const [copilotSampleDBEnabled, setCopilotSampleDBEnabled] = useState<boolean>(
+    LocalStorageUtility.getEntryString(StorageKey.CopilotSampleDBEnabled) === "true",
+  );
   const explorerVersion = configContext.gitSha;
   const shouldShowQueryPageOptions = userContext.apiType === "SQL";
   const shouldShowGraphAutoVizOption = userContext.apiType === "Gremlin";
   const shouldShowCrossPartitionOption = userContext.apiType !== "Gremlin";
   const shouldShowParallelismOption = userContext.apiType !== "Gremlin";
   const shouldShowPriorityLevelOption = PriorityBasedExecutionUtils.isFeatureEnabled();
-  const handlerOnSubmit = () => {
+  const shouldShowCopilotSampleDBOption = userContext.apiType === "SQL" && useQueryCopilot.getState().copilotEnabled;
+  const handlerOnSubmit = async () => {
     setIsExecuting(true);
 
     LocalStorageUtility.setEntryNumber(
@@ -100,6 +111,7 @@ export const SettingsPane: FunctionComponent = () => {
     LocalStorageUtility.setEntryString(StorageKey.IsCrossPartitionQueryEnabled, crossPartitionQueryEnabled.toString());
     LocalStorageUtility.setEntryNumber(StorageKey.MaxDegreeOfParellism, maxDegreeOfParallelism);
     LocalStorageUtility.setEntryString(StorageKey.PriorityLevel, priorityLevel.toString());
+    LocalStorageUtility.setEntryString(StorageKey.CopilotSampleDBEnabled, copilotSampleDBEnabled.toString());
 
     if (shouldShowGraphAutoVizOption) {
       LocalStorageUtility.setEntryBoolean(
@@ -139,6 +151,7 @@ export const SettingsPane: FunctionComponent = () => {
     logConsoleInfo(
       `Updated query setting to ${LocalStorageUtility.getEntryString(StorageKey.SetPartitionKeyUndefined)}`,
     );
+    refreshExplorer && (await explorer.refreshExplorer());
     closeSidePanel();
   };
 
@@ -218,6 +231,12 @@ export const SettingsPane: FunctionComponent = () => {
     }
   };
 
+  const handleSampleDatabaseChange = async (ev: React.MouseEvent<HTMLElement>, checked?: boolean): Promise<void> => {
+    setCopilotSampleDBEnabled(checked);
+    useQueryCopilot.getState().setCopilotSampleDBEnabled(checked);
+    setRefreshExplorer(!refreshExplorer);
+  };
+
   const choiceButtonStyles = {
     root: {
       clear: "both",
@@ -434,7 +453,7 @@ export const SettingsPane: FunctionComponent = () => {
           </div>
         </div>
         <div className="settingsSection">
-          <div className="settingsSectionPart">
+          <div className="settingsSectionPart settingsSectionInlineCheckbox">
             <div className="settingsSectionLabel">
               Enable container pagination
               <InfoTooltip>
@@ -454,7 +473,7 @@ export const SettingsPane: FunctionComponent = () => {
         </div>
         {shouldShowCrossPartitionOption && (
           <div className="settingsSection">
-            <div className="settingsSectionPart">
+            <div className="settingsSectionPart settingsSectionInlineCheckbox">
               <div className="settingsSectionLabel">
                 Enable cross-partition query
                 <InfoTooltip>
@@ -545,6 +564,30 @@ export const SettingsPane: FunctionComponent = () => {
             </div>
           </div>
         )}
+        {shouldShowCopilotSampleDBOption && (
+          <div className="settingsSection">
+            <div className="settingsSectionPart settingsSectionInlineCheckbox">
+              <div className="settingsSectionLabel">
+                Enable sample database
+                <InfoTooltip>
+                  This is a sample database and collection with synthetic product data you can use to explore using
+                  NoSQL queries and Copilot. This will appear as another database in the Data Explorer UI, and is
+                  created by, and maintained by Microsoft at no cost to you.
+                </InfoTooltip>
+              </div>
+
+              <Checkbox
+                styles={{
+                  label: { padding: 0 },
+                }}
+                className="padding"
+                ariaLabel="Enable sample db for Copilot"
+                checked={copilotSampleDBEnabled}
+                onChange={handleSampleDatabaseChange}
+              />
+            </div>
+          </div>
+        )}
         <div className="settingsSection">
           <div className="settingsSectionPart">
             <div className="settingsSectionLabel">Explorer Version</div>

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

@@ -274,7 +274,7 @@ exports[`Settings Pane should render Default properly 1`] = `
       className="settingsSection"
     >
       <div
-        className="settingsSectionPart"
+        className="settingsSectionPart settingsSectionInlineCheckbox"
       >
         <div
           className="settingsSectionLabel"
@@ -303,7 +303,7 @@ exports[`Settings Pane should render Default properly 1`] = `
       className="settingsSection"
     >
       <div
-        className="settingsSectionPart"
+        className="settingsSectionPart settingsSectionInlineCheckbox"
       >
         <div
           className="settingsSectionLabel"
@@ -521,7 +521,7 @@ exports[`Settings Pane should render Gremlin properly 1`] = `
       className="settingsSection"
     >
       <div
-        className="settingsSectionPart"
+        className="settingsSectionPart settingsSectionInlineCheckbox"
       >
         <div
           className="settingsSectionLabel"

src/Explorer/QueryCopilot/Shared/QueryCopilotClient.ts

@@ -15,7 +15,12 @@ import { MinimalQueryIterator } from "Common/IteratorUtilities";
 import { createUri } from "Common/UrlUtility";
 import { queryDocumentsPage } from "Common/dataAccess/queryDocumentsPage";
 import { configContext } from "ConfigContext";
-import { ContainerConnectionInfo, CopilotEnabledConfiguration, IProvisionData } from "Contracts/DataModels";
+import {
+  ContainerConnectionInfo,
+  CopilotEnabledConfiguration,
+  FeatureRegistration,
+  IProvisionData,
+} from "Contracts/DataModels";
 import { AuthorizationTokenHeaderMetadata, QueryResults } from "Contracts/ViewModels";
 import { useDialog } from "Explorer/Controls/Dialog";
 import Explorer from "Explorer/Explorer";
@@ -52,6 +57,28 @@ async function fetchWithTimeout(
   return response;
 }
 
+export const isCopilotFeatureRegistered = async (subscriptionId: string): Promise<boolean> => {
+  const api_version = "2021-07-01";
+  const url = `${configContext.ARM_ENDPOINT}/subscriptions/${subscriptionId}/providers/Microsoft.Features/featureProviders/Microsoft.DocumentDB/subscriptionFeatureRegistrations/MicrosoftCopilotForAzureInCDB?api-version=${api_version}`;
+  const authorizationHeader: AuthorizationTokenHeaderMetadata = getAuthorizationHeader();
+  const headers = { [authorizationHeader.header]: authorizationHeader.token };
+
+  let response;
+
+  try {
+    response = await fetchWithTimeout(url, headers);
+  } catch (error) {
+    return false;
+  }
+
+  if (!response?.ok) {
+    return false;
+  }
+
+  const featureRegistration = (await response?.json()) as FeatureRegistration;
+  return featureRegistration?.properties?.state === "Registered";
+};
+
 export const getCopilotEnabled = async (): Promise<boolean> => {
   const url = `${configContext.BACKEND_ENDPOINT}/api/portalsettings/querycopilot`;
   const authorizationHeader: AuthorizationTokenHeaderMetadata = getAuthorizationHeader();

src/Explorer/SplashScreen/SplashScreen.tsx

@@ -148,6 +148,7 @@ export class SplashScreen extends React.Component<SplashScreenProps> {
             />
           </Stack>
           <Stack horizontal tokens={{ childrenGap: 16 }}>
+            {useQueryCopilot.getState().copilotEnabled && useQueryCopilot.getState().copilotSampleDBEnabled && (
               <SplashScreenButton
                 imgSrc={CopilotIcon}
                 title={"Query faster with Copilot"}
@@ -165,6 +166,7 @@ export class SplashScreen extends React.Component<SplashScreenProps> {
                   traceOpen(Action.OpenQueryCopilotFromSplashScreen, { apiType: userContext.apiType });
                 }}
               />
+            )}
             <SplashScreenButton
               imgSrc={ConnectIcon}
               title={"Connect"}

src/Explorer/Tree/ResourceTree.tsx

@@ -769,7 +769,10 @@ export const ResourceTree: React.FC<ResourceTreeProps> = ({ container }: Resourc
 
   const dataRootNode = buildDataTree();
   const isSampleDataEnabled =
-    useQueryCopilot().copilotEnabled && userContext.sampleDataConnectionInfo && userContext.apiType === "SQL";
+    useQueryCopilot().copilotEnabled &&
+    useQueryCopilot().copilotSampleDBEnabled &&
+    userContext.sampleDataConnectionInfo &&
+    userContext.apiType === "SQL";
   const sampleDataResourceTokenCollection = useDatabases((state) => state.sampleDataResourceTokenCollection);
 
   return (

src/Shared/StorageUtility.ts

@@ -11,6 +11,7 @@ export enum StorageKey {
   MaxWaitTimeInSeconds,
   AutomaticallyCancelQueryAfterTimeout,
   ContainerPaginationEnabled,
+  CopilotSampleDBEnabled,
   CustomItemPerPage,
   DatabaseAccountId,
   EncryptedKeyToken,

src/hooks/useQueryCopilot.ts

@@ -10,6 +10,7 @@ import { ContainerInfo } from "../Contracts/DataModels";
 export interface QueryCopilotState {
   copilotEnabled: boolean;
   copilotUserDBEnabled: boolean;
+  copilotSampleDBEnabled: boolean;
   generatedQuery: string;
   likeQuery: boolean;
   userPrompt: string;
@@ -50,6 +51,7 @@ export interface QueryCopilotState {
 
   setCopilotEnabled: (copilotEnabled: boolean) => void;
   setCopilotUserDBEnabled: (copilotUserDBEnabled: boolean) => void;
+  setCopilotSampleDBEnabled: (copilotSampleDBEnabled: boolean) => void;
   openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) => void;
   closeFeedbackModal: () => void;
   setHideFeedbackModalForLikedQueries: (hideFeedbackModalForLikedQueries: boolean) => void;
@@ -96,6 +98,7 @@ type QueryCopilotStore = UseStore<QueryCopilotState>;
 export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   copilotEnabled: false,
   copilotUserDBEnabled: false,
+  copilotSampleDBEnabled: false,
   generatedQuery: "",
   likeQuery: false,
   userPrompt: "",
@@ -145,6 +148,7 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
 
   setCopilotEnabled: (copilotEnabled: boolean) => set({ copilotEnabled }),
   setCopilotUserDBEnabled: (copilotUserDBEnabled: boolean) => set({ copilotUserDBEnabled }),
+  setCopilotSampleDBEnabled: (copilotSampleDBEnabled: boolean) => set({ copilotSampleDBEnabled }),
   openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) =>
     set({ generatedQuery, likeQuery, userPrompt, showFeedbackModal: true }),
   closeFeedbackModal: () => set({ showFeedbackModal: false }),