From 19a00be4beea84dd8db6574a225b3e9328b60234 Mon Sep 17 00:00:00 2001
From: Kawika Avilla <kavilla414@gmail.com>
Date: Sun, 25 Aug 2024 13:00:18 +0000
Subject: [PATCH] ppl working better

Signed-off-by: Kawika Avilla <kavilla414@gmail.com>
---
 .../query/query_string/query_string_manager.ts      | 13 +++++++++++++
 .../data/public/ui/dataset_selector/index.tsx       | 13 +++++++++----
 .../data/public/ui/query_editor/query_editor.tsx    |  2 +-
 .../public/ui/query_editor/query_editor_top_row.tsx |  6 ------
 4 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/src/plugins/data/public/query/query_string/query_string_manager.ts b/src/plugins/data/public/query/query_string/query_string_manager.ts
index 310c6981f00..c9b11aa5158 100644
--- a/src/plugins/data/public/query/query_string/query_string_manager.ts
+++ b/src/plugins/data/public/query/query_string/query_string_manager.ts
@@ -157,6 +157,19 @@ export class QueryStringManager {
     };
   };
 
+  public getInitialQueryByDataset = (newDataset: Dataset) => {
+    const curQuery = this.query$.getValue();
+    const languageId = curQuery.language;
+    const language = this.languageService.getLanguage(languageId);
+    const newQuery = { ...curQuery, dataset: newDataset };
+    const input = language?.getQueryString(newQuery) || '';
+
+    return {
+      ...newQuery,
+      query: input,
+    };
+  };
+
   private getDefaultLanguage() {
     return (
       this.storage.get('userQueryLanguage') ||
diff --git a/src/plugins/data/public/ui/dataset_selector/index.tsx b/src/plugins/data/public/ui/dataset_selector/index.tsx
index d9a5fa4a18d..720cf6cd642 100644
--- a/src/plugins/data/public/ui/dataset_selector/index.tsx
+++ b/src/plugins/data/public/ui/dataset_selector/index.tsx
@@ -5,12 +5,16 @@
 
 import { useEffect, useState } from 'react';
 import React from 'react';
-import { Dataset } from '../../../common/datasets';
+import { Dataset, Query, TimeRange } from '../../../common';
 import { DatasetSelector } from './dataset_selector';
 import { useOpenSearchDashboards } from '../../../../opensearch_dashboards_react/public';
 import { IDataPluginServices } from '../../types';
 
-const ConnectedDatasetSelector = () => {
+interface ConnectedDatasetSelectorProps {
+  onSubmit: ((query: Query, dateRange?: TimeRange | undefined) => void) | undefined;
+}
+
+const ConnectedDatasetSelector = ({ onSubmit }: ConnectedDatasetSelectorProps) => {
   const [selectedDataset, setSelectedDataset] = useState<Dataset | undefined>();
   const { services } = useOpenSearchDashboards<IDataPluginServices>();
   const queryString = services.data.query.queryString;
@@ -29,8 +33,9 @@ const ConnectedDatasetSelector = () => {
   const handleDatasetChange = (dataset?: Dataset) => {
     setSelectedDataset(dataset);
     if (dataset) {
-      const query = queryString.getQuery();
-      queryString.setQuery({ ...query, dataset });
+      const query = queryString.getInitialQueryByDataset(dataset);
+      queryString.setQuery(query);
+      onSubmit!(queryString.getQuery());
     }
   };
 
diff --git a/src/plugins/data/public/ui/query_editor/query_editor.tsx b/src/plugins/data/public/ui/query_editor/query_editor.tsx
index 82befefac00..568176c5c45 100644
--- a/src/plugins/data/public/ui/query_editor/query_editor.tsx
+++ b/src/plugins/data/public/ui/query_editor/query_editor.tsx
@@ -395,7 +395,7 @@ export default class QueryEditorUI extends Component<Props, State> {
             onClick={() => this.setState({ isCollapsed: !this.state.isCollapsed })}
             isCollapsed={!this.state.isCollapsed}
           />
-          <DatasetSelector />
+          <DatasetSelector onSubmit={this.props.onSubmit} />
           <div className="osdQueryEditor__input">
             {this.state.isCollapsed
               ? languageEditor.TopBar.Collapsed()
diff --git a/src/plugins/data/public/ui/query_editor/query_editor_top_row.tsx b/src/plugins/data/public/ui/query_editor/query_editor_top_row.tsx
index bbe99b5858d..da3ba70b383 100644
--- a/src/plugins/data/public/ui/query_editor/query_editor_top_row.tsx
+++ b/src/plugins/data/public/ui/query_editor/query_editor_top_row.tsx
@@ -73,12 +73,6 @@ export default function QueryEditorTopRow(props: QueryEditorTopRowProps) {
   } = opensearchDashboards.services;
 
   const queryLanguage = props.query && props.query.language;
-  // const parsedQuery =
-  //   !language || isValidQuery(props.query) ? props.query! : queryString.getInitialQuery();
-  // if (!isEqual(parsedQuery?.query, props.query?.query)) {
-  //   onQueryChange(parsedQuery);
-  //   onSubmit({ query: parsedQuery, dateRange: getDateRange() });
-  // }
   const persistedLog: PersistedLog | undefined = React.useMemo(
     () =>
       queryLanguage && uiSettings && storage && appName