Adjusted Extecute Query logic (#1605)

Co-authored-by: Predrag Klepic <v-prklepic@microsoft.com>

src/Explorer/QueryCopilot/QueryCopilotTab.tsx

@@ -1,5 +1,4 @@
 /* eslint-disable no-console */
-import { FeedOptions } from "@azure/cosmos";
 import {
   Callout,
   CommandBarButton,
@@ -20,16 +19,11 @@ import { useBoolean } from "@fluentui/react-hooks";
 import {
   ContainerStatusType,
   PoolIdType,
-  QueryCopilotSampleContainerId,
   QueryCopilotSampleContainerSchema,
   ShortenedQueryCopilotSampleContainerSchema,
 } from "Common/Constants";
-import { getErrorMessage, handleError } from "Common/ErrorHandlingUtils";
-import { shouldEnableCrossPartitionKey } from "Common/HeadersUtility";
-import { MinimalQueryIterator } from "Common/IteratorUtilities";
+import { handleError } from "Common/ErrorHandlingUtils";
 import { createUri } from "Common/UrlUtility";
-import { queryDocumentsPage } from "Common/dataAccess/queryDocumentsPage";
-import { QueryResults } from "Contracts/ViewModels";
 import { CommandButtonComponentProps } from "Explorer/Controls/CommandButton/CommandButtonComponent";
 import { EditorReact } from "Explorer/Controls/Editor/EditorReact";
 import { useCommandBar } from "Explorer/Menus/CommandBar/CommandBarComponentAdapter";
@@ -37,15 +31,11 @@ import { SaveQueryPane } from "Explorer/Panes/SaveQueryPane/SaveQueryPane";
 import { WelcomeModal } from "Explorer/QueryCopilot/Modal/WelcomeModal";
 import { CopyPopup } from "Explorer/QueryCopilot/Popup/CopyPopup";
 import { DeletePopup } from "Explorer/QueryCopilot/Popup/DeletePopup";
-import { querySampleDocuments } from "Explorer/QueryCopilot/QueryCopilotUtilities";
-import { SubmitFeedback } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
+import { OnExecuteQueryClick, SubmitFeedback } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
 import { GenerateSQLQueryResponse, QueryCopilotProps } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces";
+import { QueryCopilotResults } from "Explorer/QueryCopilot/Shared/QueryCopilotResults";
 import { SamplePrompts, SamplePromptsProps } from "Explorer/QueryCopilot/Shared/SamplePrompts/SamplePrompts";
-import { QueryResultSection } from "Explorer/Tabs/QueryTab/QueryResultSection";
-import { Action } from "Shared/Telemetry/TelemetryConstants";
-import { traceFailure, traceStart, traceSuccess } from "Shared/Telemetry/TelemetryProcessor";
 import { userContext } from "UserContext";
-import { queryPagesUntilContentPresent } from "Utils/QueryUtils";
 import { useQueryCopilot } from "hooks/useQueryCopilot";
 import { useSidePanel } from "hooks/useSidePanel";
 import React, { useRef, useState } from "react";
@@ -83,8 +73,6 @@ export const QueryCopilotTab: React.FC<QueryCopilotProps> = ({ explorer }: Query
     setSelectedQuery,
     isGeneratingQuery,
     setIsGeneratingQuery,
-    isExecuting,
-    setIsExecuting,
     likeQuery,
     setLikeQuery,
     dislikeQuery,
@@ -93,12 +81,6 @@ export const QueryCopilotTab: React.FC<QueryCopilotProps> = ({ explorer }: Query
     setShowCallout,
     showSamplePrompts,
     setShowSamplePrompts,
-    queryIterator,
-    setQueryIterator,
-    queryResults,
-    setQueryResults,
-    errorMessage,
-    setErrorMessage,
     isSamplePromptsOpen,
     setIsSamplePromptsOpen,
     showDeletePopup,
@@ -109,7 +91,6 @@ export const QueryCopilotTab: React.FC<QueryCopilotProps> = ({ explorer }: Query
     setshowCopyPopup,
     showErrorMessageBar,
     setShowErrorMessageBar,
-    generatedQueryComments,
     setGeneratedQueryComments,
   } = useQueryCopilot();
 
@@ -238,64 +219,12 @@ export const QueryCopilotTab: React.FC<QueryCopilotProps> = ({ explorer }: Query
     }
   };
 
-  const onExecuteQueryClick = async (): Promise<void> => {
-    traceStart(Action.ExecuteQueryGeneratedFromQueryCopilot, {
-      correlationId: useQueryCopilot.getState().correlationId,
-      userPrompt: userPrompt,
-      generatedQuery: generatedQuery,
-      generatedQueryComments: generatedQueryComments,
-      executedQuery: selectedQuery || query,
-    });
-    const queryToExecute = selectedQuery || query;
-    const queryIterator = querySampleDocuments(queryToExecute, {
-      enableCrossPartitionQuery: shouldEnableCrossPartitionKey(),
-    } as FeedOptions);
-    setQueryIterator(queryIterator);
-
-    setTimeout(async () => {
-      await queryDocumentsPerPage(0, queryIterator);
-    }, 100);
-  };
-
-  const queryDocumentsPerPage = async (firstItemIndex: number, queryIterator: MinimalQueryIterator): Promise<void> => {
-    try {
-      setIsExecuting(true);
-      useTabs.getState().setIsTabExecuting(true);
-      useTabs.getState().setIsQueryErrorThrown(false);
-      const queryResults: QueryResults = await queryPagesUntilContentPresent(
-        firstItemIndex,
-        async (firstItemIndex: number) =>
-          queryDocumentsPage(QueryCopilotSampleContainerId, queryIterator, firstItemIndex)
-      );
-
-      setQueryResults(queryResults);
-      setErrorMessage("");
-      setShowErrorMessageBar(false);
-      traceSuccess(Action.ExecuteQueryGeneratedFromQueryCopilot, {
-        correlationId: useQueryCopilot.getState().correlationId,
-      });
-    } catch (error) {
-      const errorMessage = getErrorMessage(error);
-      traceFailure(Action.ExecuteQueryGeneratedFromQueryCopilot, {
-        correlationId: useQueryCopilot.getState().correlationId,
-        errorMessage: errorMessage,
-      });
-      setErrorMessage(errorMessage);
-      handleError(errorMessage, "executeQueryCopilotTab");
-      useTabs.getState().setIsQueryErrorThrown(true);
-      setShowErrorMessageBar(true);
-    } finally {
-      setIsExecuting(false);
-      useTabs.getState().setIsTabExecuting(false);
-    }
-  };
-
   const getCommandbarButtons = (): CommandButtonComponentProps[] => {
     const executeQueryBtnLabel = selectedQuery ? "Execute Selection" : "Execute Query";
     const executeQueryBtn = {
       iconSrc: ExecuteQueryIcon,
       iconAlt: executeQueryBtnLabel,
-      onCommandClick: () => onExecuteQueryClick(),
+      onCommandClick: () => OnExecuteQueryClick(),
       commandButtonLabel: executeQueryBtnLabel,
       ariaLabel: executeQueryBtnLabel,
       hasPopup: false,
@@ -622,16 +551,7 @@ export const QueryCopilotTab: React.FC<QueryCopilotProps> = ({ explorer }: Query
               onContentChanged={(newQuery: string) => setQuery(newQuery)}
               onContentSelected={(selectedQuery: string) => setSelectedQuery(selectedQuery)}
             />
-            <QueryResultSection
-              isMongoDB={false}
-              queryEditorContent={selectedQuery || query}
-              error={errorMessage}
-              queryResults={queryResults}
-              isExecuting={isExecuting}
-              executeQueryDocumentsPage={(firstItemIndex: number) =>
-                queryDocumentsPerPage(firstItemIndex, queryIterator)
-              }
-            />
+            <QueryCopilotResults />
           </SplitterLayout>
         </Stack>
         <WelcomeModal visible={localStorage.getItem("hideWelcomeModal") !== "true"} />

src/Explorer/QueryCopilot/Shared/QueryCopilotClient.ts

@@ -1,14 +1,24 @@
+import { FeedOptions } from "@azure/cosmos";
 import {
   ContainerStatusType,
   PoolIdType,
+  QueryCopilotSampleContainerId,
   QueryCopilotSampleContainerSchema,
   ShortenedQueryCopilotSampleContainerSchema,
 } from "Common/Constants";
-import { handleError } from "Common/ErrorHandlingUtils";
+import { getErrorMessage, handleError } from "Common/ErrorHandlingUtils";
+import { shouldEnableCrossPartitionKey } from "Common/HeadersUtility";
+import { MinimalQueryIterator } from "Common/IteratorUtilities";
 import { createUri } from "Common/UrlUtility";
+import { queryDocumentsPage } from "Common/dataAccess/queryDocumentsPage";
+import { QueryResults } from "Contracts/ViewModels";
 import Explorer from "Explorer/Explorer";
+import { querySampleDocuments } from "Explorer/QueryCopilot/QueryCopilotUtilities";
 import { FeedbackParams, GenerateSQLQueryResponse } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces";
+import { Action } from "Shared/Telemetry/TelemetryConstants";
+import { traceFailure, traceStart, traceSuccess } from "Shared/Telemetry/TelemetryProcessor";
 import { userContext } from "UserContext";
+import { queryPagesUntilContentPresent } from "Utils/QueryUtils";
 import { useQueryCopilot } from "hooks/useQueryCopilot";
 import { useTabs } from "hooks/useTabs";
 
@@ -133,3 +143,57 @@ export const SubmitFeedback = async ({
     handleError(error, "copilotSubmitFeedback");
   }
 };
+
+export const OnExecuteQueryClick = async (): Promise<void> => {
+  traceStart(Action.ExecuteQueryGeneratedFromQueryCopilot, {
+    correlationId: useQueryCopilot.getState().correlationId,
+    userPrompt: useQueryCopilot.getState().userPrompt,
+    generatedQuery: useQueryCopilot.getState().generatedQuery,
+    generatedQueryComments: useQueryCopilot.getState().generatedQueryComments,
+    executedQuery: useQueryCopilot.getState().selectedQuery || useQueryCopilot.getState().query,
+  });
+  const queryToExecute = useQueryCopilot.getState().selectedQuery || useQueryCopilot.getState().query;
+  const queryIterator = querySampleDocuments(queryToExecute, {
+    enableCrossPartitionQuery: shouldEnableCrossPartitionKey(),
+  } as FeedOptions);
+  useQueryCopilot.getState().setQueryIterator(queryIterator);
+
+  setTimeout(async () => {
+    await QueryDocumentsPerPage(0, queryIterator);
+  }, 100);
+};
+
+export const QueryDocumentsPerPage = async (
+  firstItemIndex: number,
+  queryIterator: MinimalQueryIterator
+): Promise<void> => {
+  try {
+    useQueryCopilot.getState().setIsExecuting(true);
+    useTabs.getState().setIsTabExecuting(true);
+    useTabs.getState().setIsQueryErrorThrown(false);
+    const queryResults: QueryResults = await queryPagesUntilContentPresent(
+      firstItemIndex,
+      async (firstItemIndex: number) => queryDocumentsPage(QueryCopilotSampleContainerId, queryIterator, firstItemIndex)
+    );
+
+    useQueryCopilot.getState().setQueryResults(queryResults);
+    useQueryCopilot.getState().setErrorMessage("");
+    useQueryCopilot.getState().setShowErrorMessageBar(false);
+    traceSuccess(Action.ExecuteQueryGeneratedFromQueryCopilot, {
+      correlationId: useQueryCopilot.getState().correlationId,
+    });
+  } catch (error) {
+    const errorMessage = getErrorMessage(error);
+    traceFailure(Action.ExecuteQueryGeneratedFromQueryCopilot, {
+      correlationId: useQueryCopilot.getState().correlationId,
+      errorMessage: errorMessage,
+    });
+    useQueryCopilot.getState().setErrorMessage(errorMessage);
+    handleError(errorMessage, "executeQueryCopilotTab");
+    useTabs.getState().setIsQueryErrorThrown(true);
+    useQueryCopilot.getState().setShowErrorMessageBar(true);
+  } finally {
+    useQueryCopilot.getState().setIsExecuting(false);
+    useTabs.getState().setIsTabExecuting(false);
+  }
+};

src/Explorer/QueryCopilot/Shared/QueryCopilotResults.tsx

@@ -0,0 +1,19 @@
+import { QueryDocumentsPerPage } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
+import { QueryResultSection } from "Explorer/Tabs/QueryTab/QueryResultSection";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
+import React from "react";
+
+export const QueryCopilotResults: React.FC = (): JSX.Element => {
+  return (
+    <QueryResultSection
+      isMongoDB={false}
+      queryEditorContent={useQueryCopilot.getState().selectedQuery || useQueryCopilot.getState().query}
+      error={useQueryCopilot.getState().errorMessage}
+      queryResults={useQueryCopilot.getState().queryResults}
+      isExecuting={useQueryCopilot.getState().isExecuting}
+      executeQueryDocumentsPage={(firstItemIndex: number) =>
+        QueryDocumentsPerPage(firstItemIndex, useQueryCopilot.getState().queryIterator)
+      }
+    />
+  );
+};

src/Explorer/QueryCopilot/__snapshots__/QueryCopilotTab.test.tsx.snap

@@ -125,25 +125,26 @@ exports[`Query copilot tab snapshot test should render with initial input 1`] =
       >
         <EditorReact
           ariaLabel="Editing Query"
-          content=""
+          content="SELECT * FROM c"
           isReadOnly={false}
           language="sql"
           lineNumbers="on"
           onContentChanged={[Function]}
           onContentSelected={[Function]}
         />
-        <QueryResultSection
-          error=""
-          executeQueryDocumentsPage={[Function]}
-          isExecuting={false}
-          isMongoDB={false}
-          queryEditorContent=""
-        />
+        <QueryCopilotResults />
       </t>
     </Stack>
     <WelcomeModal
       visible={true}
     />
+    <DeletePopup
+      clearFeedback={[Function]}
+      setQuery={[Function]}
+      setShowDeletePopup={[Function]}
+      showDeletePopup={false}
+      showFeedbackBar={[Function]}
+    />
     <CopyPopup
       setShowCopyPopup={[Function]}
       showCopyPopup={false}

src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx

@@ -1,4 +1,6 @@
 import { FeedOptions } from "@azure/cosmos";
+import { OnExecuteQueryClick } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
+import { QueryCopilotResults } from "Explorer/QueryCopilot/Shared/QueryCopilotResults";
 import { QueryCopilotSidebar } from "Explorer/QueryCopilot/V2/Sidebar/QueryCopilotSidebar";
 import { QueryResultSection } from "Explorer/Tabs/QueryTab/QueryResultSection";
 import { QueryCopilotState, useQueryCopilot } from "hooks/useQueryCopilot";
@@ -276,7 +278,7 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
       buttons.push({
         iconSrc: ExecuteQueryIcon,
         iconAlt: label,
-        onCommandClick: this.onExecuteQueryClick,
+        onCommandClick: this.isCopilotTabActive ? () => OnExecuteQueryClick() : this.onExecuteQueryClick,
         commandButtonLabel: label,
         ariaLabel: label,
         hasPopup: false,
@@ -365,11 +367,18 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
         selectedContent: "",
       });
     }
+
+    if (this.isCopilotTabActive) {
+      selectedContent.trim().length > 0
+        ? useQueryCopilot.getState().setSelectedQuery(selectedContent)
+        : useQueryCopilot.getState().setSelectedQuery("");
+    }
+
     useCommandBar.getState().setContextButtons(this.getTabsButtons());
   }
 
   public setEditorContent(): string {
-    if (this.state.queryCopilotGeneratedQuery) {
+    if (this.isCopilotTabActive && this.state.queryCopilotGeneratedQuery) {
       return this.state.queryCopilotGeneratedQuery;
     }
 
@@ -416,14 +425,20 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
                   />
                 </div>
               </Fragment>
-              <QueryResultSection
-                isMongoDB={this.props.isPreferredApiMongoDB}
-                queryEditorContent={this.state.sqlQueryEditorContent}
-                error={this.state.error}
-                queryResults={this.state.queryResults}
-                isExecuting={this.state.isExecuting}
-                executeQueryDocumentsPage={(firstItemIndex: number) => this._executeQueryDocumentsPage(firstItemIndex)}
-              />
+              {this.isCopilotTabActive ? (
+                <QueryCopilotResults />
+              ) : (
+                <QueryResultSection
+                  isMongoDB={this.props.isPreferredApiMongoDB}
+                  queryEditorContent={this.state.sqlQueryEditorContent}
+                  error={this.state.error}
+                  queryResults={this.state.queryResults}
+                  isExecuting={this.state.isExecuting}
+                  executeQueryDocumentsPage={(firstItemIndex: number) =>
+                    this._executeQueryDocumentsPage(firstItemIndex)
+                  }
+                />
+              )}
             </SplitterLayout>
           </div>
         </div>

src/hooks/useQueryCopilot.ts

@@ -88,7 +88,7 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   showFeedbackModal: false,
   hideFeedbackModalForLikedQueries: false,
   correlationId: "",
-  query: "",
+  query: "SELECT * FROM c",
   selectedQuery: "",
   isGeneratingQuery: false,
   isGeneratingExplanation: false,