From 471a4ea4ad1735a52559742aa33981f50bb07304 Mon Sep 17 00:00:00 2001
From: Sean Li <lnse@amazon.com>
Date: Fri, 7 Jun 2024 15:26:45 -0700
Subject: [PATCH] updating datasource in request

Signed-off-by: Sean Li <lnse@amazon.com>
---
 .../public/search/sql_async_search_interceptor.ts    | 12 +++++++-----
 .../query_enhancements/server/routes/index.ts        |  1 +
 .../server/search/async/sql_async_facet.ts           |  2 +-
 .../search/search_source/fetch/get_search_params.ts  |  4 ++--
 4 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/plugins-extra/query_enhancements/public/search/sql_async_search_interceptor.ts b/plugins-extra/query_enhancements/public/search/sql_async_search_interceptor.ts
index 54075dc607f1..6f73983a04ed 100644
--- a/plugins-extra/query_enhancements/public/search/sql_async_search_interceptor.ts
+++ b/plugins-extra/query_enhancements/public/search/sql_async_search_interceptor.ts
@@ -37,9 +37,10 @@ export class SQLAsyncQlSearchInterceptor extends SearchInterceptor {
     const { id, ...searchRequest } = request;
     const path = trimEnd('/api/sqlasyncql/jobs');
 
-    const fetchDataFrame = (queryString: string | undefined, df = null) => {
+    const fetchDataFrame = (queryString: string | undefined, dataSource: string, df = null) => {
       const body = stringify({
         query: { qs: queryString, format: 'jdbc' },
+        dataSource,
         df,
       });
       return from(
@@ -53,14 +54,15 @@ export class SQLAsyncQlSearchInterceptor extends SearchInterceptor {
     };
 
     const rawDataFrame = getRawDataFrame(searchRequest);
-    const dataFrame = fetchDataFrame(getRawQueryString(searchRequest), rawDataFrame);
+    const dataFrame = fetchDataFrame(
+      getRawQueryString(searchRequest),
+      searchRequest.params.index,
+      rawDataFrame
+    );
 
     // subscribe to dataFrame to see if an error is returned, display a toast message if so
     dataFrame.subscribe((df) => {
       // TODO: MQL Async: clean later
-      // if (!df.body.error) {
-      //   console.log('SEARCH INTERCEPTOR:', df);
-      // }
       if (!df.body.error) return;
       const jsError = new Error(df.body.error.response);
       this.deps.toasts.addError(jsError, {
diff --git a/plugins-extra/query_enhancements/server/routes/index.ts b/plugins-extra/query_enhancements/server/routes/index.ts
index c34246b25119..e4cb589f3b8b 100644
--- a/plugins-extra/query_enhancements/server/routes/index.ts
+++ b/plugins-extra/query_enhancements/server/routes/index.ts
@@ -106,6 +106,7 @@ export function defineRoutes(
             format: schema.string(),
           }),
           df: schema.any(),
+          dataSource: schema.string(),
         }),
       },
     },
diff --git a/plugins-extra/query_enhancements/server/search/async/sql_async_facet.ts b/plugins-extra/query_enhancements/server/search/async/sql_async_facet.ts
index 45b7e44e606b..1afcdc56351a 100644
--- a/plugins-extra/query_enhancements/server/search/async/sql_async_facet.ts
+++ b/plugins-extra/query_enhancements/server/search/async/sql_async_facet.ts
@@ -18,7 +18,7 @@ export class SQLAsyncFacet {
       const params = {
         body: {
           query: request.body.query.qs,
-          datasource: df?.name,
+          datasource: request.body.dataSource,
           lang: 'sql',
           sessionId: df?.meta?.sessionId,
         },
diff --git a/src/plugins/data/common/search/search_source/fetch/get_search_params.ts b/src/plugins/data/common/search/search_source/fetch/get_search_params.ts
index c0ad36ffcaa2..d7cb331fbb7b 100644
--- a/src/plugins/data/common/search/search_source/fetch/get_search_params.ts
+++ b/src/plugins/data/common/search/search_source/fetch/get_search_params.ts
@@ -77,9 +77,9 @@ export async function getExternalSearchParamsFromRequest(
   const dataFrame =
     getDataFrame() ??
     (dataSource?.name ? getDataFrameBySource(dataSource.name) : null) ??
-    (await setDataFrame(createDataFrame({ name: dataSource?.name ?? '', fields: [] })));
+    (await setDataFrame(createDataFrame({ name: dataSource?.name ?? indexTitle, fields: [] })));
   return {
-    index: indexTitle,
+    index: dataSource ? dataSource.name : indexTitle,
     body: {
       ...searchRequest.body,
       df: dataFrame,