copilot db exceptions plus zustand removal

src/Explorer/Tabs/QueryTab/QueryResultSection.tsx

@@ -3,18 +3,21 @@ import QueryError from "Common/QueryError";
 import { IndeterminateProgressBar } from "Explorer/Controls/IndeterminateProgressBar";
 import { MessageBanner } from "Explorer/Controls/MessageBanner";
 import { useQueryTabStyles } from "Explorer/Tabs/QueryTab/Styles";
+import useZoomLevel from "hooks/useZoomLevel";
 import React from "react";
+import { conditionalClass } from "Utils/StyleUtils";
 import RunQuery from "../../../../images/RunQuery.png";
 import { QueryResults } from "../../../Contracts/ViewModels";
 import { ErrorList } from "./ErrorList";
 import { ResultsView } from "./ResultsView";
-import useZoomLevel from "hooks/useZoomLevel";
-import { conditionalClass } from "Utils/StyleUtils";
 
 export interface ResultsViewProps {
   isMongoDB: boolean;
   queryResults: QueryResults;
   executeQueryDocumentsPage: (firstItemIndex: number) => Promise<void>;
+  queryEditorContent?: string;
+  databaseId?: string;
+  containerId?: string;
 }
 
 interface QueryResultProps extends ResultsViewProps {
@@ -49,6 +52,8 @@ export const QueryResultSection: React.FC<QueryResultProps> = ({
   queryResults,
   executeQueryDocumentsPage,
   isExecuting,
+  databaseId,
+  containerId,
 }: QueryResultProps): JSX.Element => {
   const styles = useQueryTabStyles();
   const maybeSubQuery = queryEditorContent && /.*\(.*SELECT.*\)/i.test(queryEditorContent);
@@ -91,6 +96,9 @@ export const QueryResultSection: React.FC<QueryResultProps> = ({
           queryResults={queryResults}
           executeQueryDocumentsPage={executeQueryDocumentsPage}
           isMongoDB={isMongoDB}
+          queryEditorContent={queryEditorContent}
+          databaseId={databaseId}
+          containerId={containerId}
         />
       ) : (
         <ExecuteQueryCallToAction />

src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx

@@ -375,9 +375,9 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
       this._iterator = this.props.isPreferredApiMongoDB
         ? queryIterator(this.props.collection.databaseId, this.props.viewModelcollection, query)
         : queryDocuments(this.props.collection.databaseId, this.props.collection.id(), query, {
-          enableCrossPartitionQuery: HeadersUtility.shouldEnableCrossPartitionKey(),
+            enableCrossPartitionQuery: HeadersUtility.shouldEnableCrossPartitionKey(),
-          abortSignal: this.queryAbortController.signal,
+            abortSignal: this.queryAbortController.signal,
-        } as unknown as FeedOptions);
+          } as unknown as FeedOptions);
     }
 
     await this._queryDocumentsPage(firstItemIndex);
@@ -794,6 +794,8 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
                   errors={this.props.copilotStore?.errors}
                   isExecuting={this.props.copilotStore?.isExecuting}
                   queryResults={this.props.copilotStore?.queryResults}
+                  databaseId={this.props.collection.databaseId}
+                  containerId={this.props.collection.id()}
                   executeQueryDocumentsPage={(firstItemIndex: number) =>
                     QueryDocumentsPerPage(
                       firstItemIndex,
@@ -809,6 +811,8 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
                   errors={this.state.errors}
                   isExecuting={this.state.isExecuting}
                   queryResults={this.state.queryResults}
+                  databaseId={this.props.collection.databaseId}
+                  containerId={this.props.collection.id()}
                   executeQueryDocumentsPage={(firstItemIndex: number) =>
                     this._executeQueryDocumentsPage(firstItemIndex)
                   }

src/Explorer/Tabs/QueryTab/ResultsView.test.tsx

@@ -9,17 +9,27 @@ const mockRead = jest.fn();
 const mockLogConsoleProgress = jest.fn();
 const mockHandleError = jest.fn();
 
-const indexMetricsString = `
+const indexMetricsResponse = {
-Utilized Single Indexes
+  UtilizedIndexes: {
-Index Spec: /foo/?
+    SingleIndexes: [{ IndexSpec: "/foo/?", IndexImpactScore: "High" }],
-Index Impact Score: High
+    CompositeIndexes: [{ IndexSpecs: ["/baz/? DESC", "/qux/? ASC"], IndexImpactScore: "Low" }],
-Potential Single Indexes
+  },
-Index Spec: /bar/?
+  PotentialIndexes: {
-Index Impact Score: Medium
+    SingleIndexes: [{ IndexSpec: "/bar/?", IndexImpactScore: "Medium" }],
-Utilized Composite Indexes
+    CompositeIndexes: [] as Array<{ IndexSpecs: string[]; IndexImpactScore?: string }>,
-Index Spec: /baz/? DESC, /qux/? ASC
+  },
-Index Impact Score: Low
+};
-`;
+
+const mockQueryResults = {
+  documents: [] as unknown[],
+  hasMoreResults: false,
+  itemCount: 0,
+  firstItemIndex: 0,
+  lastItemIndex: 0,
+  requestCharge: 0,
+  activityId: "test-activity-id",
+};
+
 mockRead.mockResolvedValue({
   resource: {
     indexingPolicy: {
@@ -42,20 +52,13 @@ mockReplace.mockResolvedValue({
   },
 });
 
-jest.mock("./QueryTabComponent", () => ({
-  useQueryMetadataStore: () => ({
-    userQuery: "SELECT * FROM c",
-    databaseId: "db1",
-    containerId: "col1",
-  }),
-}));
 jest.mock("Common/CosmosClient", () => ({
   client: () => ({
     database: () => ({
       container: () => ({
         items: {
           query: () => ({
-            fetchAll: mockFetchAll.mockResolvedValueOnce({ indexMetrics: indexMetricsString }),
+            fetchAll: mockFetchAll.mockResolvedValue({ indexMetrics: indexMetricsResponse }),
           }),
         },
         read: mockRead,
@@ -71,7 +74,7 @@ jest.mock("./StylesAdvisor", () => ({
 jest.mock("../../../Utils/NotificationConsoleUtils", () => ({
   logConsoleProgress: (...args: unknown[]) => {
     mockLogConsoleProgress(...args);
-    return () => {};
+    return () => { };
   },
 }));
 
@@ -82,18 +85,32 @@ jest.mock("../../../Common/ErrorHandlingUtils", () => {
 });
 
 test("logs progress message when fetching index metrics", async () => {
-  render(<IndexAdvisorTab />);
+  render(
+    <IndexAdvisorTab
+      queryResults={mockQueryResults}
+      queryEditorContent="SELECT * FROM c"
+      databaseId="db1"
+      containerId="col1"
+    />,
+  );
   await waitFor(() => expect(mockLogConsoleProgress).toHaveBeenCalledWith(expect.stringContaining("IndexMetrics")));
 });
 test("renders both Included and Not Included sections after loading", async () => {
-  render(<IndexAdvisorTab />);
+  render(
+    <IndexAdvisorTab
+      queryResults={mockQueryResults}
+      queryEditorContent="SELECT * FROM c"
+      databaseId="db1"
+      containerId="col1"
+    />,
+  );
   await waitFor(() => expect(screen.getByText("Included in Current Policy")).toBeInTheDocument());
   expect(screen.getByText("Not Included in Current Policy")).toBeInTheDocument();
   expect(screen.getByText("/foo/?")).toBeInTheDocument();
   expect(screen.getByText("/bar/?")).toBeInTheDocument();
 });
 test("shows update button only when an index is selected", async () => {
-  render(<IndexAdvisorTab />);
+  render(<IndexAdvisorTab queryEditorContent="SELECT * FROM c" databaseId="db1" containerId="col1" />);
   await waitFor(() => expect(screen.getByText("/bar/?")).toBeInTheDocument());
   const checkboxes = screen.getAllByRole("checkbox");
   expect(checkboxes.length).toBeGreaterThan(1);
@@ -104,7 +121,7 @@ test("shows update button only when an index is selected", async () => {
   expect(screen.queryByText(/Update Indexing Policy/)).not.toBeInTheDocument();
 });
 test("calls replace when update policy is confirmed", async () => {
-  render(<IndexAdvisorTab />);
+  render(<IndexAdvisorTab queryEditorContent="SELECT * FROM c" databaseId="db1" containerId="col1" />);
   await waitFor(() => expect(screen.getByText("/bar/?")).toBeInTheDocument());
   const checkboxes = screen.getAllByRole("checkbox");
   fireEvent.click(checkboxes[1]);
@@ -114,7 +131,7 @@ test("calls replace when update policy is confirmed", async () => {
 });
 
 test("calls replace when update button is clicked", async () => {
-  render(<IndexAdvisorTab />);
+  render(<IndexAdvisorTab queryEditorContent="SELECT * FROM c" databaseId="db1" containerId="col1" />);
   await waitFor(() => expect(screen.getByText("/bar/?")).toBeInTheDocument());
   const checkboxes = screen.getAllByRole("checkbox");
   fireEvent.click(checkboxes[1]); // Select /bar/?
@@ -123,14 +140,14 @@ test("calls replace when update button is clicked", async () => {
 });
 
 test("fetches indexing policy via read", async () => {
-  render(<IndexAdvisorTab />);
+  render(<IndexAdvisorTab queryEditorContent="SELECT * FROM c" databaseId="db1" containerId="col1" />);
   await waitFor(() => {
     expect(mockRead).toHaveBeenCalled();
   });
 });
 
 test("selects all indexes when select-all is clicked", async () => {
-  render(<IndexAdvisorTab />);
+  render(<IndexAdvisorTab queryEditorContent="SELECT * FROM c" databaseId="db1" containerId="col1" />);
   await waitFor(() => expect(screen.getByText("/bar/?")).toBeInTheDocument());
   const checkboxes = screen.getAllByRole("checkbox");
 
@@ -141,22 +158,22 @@ test("selects all indexes when select-all is clicked", async () => {
 });
 
 test("shows spinner while loading and hides after fetchIndexMetrics resolves", async () => {
-  render(<IndexAdvisorTab />);
+  render(<IndexAdvisorTab queryEditorContent="SELECT * FROM c" databaseId="db1" containerId="col1" />);
   expect(screen.getByRole("progressbar")).toBeInTheDocument();
   await waitFor(() => expect(screen.queryByRole("progressbar")).not.toBeInTheDocument());
 });
 
 test("calls fetchAll with correct query and options", async () => {
-  render(<IndexAdvisorTab />);
+  render(<IndexAdvisorTab queryEditorContent="SELECT * FROM c" databaseId="db1" containerId="col1" />);
   await waitFor(() => expect(mockFetchAll).toHaveBeenCalled());
 });
 test("renders IndexAdvisorTab when clicked from ResultsView", async () => {
-  render(<IndexAdvisorTab />);
+  render(<IndexAdvisorTab queryEditorContent="SELECT * FROM c" databaseId="db1" containerId="col1" />);
   await waitFor(() => expect(screen.getByText("Included in Current Policy")).toBeInTheDocument());
   expect(screen.getByText("/foo/?")).toBeInTheDocument();
 });
 test("renders index metrics from SDK response", async () => {
-  render(<IndexAdvisorTab />);
+  render(<IndexAdvisorTab queryEditorContent="SELECT * FROM c" databaseId="db1" containerId="col1" />);
   await waitFor(() => expect(screen.getByText("/foo/?")).toBeInTheDocument());
   expect(screen.getByText("/bar/?")).toBeInTheDocument();
   expect(screen.getByText("/baz/? DESC, /qux/? ASC")).toBeInTheDocument();
@@ -164,20 +181,20 @@ test("renders index metrics from SDK response", async () => {
 
 test("calls handleError if fetchIndexMetrics throws", async () => {
   mockFetchAll.mockRejectedValueOnce(new Error("fail"));
-  render(<IndexAdvisorTab />);
+  render(<IndexAdvisorTab queryEditorContent="SELECT * FROM c" databaseId="db1" containerId="col1" />);
   await waitFor(() => expect(mockHandleError).toHaveBeenCalled());
 });
 
 test("calls handleError if fetchIndexMetrics throws2nd", async () => {
   mockFetchAll.mockRejectedValueOnce(new Error("fail"));
 
-  render(<IndexAdvisorTab />);
+  render(<IndexAdvisorTab queryEditorContent="SELECT * FROM c" databaseId="db1" containerId="col1" />);
   await waitFor(() => expect(mockHandleError).toHaveBeenCalled());
   expect(screen.queryByRole("status")).not.toBeInTheDocument();
 });
 
 test("IndexingPolicyStore stores updated policy on componentDidMount", async () => {
-  render(<IndexAdvisorTab />);
+  render(<IndexAdvisorTab queryEditorContent="SELECT * FROM c" databaseId="db1" containerId="col1" />);
   await waitFor(() => expect(mockRead).toHaveBeenCalled());
 
   const readResult = await mockRead.mock.results[0].value;
@@ -190,7 +207,7 @@ test("IndexingPolicyStore stores updated policy on componentDidMount", async ()
 });
 
 test("refreshCollectionData updates observable and re-renders", async () => {
-  render(<IndexAdvisorTab />);
+  render(<IndexAdvisorTab queryEditorContent="SELECT * FROM c" databaseId="db1" containerId="col1" />);
   await waitFor(() => expect(screen.getByText("/bar/?")).toBeInTheDocument());
 
   const checkboxes = screen.getAllByRole("checkbox");

src/Explorer/Tabs/QueryTab/ResultsView.tsx

@@ -29,8 +29,12 @@ import MongoUtility from "Common/MongoUtility";
 import { QueryMetrics } from "Contracts/DataModels";
 import { QueryResults } from "Contracts/ViewModels";
 import { EditorReact } from "Explorer/Controls/Editor/EditorReact";
-import { parseIndexMetrics, renderImpactDots, type IndexMetricsResponse } from "Explorer/Tabs/QueryTab/IndexAdvisorUtils";
+import {
-import { IDocument, useQueryMetadataStore } from "Explorer/Tabs/QueryTab/QueryTabComponent";
+  parseIndexMetrics,
+  renderImpactDots,
+  type IndexMetricsResponse,
+} from "Explorer/Tabs/QueryTab/IndexAdvisorUtils";
+import { IDocument } from "Explorer/Tabs/QueryTab/QueryTabComponent";
 import { useQueryTabStyles } from "Explorer/Tabs/QueryTab/Styles";
 import React, { useCallback, useEffect, useState } from "react";
 import { userContext } from "UserContext";
@@ -38,6 +42,7 @@ import { logConsoleProgress } from "Utils/NotificationConsoleUtils";
 import create from "zustand";
 import { client } from "../../../Common/CosmosClient";
 import { handleError } from "../../../Common/ErrorHandlingUtils";
+import { sampleDataClient } from "../../../Common/SampleDataClient";
 import { ResultsViewProps } from "./QueryResultSection";
 import { useIndexAdvisorStyles } from "./StylesAdvisor";
 enum ResultsTabs {
@@ -544,9 +549,14 @@ export interface IIndexMetric {
   path?: string;
   composite?: { path: string; order: string }[];
 }
-export const IndexAdvisorTab: React.FC<{ queryResults?: QueryResults }> = ({ queryResults }) => {
+export const IndexAdvisorTab: React.FC<{
+  queryResults?: QueryResults;
+  queryEditorContent?: string;
+  databaseId?: string;
+  containerId?: string;
+}> = ({ queryResults, queryEditorContent, databaseId, containerId }) => {
   const style = useIndexAdvisorStyles();
-  const { userQuery, databaseId, containerId } = useQueryMetadataStore();
+
   const [loading, setLoading] = useState(false);
   const [indexMetrics, setIndexMetrics] = useState<IndexMetricsResponse | null>(null);
   const [showIncluded, setShowIncluded] = useState(true);
@@ -561,20 +571,21 @@ export const IndexAdvisorTab: React.FC<{ queryResults?: QueryResults }> = ({ que
   const indexingMetricsDocLink = "https://learn.microsoft.com/azure/cosmos-db/nosql/index-metrics";
 
   const fetchIndexMetrics = async () => {
-    if (!userQuery || !databaseId || !containerId) {
+    if (!queryEditorContent || !databaseId || !containerId) {
       return;
     }
 
     setLoading(true);
     const clearMessage = logConsoleProgress(`Querying items with IndexMetrics in container ${containerId}`);
     try {
-      const containerRef = client().database(databaseId).container(containerId);
-      const { resource: containerDef } = await containerRef.read();
-
       const querySpec = {
-        query: userQuery,
+        query: queryEditorContent,
       };
-      const sdkResponse = await client()
+
+      // Use sampleDataClient for CopilotSampleDB, regular client for other databases
+      const cosmosClient = databaseId === "CopilotSampleDB" ? sampleDataClient() : client();
+
+      const sdkResponse = await cosmosClient
         .database(databaseId)
         .container(containerId)
         .items.query(querySpec, {
@@ -582,9 +593,8 @@ export const IndexAdvisorTab: React.FC<{ queryResults?: QueryResults }> = ({ que
         })
         .fetchAll();
 
-      const parsedMetrics = typeof sdkResponse.indexMetrics === 'string'
+      const parsedMetrics =
-        ? JSON.parse(sdkResponse.indexMetrics)
+        typeof sdkResponse.indexMetrics === "string" ? JSON.parse(sdkResponse.indexMetrics) : sdkResponse.indexMetrics;
-        : sdkResponse.indexMetrics;
 
       setIndexMetrics(parsedMetrics);
     } catch (error) {
@@ -597,7 +607,7 @@ export const IndexAdvisorTab: React.FC<{ queryResults?: QueryResults }> = ({ que
 
   // Fetch index metrics when query results change (i.e., when Execute Query is clicked)
   useEffect(() => {
-    if (userQuery && databaseId && containerId && queryResults) {
+    if (queryEditorContent && databaseId && containerId && queryResults) {
       fetchIndexMetrics();
     }
   }, [queryResults]);
@@ -843,7 +853,14 @@ export const IndexAdvisorTab: React.FC<{ queryResults?: QueryResults }> = ({ que
     </div>
   );
 };
-export const ResultsView: React.FC<ResultsViewProps> = ({ isMongoDB, queryResults, executeQueryDocumentsPage }) => {
+export const ResultsView: React.FC<ResultsViewProps> = ({
+  isMongoDB,
+  queryResults,
+  executeQueryDocumentsPage,
+  queryEditorContent,
+  databaseId,
+  containerId,
+}) => {
   const styles = useQueryTabStyles();
   const [activeTab, setActiveTab] = useState<ResultsTabs>(ResultsTabs.Results);
 
@@ -884,7 +901,14 @@ export const ResultsView: React.FC<ResultsViewProps> = ({ isMongoDB, queryResult
           />
         )}
         {activeTab === ResultsTabs.QueryStats && <QueryStatsTab queryResults={queryResults} />}
-        {activeTab === ResultsTabs.IndexAdvisor && <IndexAdvisorTab queryResults={queryResults} />}
+        {activeTab === ResultsTabs.IndexAdvisor && (
+          <IndexAdvisorTab
+            queryResults={queryResults}
+            queryEditorContent={queryEditorContent}
+            databaseId={databaseId}
+            containerId={containerId}
+          />
+        )}
       </div>
     </div>
   );