copilot db exceptions plus zustand removal

2025-12-27 12:51:41 +00:00 · 2025-10-16 01:45:27 +05:30
parent f77666b86c
commit 129ffc57d8
4 changed files with 109 additions and 56 deletions
--- a/src/Explorer/Tabs/QueryTab/QueryResultSection.tsx
+++ b/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 />
--- a/src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx
+++ b/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(),
-          abortSignal: this.queryAbortController.signal,
-        } as unknown as FeedOptions);
+            enableCrossPartitionQuery: HeadersUtility.shouldEnableCrossPartitionKey(),
+            abortSignal: this.queryAbortController.signal,
+          } 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)
                  }
--- a/src/Explorer/Tabs/QueryTab/ResultsView.test.tsx
+++ b/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 = `
-Utilized Single Indexes
-Index Spec: /foo/?
-Index Impact Score: High
-Potential Single Indexes
-Index Spec: /bar/?
-Index Impact Score: Medium
-Utilized Composite Indexes
-Index Spec: /baz/? DESC, /qux/? ASC
-Index Impact Score: Low
-`;
+const indexMetricsResponse = {
+  UtilizedIndexes: {
+    SingleIndexes: [{ IndexSpec: "/foo/?", IndexImpactScore: "High" }],
+    CompositeIndexes: [{ IndexSpecs: ["/baz/? DESC", "/qux/? ASC"], IndexImpactScore: "Low" }],
+  },
+  PotentialIndexes: {
+    SingleIndexes: [{ IndexSpec: "/bar/?", IndexImpactScore: "Medium" }],
+    CompositeIndexes: [] as Array<{ IndexSpecs: string[]; IndexImpactScore?: string }>,
+  },
+};
+
+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");
--- a/src/Explorer/Tabs/QueryTab/ResultsView.tsx
+++ b/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 { IDocument, useQueryMetadataStore } from "Explorer/Tabs/QueryTab/QueryTabComponent";
+import {
+  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'
-        ? JSON.parse(sdkResponse.indexMetrics)
-        : sdkResponse.indexMetrics;
+      const parsedMetrics =
+        typeof sdkResponse.indexMetrics === "string" ? JSON.parse(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>
  );