Fix "items" button in sample data resource tree does not open documents tab (#1528)

src/Explorer/QueryCopilot/QueryCopilotTab.tsx

@@ -1,5 +1,5 @@
 /* eslint-disable no-console */
-import { FeedOptions, ItemDefinition, QueryIterator, Resource } from "@azure/cosmos";
+import { FeedOptions } from "@azure/cosmos";
 import {
   Callout,
   CommandBarButton,
@@ -17,16 +17,10 @@ import {
   TextField,
 } from "@fluentui/react";
 import { useBoolean } from "@fluentui/react-hooks";
-import {
+import { QueryCopilotSampleContainerId, QueryCopilotSampleContainerSchema } from "Common/Constants";
-  QueryCopilotSampleContainerId,
-  QueryCopilotSampleContainerSchema,
-  QueryCopilotSampleDatabaseId,
-} from "Common/Constants";
 import { getErrorMessage, handleError } from "Common/ErrorHandlingUtils";
 import { shouldEnableCrossPartitionKey } from "Common/HeadersUtility";
 import { MinimalQueryIterator } from "Common/IteratorUtilities";
-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";
@@ -37,7 +31,7 @@ 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 { submitFeedback } from "Explorer/QueryCopilot/QueryCopilotUtilities";
+import { querySampleDocuments, submitFeedback } from "Explorer/QueryCopilot/QueryCopilotUtilities";
 import { SamplePrompts, SamplePromptsProps } from "Explorer/QueryCopilot/SamplePrompts/SamplePrompts";
 import { QueryResultSection } from "Explorer/Tabs/QueryTab/QueryResultSection";
 import { userContext } from "UserContext";
@@ -200,14 +194,6 @@ 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 = querySampleDocuments(queryToExecute, {

src/Explorer/QueryCopilot/QueryCopilotUtilities.ts

@@ -1,5 +1,15 @@
-import { QueryCopilotSampleContainerSchema } from "Common/Constants";
+import { FeedOptions, Item, ItemDefinition, QueryIterator, Resource } from "@azure/cosmos";
+import {
+  QueryCopilotSampleContainerId,
+  QueryCopilotSampleContainerSchema,
+  QueryCopilotSampleDatabaseId,
+} from "Common/Constants";
 import { handleError } from "Common/ErrorHandlingUtils";
+import { sampleDataClient } from "Common/SampleDataClient";
+import { getPartitionKeyValue } from "Common/dataAccess/getPartitionKeyValue";
+import { getCommonQueryOptions } from "Common/dataAccess/queryDocuments";
+import DocumentId from "Explorer/Tree/DocumentId";
+import { logConsoleProgress } from "Utils/NotificationConsoleUtils";
 
 interface FeedbackParams {
   likeQuery: boolean;
@@ -21,16 +31,41 @@ export const submitFeedback = async (params: FeedbackParams): Promise<void> => {
       contact: contact || "",
     };
 
-    const response = await fetch("https://copilotorchestrater.azurewebsites.net/feedback", {
+    await fetch("https://copilotorchestrater.azurewebsites.net/feedback", {
       method: "POST",
       headers: {
         "content-type": "application/json",
       },
       body: JSON.stringify(payload),
     });
-    // eslint-disable-next-line no-console
-    console.log(response);
   } catch (error) {
     handleError(error, "copilotSubmitFeedback");
   }
 };
+
+export const querySampleDocuments = (query: string, options: FeedOptions): QueryIterator<ItemDefinition & Resource> => {
+  options = getCommonQueryOptions(options);
+  return sampleDataClient()
+    .database(QueryCopilotSampleDatabaseId)
+    .container(QueryCopilotSampleContainerId)
+    .items.query(query, options);
+};
+
+export const readSampleDocument = async (documentId: DocumentId): Promise<Item> => {
+  const clearMessage = logConsoleProgress(`Reading item ${documentId.id()}`);
+
+  try {
+    const response = await sampleDataClient()
+      .database(QueryCopilotSampleDatabaseId)
+      .container(QueryCopilotSampleContainerId)
+      .item(documentId.id(), getPartitionKeyValue(documentId))
+      .read();
+
+    return response?.resource;
+  } catch (error) {
+    handleError(error, "ReadDocument", `Failed to read item ${documentId.id()}`);
+    throw error;
+  } finally {
+    clearMessage();
+  }
+};

src/Explorer/Tabs/DocumentsTab.ts

@@ -1,4 +1,5 @@
 import { extractPartitionKey, ItemDefinition, PartitionKeyDefinition, QueryIterator, Resource } from "@azure/cosmos";
+import { querySampleDocuments, readSampleDocument } from "Explorer/QueryCopilot/QueryCopilotUtilities";
 import * as ko from "knockout";
 import Q from "q";
 import DeleteDocumentIcon from "../../../images/DeleteDocument.svg";
@@ -7,7 +8,12 @@ import NewDocumentIcon from "../../../images/NewDocument.svg";
 import SaveIcon from "../../../images/save-cosmos.svg";
 import UploadIcon from "../../../images/Upload_16x16.svg";
 import * as Constants from "../../Common/Constants";
-import { DocumentsGridMetrics, KeyCodes } from "../../Common/Constants";
+import {
+  DocumentsGridMetrics,
+  KeyCodes,
+  QueryCopilotSampleContainerId,
+  QueryCopilotSampleDatabaseId,
+} from "../../Common/Constants";
 import { createDocument } from "../../Common/dataAccess/createDocument";
 import { deleteDocument } from "../../Common/dataAccess/deleteDocument";
 import { queryDocuments } from "../../Common/dataAccess/queryDocuments";
@@ -71,6 +77,7 @@ export default class DocumentsTab extends TabsBase {
 
   private _documentsIterator: QueryIterator<ItemDefinition & Resource>;
   private _resourceTokenPartitionKey: string;
+  private _isQueryCopilotSampleContainer: boolean;
 
   constructor(options: ViewModels.DocumentsTabOptions) {
     super(options);
@@ -317,6 +324,9 @@ export default class DocumentsTab extends TabsBase {
     this.selectedDocumentContent.subscribe((newContent: string) => this._onEditorContentChange(newContent));
 
     this.showPartitionKey = this._shouldShowPartitionKey();
+    this._isQueryCopilotSampleContainer =
+      this.collection?.databaseId === QueryCopilotSampleDatabaseId &&
+      this.collection?.id() === QueryCopilotSampleContainerId;
   }
 
   private _shouldShowPartitionKey(): boolean {
@@ -678,7 +688,6 @@ export default class DocumentsTab extends TabsBase {
   }
 
   public createIterator(): QueryIterator<ItemDefinition & Resource> {
-    let filters = this.lastFilterContents();
     const filter: string = this.filterContent().trim();
     const query: string = this.buildQuery(filter);
     let options: any = {};
@@ -688,12 +697,16 @@ export default class DocumentsTab extends TabsBase {
       options.partitionKey = this._resourceTokenPartitionKey;
     }
 
-    return queryDocuments(this.collection.databaseId, this.collection.id(), query, options);
+    return this._isQueryCopilotSampleContainer
+      ? querySampleDocuments(query, options)
+      : queryDocuments(this.collection.databaseId, this.collection.id(), query, options);
   }
 
   public async selectDocument(documentId: DocumentId): Promise<void> {
     this.selectedDocumentId(documentId);
-    const content = await readDocument(this.collection, documentId);
+    const content = await (this._isQueryCopilotSampleContainer
+      ? readSampleDocument(documentId)
+      : readDocument(this.collection, documentId));
     this.initDocumentEditor(documentId, content);
   }
 

src/Explorer/Tree/ResourceTokenCollection.ts

@@ -2,10 +2,10 @@ import * as ko from "knockout";
 import * as Constants from "../../Common/Constants";
 import * as DataModels from "../../Contracts/DataModels";
 import * as ViewModels from "../../Contracts/ViewModels";
-import { useTabs } from "../../hooks/useTabs";
 import { Action, ActionModifiers } from "../../Shared/Telemetry/TelemetryConstants";
 import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";
 import { userContext } from "../../UserContext";
+import { useTabs } from "../../hooks/useTabs";
 import Explorer from "../Explorer";
 import DocumentsTab from "../Tabs/DocumentsTab";
 import { NewQueryTab } from "../Tabs/QueryTab/QueryTab";
@@ -139,7 +139,7 @@ export default class ResourceTokenCollection implements ViewModels.CollectionBas
 
       documentsTab = new DocumentsTab({
         partitionKey: this.partitionKey,
-        resourceTokenPartitionKey: userContext.parsedResourceToken.partitionKey,
+        resourceTokenPartitionKey: userContext.parsedResourceToken?.partitionKey,
         documentIds: ko.observableArray<DocumentId>([]),
         tabKind: ViewModels.CollectionTabKind.Documents,
         title: "Items",

src/Explorer/Tree/SampleDataTree.tsx

@@ -34,11 +34,13 @@ export const SampleDataTree = ({
               // Rewritten version of expandCollapseCollection
               useSelectedNode.getState().setSelectedNode(sampleDataResourceTokenCollection);
               useCommandBar.getState().setContextButtons([]);
-              useTabs().refreshActiveTab(
+              useTabs
-                (tab: TabsBase) =>
+                .getState()
-                  tab.collection?.id() === sampleDataResourceTokenCollection.id() &&
+                .refreshActiveTab(
-                  tab.collection.databaseId === sampleDataResourceTokenCollection.databaseId
+                  (tab: TabsBase) =>
-              );
+                    tab.collection?.id() === sampleDataResourceTokenCollection.id() &&
+                    tab.collection.databaseId === sampleDataResourceTokenCollection.databaseId
+                );
             },
             isSelected: () =>
               useSelectedNode
@@ -51,6 +53,15 @@ export const SampleDataTree = ({
             children: [
               {
                 label: "Items",
+                onClick: () => sampleDataResourceTokenCollection.onDocumentDBDocumentsClick(),
+                isSelected: () =>
+                  useSelectedNode
+                    .getState()
+                    .isDataNodeSelected(
+                      sampleDataResourceTokenCollection.databaseId,
+                      sampleDataResourceTokenCollection.id(),
+                      [ViewModels.CollectionTabKind.Documents]
+                    ),
               },
             ],
           },