[Query Copilot] Resource tree styling and New query button (#1519)

* Styling implemented related with the work item

* Sample container New query button implementation

* Fixing related with the not rendering Sample Data

* Fix race condition when rendering sample data resource tree

* Remove export keyword for updateContextForSampleData

* Copilot New Query should open Copilot tab

* showing buttons in sample command bar

---------

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

src/Explorer/ContextMenuButtonFactory.tsx

@@ -1,3 +1,4 @@
+import { ReactTabKind, useTabs } from "hooks/useTabs";
 import React from "react";
 import AddCollectionIcon from "../../images/AddCollection.svg";
 import AddSqlQueryIcon from "../../images/AddSqlQuery_16x16.svg";
@@ -140,6 +141,21 @@ export const createCollectionContextMenuButton = (
   return items;
 };
 
+export const createSampleCollectionContextMenuButton = (
+  selectedCollection: ViewModels.CollectionBase
+): TreeNodeMenuItem[] => {
+  const items: TreeNodeMenuItem[] = [];
+  if (userContext.apiType === "SQL") {
+    items.push({
+      iconSrc: AddSqlQueryIcon,
+      onClick: () => selectedCollection && useTabs.getState().openAndActivateReactTab(ReactTabKind.QueryCopilot),
+      label: "New SQL Query",
+    });
+  }
+
+  return items;
+};
+
 export const createStoreProcedureContextMenuItems = (
   container: Explorer,
   storedProcedure: StoredProcedure

src/Explorer/Controls/Accordion/AccordionComponent.tsx

@@ -25,6 +25,7 @@ export class AccordionComponent extends React.Component<AccordionComponentProps>
 export interface AccordionItemComponentProps {
   title: string;
   isExpanded?: boolean;
+  styles?: React.CSSProperties;
 }
 
 interface AccordionItemComponentState {
@@ -53,6 +54,7 @@ export class AccordionItemComponent extends React.Component<AccordionItemCompone
   }
 
   public render(): JSX.Element {
+    const { styles } = this.props;
     return (
       <div className="accordionItemContainer">
         <div className="accordionItemHeader" onClick={this.onHeaderClick} onKeyPress={this.onHeaderKeyPress}>
@@ -60,7 +62,11 @@ export class AccordionItemComponent extends React.Component<AccordionItemCompone
           {this.props.title}
         </div>
         <div className="accordionItemContent">
-          <AnimateHeight duration={AccordionItemComponent.durationMS} height={this.state.isExpanded ? "auto" : 0}>
+          <AnimateHeight
+            style={{ ...styles }}
+            duration={AccordionItemComponent.durationMS}
+            height={this.state.isExpanded ? "auto" : 0}
+          >
             {this.props.children}
           </AnimateHeight>
         </div>

src/Explorer/Tree/ResourceTree.tsx

@@ -790,7 +790,11 @@ export const ResourceTree: React.FC<ResourceTreeProps> = ({ container }: Resourc
       {!isNotebookEnabled && isSampleDataEnabled && (
         <>
           <AccordionComponent>
-            <AccordionItemComponent title={"MY DATA"} isExpanded={!gitHubNotebooksContentRoot}>
+            <AccordionItemComponent
+              title={"MY DATA"}
+              isExpanded={!gitHubNotebooksContentRoot}
+              styles={{ maxHeight: 230 }}
+            >
               <TreeComponent className="dataResourceTree" rootNode={dataRootNode} />
             </AccordionItemComponent>
             <AccordionItemComponent title={"SAMPLE DATA"}>
@@ -804,7 +808,11 @@ export const ResourceTree: React.FC<ResourceTreeProps> = ({ container }: Resourc
       {isNotebookEnabled && isSampleDataEnabled && (
         <>
           <AccordionComponent>
-            <AccordionItemComponent title={"MY DATA"} isExpanded={!gitHubNotebooksContentRoot}>
+            <AccordionItemComponent
+              title={"MY DATA"}
+              isExpanded={!gitHubNotebooksContentRoot}
+              styles={{ maxHeight: 130 }}
+            >
               <TreeComponent className="dataResourceTree" rootNode={dataRootNode} />
             </AccordionItemComponent>
             <AccordionItemComponent title={"SAMPLE DATA"}>

src/Explorer/Tree/SampleDataTree.tsx

@@ -1,7 +1,12 @@
+import { useCommandBar } from "Explorer/Menus/CommandBar/CommandBarComponentAdapter";
+import TabsBase from "Explorer/Tabs/TabsBase";
+import { useSelectedNode } from "Explorer/useSelectedNode";
+import { useTabs } from "hooks/useTabs";
 import React, { useEffect, useState } from "react";
 import CosmosDBIcon from "../../../images/Azure-Cosmos-DB.svg";
 import CollectionIcon from "../../../images/tree-collection.svg";
 import * as ViewModels from "../../Contracts/ViewModels";
+import * as ResourceTreeContextMenuButtonFactory from "../ContextMenuButtonFactory";
 import { TreeComponent, TreeNode } from "../Controls/TreeComponent/TreeComponent";
 
 export const SampleDataTree = ({
@@ -15,14 +20,36 @@ export const SampleDataTree = ({
     if (sampleDataResourceTokenCollection) {
       const updatedSampleTree: TreeNode = {
         label: sampleDataResourceTokenCollection.databaseId,
-        isExpanded: true,
+        isExpanded: false,
         iconSrc: CosmosDBIcon,
+        className: "databaseHeader",
         children: [
           {
             label: sampleDataResourceTokenCollection.id(),
             iconSrc: CollectionIcon,
-            isExpanded: true,
+            isExpanded: false,
-            className: "collectionHeader",
+            className: "dataResourceTree",
+            contextMenu: ResourceTreeContextMenuButtonFactory.createSampleCollectionContextMenuButton(
+              sampleDataResourceTokenCollection
+            ),
+            onClick: () => {
+              // Rewritten version of expandCollapseCollection
+              useSelectedNode.getState().setSelectedNode(sampleDataResourceTokenCollection);
+              useCommandBar.getState().setContextButtons([]);
+              useTabs().refreshActiveTab(
+                (tab: TabsBase) =>
+                  tab.collection?.id() === sampleDataResourceTokenCollection.id() &&
+                  tab.collection.databaseId === sampleDataResourceTokenCollection.databaseId
+              );
+            },
+            isSelected: () =>
+              useSelectedNode
+                .getState()
+                .isDataNodeSelected(
+                  sampleDataResourceTokenCollection.databaseId,
+                  sampleDataResourceTokenCollection.id()
+                ),
+            onContextMenuOpen: () => useSelectedNode.getState().setSelectedNode(sampleDataResourceTokenCollection),
             children: [
               {
                 label: "Items",
@@ -35,7 +62,5 @@ export const SampleDataTree = ({
     }
   }, [sampleDataResourceTokenCollection]);
 
-  return (
+  return <TreeComponent className="dataResourceTree" rootNode={root || { label: "Sample data not initialized." }} />;
-    <TreeComponent className="sampleDataResourceTree" rootNode={root || { label: "Sample data not initialized." }} />
-  );
 };

src/hooks/useKnockoutExplorer.ts

@@ -62,10 +62,6 @@ export function useKnockoutExplorer(platform: Platform): Explorer {
           setExplorer(explorer);
         }
       }
-
-      if (userContext.features.enableCopilot) {
-        await updateContextForSampleData();
-      }
     };
     effect();
   }, [platform]);
@@ -73,6 +69,9 @@ export function useKnockoutExplorer(platform: Platform): Explorer {
   useEffect(() => {
     if (explorer) {
       applyExplorerBindings(explorer);
+      if (userContext.features.enableCopilot) {
+        updateContextForSampleData(explorer);
+      }
     }
   }, [explorer]);
 
@@ -415,7 +414,7 @@ interface PortalMessage {
   inputs?: DataExplorerInputsFrame;
 }
 
-async function updateContextForSampleData(): Promise<void> {
+async function updateContextForSampleData(explorer: Explorer): Promise<void> {
   if (!userContext.features.enableCopilot) {
     return;
   }
@@ -435,6 +434,8 @@ async function updateContextForSampleData(): Promise<void> {
   const data: SampledataconnectionResponse = await response.json();
   const sampleDataConnectionInfo = parseResourceTokenConnectionString(data.connectionString);
   updateUserContext({ sampleDataConnectionInfo });
+
+  await explorer.refreshSampleData();
 }
 
 interface SampledataconnectionResponse {