From 1d6c0bbd1eabe39cb9d27589f9f902d30e4e7978 Mon Sep 17 00:00:00 2001
From: victor-meng <56978073+victor-meng@users.noreply.github.com>
Date: Tue, 11 Jul 2023 11:59:26 -0700
Subject: [PATCH] Fix issue with executing queries in the query copilot tab
 (#1522)

---
 src/Explorer/QueryCopilot/QueryCopilotTab.tsx | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/Explorer/QueryCopilot/QueryCopilotTab.tsx b/src/Explorer/QueryCopilot/QueryCopilotTab.tsx
index 594d3d1b2..28857db4d 100644
--- a/src/Explorer/QueryCopilot/QueryCopilotTab.tsx
+++ b/src/Explorer/QueryCopilot/QueryCopilotTab.tsx
@@ -1,5 +1,5 @@
 /* eslint-disable no-console */
-import { FeedOptions } from "@azure/cosmos";
+import { FeedOptions, ItemDefinition, QueryIterator, Resource } from "@azure/cosmos";
 import {
   Callout,
   CommandBarButton,
@@ -25,7 +25,8 @@ import {
 import { getErrorMessage, handleError } from "Common/ErrorHandlingUtils";
 import { shouldEnableCrossPartitionKey } from "Common/HeadersUtility";
 import { MinimalQueryIterator } from "Common/IteratorUtilities";
-import { queryDocuments } from "Common/dataAccess/queryDocuments";
+import { sampleDataClient } from "Common/SampleDataClient";
+import { getCommonQueryOptions } from "Common/dataAccess/queryDocuments";
 import { queryDocumentsPage } from "Common/dataAccess/queryDocumentsPage";
 import { QueryResults } from "Contracts/ViewModels";
 import { CommandButtonComponentProps } from "Explorer/Controls/CommandButton/CommandButtonComponent";
@@ -199,9 +200,17 @@ export const QueryCopilotTab: React.FC<QueryCopilotTabProps> = ({
     }
   };
 
+  const querySampleDocuments = (query: string, options: FeedOptions): QueryIterator<ItemDefinition & Resource> => {
+    options = getCommonQueryOptions(options);
+    return sampleDataClient()
+      .database(QueryCopilotSampleDatabaseId)
+      .container(QueryCopilotSampleContainerId)
+      .items.query(query, options);
+  };
+
   const onExecuteQueryClick = async (): Promise<void> => {
     const queryToExecute = selectedQuery || query;
-    const queryIterator = queryDocuments(QueryCopilotSampleDatabaseId, QueryCopilotSampleContainerId, queryToExecute, {
+    const queryIterator = querySampleDocuments(queryToExecute, {
       enableCrossPartitionQuery: shouldEnableCrossPartitionKey(),
     } as FeedOptions);
     setQueryIterator(queryIterator);