Fix query assist for query editor (#7552)

After the following PRs:  #7492, #7546, #7540

this commit added skip(1) back to dataset manager observable: fef6156, we need to revert changes done in
fix(query assist): update reading data source id from dataset manager #7464 (comment)

    revert dataset manager observable usage in query assist to support skip(1)
    revert dataset manager tests

[Discover Next] Fixes Discover styles #7546 removed query editor header div, this PR adds it back to enable query editor extensions

Signed-off-by: Joshua Li <joshuali925@gmail.com>

src/plugins/data/public/query/dataset_manager/dataset_manager.mock.ts

@@ -12,6 +12,8 @@ const createSetupContractMock = () => {
     setDataSet: jest.fn(),
     getUpdates$: jest.fn(),
     getDefaultDataSet: jest.fn(),
+    fetchDefaultDataSet: jest.fn(),
+    initWithIndexPattern: jest.fn(),
   };
   return dataSetManagerMock;
 };

src/plugins/data/public/query/dataset_manager/dataset_manager.test.ts

@@ -15,11 +15,11 @@ describe('DataSetManager', () => {
     service = new DataSetManager(uiSettingsMock);
   });
 
-  test('getUpdates$ emits initially and after data set changes', () => {
+  test('getUpdates$ is a cold emits only after dataset changes', () => {
     const obs$ = service.getUpdates$();
-    const emittedValues: Array<SimpleDataSet | undefined> = [];
+    const emittedValues: SimpleDataSet[] = [];
     obs$.subscribe((v) => {
-      emittedValues.push(v);
+      emittedValues.push(v!);
     });
     expect(emittedValues).toHaveLength(0);
     expect(emittedValues[0]).toEqual(undefined);

src/plugins/data/public/query/mocks.ts

@@ -42,7 +42,7 @@ const createSetupContractMock = () => {
     filterManager: createFilterManagerMock(),
     timefilter: timefilterServiceMock.createSetupContract(),
     queryString: queryStringManagerMock.createSetupContract(),
-    dataSet: dataSetManagerMock.createSetupContract(),
+    dataSetManager: dataSetManagerMock.createSetupContract(),
     state$: new Observable(),
   };
 
@@ -57,7 +57,7 @@ const createStartContractMock = () => {
     savedQueries: jest.fn() as any,
     state$: new Observable(),
     timefilter: timefilterServiceMock.createStartContract(),
-    dataSet: dataSetManagerMock.createStartContract(),
+    dataSetManager: dataSetManagerMock.createStartContract(),
     getOpenSearchQuery: jest.fn(),
   };
 

src/plugins/data/public/ui/dataset_navigator/create_dataset_navigator.tsx

@@ -15,7 +15,7 @@ export function createDataSetNavigator(
   dataSetManager: DataSetContract
 ) {
   // Return a function that takes props, omitting the dependencies from the props type
-  return (props: Omit<DataSetNavigatorProps, 'savedObjectsClient' | 'http' | 'dataSet'>) => (
+  return (props: Omit<DataSetNavigatorProps, 'savedObjectsClient' | 'http' | 'dataSetManager'>) => (
     <DataSetNavigator
       {...props}
       savedObjectsClient={savedObjectsClient}

src/plugins/data/public/ui/query_editor/query_editor.tsx

@@ -406,6 +406,10 @@ export default class QueryEditorUI extends Component<Props, State> {
           {languageSelector}
           {this.props.queryActions}
         </div>
+        <div
+          ref={this.headerRef}
+          className={classNames('osdQueryEditor__header', this.props.headerClassName)}
+        />
         {!this.state.isCollapsed && (
           <div className="osdQueryEditor__body">{languageEditor.Body()}</div>
         )}
@@ -505,7 +509,6 @@ export default class QueryEditorUI extends Component<Props, State> {
            </EuiFlexItem>
 
            <EuiFlexItem onClick={this.onClickInput} grow={true}>
-             <div ref={this.headerRef} className={headerClassName} />
              {!this.state.isCollapsed && useQueryEditor && (
                <CodeEditor
                  height={70}

src/plugins/query_enhancements/public/index.scss

@@ -1 +1,6 @@
-/* stylelint-disable no-empty-source */
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+@import "./query_assist";

src/plugins/query_enhancements/public/query_assist/_index.scss

@@ -0,0 +1,14 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+.queryAssist {
+  &.queryAssist__form {
+    padding: 0 $euiSizeXS $euiSizeXS;
+  }
+
+  &.queryAssist__callout {
+    margin-top: $euiSizeXS;
+  }
+}

src/plugins/query_enhancements/public/query_assist/components/call_outs.tsx

@@ -22,6 +22,7 @@ export type QueryAssistCallOutType =
 
 const EmptyIndexCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
   <EuiCallOut
+    className="queryAssist queryAssist__callout"
     data-test-subj="query-assist-empty-index-callout"
     title={
       <FormattedMessage
@@ -39,6 +40,7 @@ const EmptyIndexCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
 
 const ProhibitedQueryCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
   <EuiCallOut
+    className="queryAssist queryAssist__callout"
     data-test-subj="query-assist-guard-callout"
     title={
       <FormattedMessage
@@ -56,6 +58,7 @@ const ProhibitedQueryCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
 
 const EmptyQueryCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
   <EuiCallOut
+    className="queryAssist queryAssist__callout"
     data-test-subj="query-assist-empty-query-callout"
     title={
       <FormattedMessage
@@ -73,6 +76,7 @@ const EmptyQueryCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
 
 const QueryGeneratedCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
   <EuiCallOut
+    className="queryAssist queryAssist__callout"
     data-test-subj="query-assist-query-generated-callout"
     title={
       <FormattedMessage

src/plugins/query_enhancements/public/query_assist/components/query_assist_banner.tsx

@@ -40,6 +40,7 @@ export const QueryAssistBanner: React.FC<QueryAssistBannerProps> = (props) => {
 
   return (
     <EuiCallOut
+      className="queryAssist"
       size="s"
       title={
         <EuiFlexGroup responsive={false} gutterSize="s" alignItems="center">

src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.tsx

@@ -35,7 +35,9 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
   const { generateQuery, loading } = useGenerateQuery();
   const [callOutType, setCallOutType] = useState<QueryAssistCallOutType>();
   const dismissCallout = () => setCallOutType(undefined);
-  const [selectedDataSet, setSelectedDataSet] = useState<SimpleDataSet>();
+  const [selectedDataSet, setSelectedDataSet] = useState<SimpleDataSet | undefined>(
+    services.data.query.dataSetManager.getDataSet()
+  );
   const selectedIndex = selectedDataSet?.title;
   const previousQuestionRef = useRef<string>();
 
@@ -85,7 +87,7 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
   if (props.dependencies.isCollapsed) return null;
 
   return (
-    <EuiForm component="form" onSubmit={onSubmit}>
+    <EuiForm component="form" onSubmit={onSubmit} className="queryAssist queryAssist__form">
       <EuiFormRow fullWidth>
         <EuiFlexGroup gutterSize="s" responsive={false} alignItems="center">
           <EuiFlexItem>

src/plugins/query_enhancements/public/query_assist/utils/create_extension.test.tsx

@@ -15,24 +15,6 @@ import { DataSetContract } from '../../../../data/public/query';
 import { ConfigSchema } from '../../../common/config';
 import { createQueryAssistExtension } from './create_extension';
 
-const mockSimpleDataSet = {
-  id: 'mock-data-set-id',
-  title: 'mock-title',
-  dataSourceRef: {
-    id: 'mock-data-source-id',
-  },
-} as SimpleDataSet;
-
-const mockDataSetManager: jest.Mocked<DataSetContract> = {
-  getUpdates$: jest.fn().mockReturnValue(of(mockSimpleDataSet)),
-  getDataSet: jest.fn().mockReturnValue(mockSimpleDataSet),
-  setDataSet: jest.fn(),
-  getDefaultDataSet: jest.fn().mockReturnValue(mockSimpleDataSet),
-  fetchDefaultDataSet: jest.fn().mockResolvedValue(mockSimpleDataSet),
-  init: jest.fn(),
-  initWithIndexPattern: jest.fn(),
-};
-
 const coreSetupMock = coreMock.createSetup({
   pluginStartDeps: {
     data: {
@@ -41,13 +23,19 @@ const coreSetupMock = coreMock.createSetup({
   },
 });
 const httpMock = coreSetupMock.http;
-const dataMock = {
-  ...dataPluginMock.createSetupContract(),
-  query: {
-    ...dataPluginMock.createSetupContract().query,
-    dataSetManager: mockDataSetManager,
+const dataMock = dataPluginMock.createSetupContract();
+const dataSetMock = (dataMock.query.dataSetManager as unknown) as jest.Mocked<DataSetContract>;
+
+const mockSimpleDataSet = {
+  id: 'mock-data-set-id',
+  title: 'mock-title',
+  dataSourceRef: {
+    id: 'mock-data-source-id',
   },
-};
+} as SimpleDataSet;
+
+dataSetMock.getDataSet.mockReturnValue(mockSimpleDataSet);
+dataSetMock.getUpdates$.mockReturnValue(of(mockSimpleDataSet));
 
 jest.mock('../components', () => ({
   QueryAssistBar: jest.fn(() => <div>QueryAssistBar</div>),
@@ -71,17 +59,7 @@ describe('CreateExtension', () => {
 
   it('should be enabled if at least one language is configured', async () => {
     httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
-    const extension = createQueryAssistExtension(
-      httpMock,
-      {
-        ...dataMock,
-        query: {
-          ...dataMock.query,
-          dataSetManager: mockDataSetManager,
-        },
-      },
-      config
-    );
+    const extension = createQueryAssistExtension(httpMock, dataMock, config);
     const isEnabled = await firstValueFrom(extension.isEnabled$(dependencies));
     expect(isEnabled).toBeTruthy();
     expect(httpMock.get).toBeCalledWith('/api/enhancements/assist/languages', {

src/plugins/query_enhancements/public/query_assist/utils/create_extension.tsx

@@ -5,7 +5,7 @@
 
 import { HttpSetup } from 'opensearch-dashboards/public';
 import React, { useEffect, useState } from 'react';
-import { distinctUntilChanged, map, switchMap } from 'rxjs/operators';
+import { distinctUntilChanged, map, startWith, switchMap } from 'rxjs/operators';
 import { SIMPLE_DATA_SOURCE_TYPES } from '../../../../data/common';
 import {
   DataPublicPluginSetup,
@@ -26,6 +26,7 @@ const getAvailableLanguages$ = (
   data: DataPublicPluginSetup
 ) =>
   data.query.dataSetManager.getUpdates$().pipe(
+    startWith(data.query.dataSetManager.getDataSet()),
     distinctUntilChanged(),
     switchMap(async (simpleDataSet) => {
       // currently query assist tool relies on opensearch API to get index