[Explorer] Supports session for s3 direct query (#1178)

* remove unused files Signed-off-by: Eric Wei <menwe@amazon.com> * missing snapshots Signed-off-by: Eric Wei <menwe@amazon.com> * remove unused files Signed-off-by: Eric Wei <menwe@amazon.com> * supports session id for direct query Signed-off-by: Eric <menwe@amazon.com> * update snapshots Signed-off-by: Eric <menwe@amazon.com> * remove logging Signed-off-by: Eric <menwe@amazon.com> --------- Signed-off-by: Eric Wei <menwe@amazon.com> Signed-off-by: Eric <menwe@amazon.com>
opensearch-project · Oct 25, 2023 · 365a2fd · 365a2fd
1 parent a6c0333
commit 365a2fd
Show file tree

Hide file tree

Showing 9 changed files with 104 additions and 21 deletions.
diff --git a/common/constants/data_sources.ts b/common/constants/data_sources.ts
@@ -14,3 +14,8 @@ export const enum QUERY_LANGUAGE {
   SQL = 'SQL',
   DQL = 'DQL',
 }
+export enum DATA_SOURCE_TYPES {
+  DEFAULT_CLUSTER_TYPE = DEFAULT_DATA_SOURCE_TYPE,
+  SPARK = 'spark',
+  S3Glue = 's3glue',
+}
diff --git a/common/constants/shared.ts b/common/constants/shared.ts
@@ -246,3 +246,7 @@ export const VISUALIZATION_ERROR = {
 };
 
 export const S3_DATASOURCE_TYPE = 'S3_DATASOURCE';
+
+export const ASYNC_QUERY_SESSION_ID = 'async-query-session-id';
+
+export const DIRECT_DUMMY_QUERY = 'select 1';
diff --git a/common/types/explorer.ts b/common/types/explorer.ts
@@ -427,3 +427,10 @@ export enum DirectQueryLoadingStatus {
   SCHEDULED = 'SCHEDULED',
   CANCELED = 'CANCELED',
 }
+
+export interface DirectQueryRequest {
+  query: string;
+  lang: string;
+  datasource: string;
+  sessionId?: string;
+}
diff --git a/common/utils/query_session_utils.ts b/common/utils/query_session_utils.ts
@@ -0,0 +1,16 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { ASYNC_QUERY_SESSION_ID } from '../constants/shared';
+
+export const setAsyncSessionId = (value: string | null) => {
+  if (value !== null) {
+    sessionStorage.setItem(ASYNC_QUERY_SESSION_ID, value);
+  }
+};
+
+export const getAsyncSessionId = () => {
+  return sessionStorage.getItem(ASYNC_QUERY_SESSION_ID);
+};
diff --git a/common/utils/shared.ts b/common/utils/shared.ts
@@ -0,0 +1,8 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export function get<T = unknown>(obj: Record<string, any>, path: string, defaultValue?: T): T {
+  return path.split('.').reduce((acc: any, part: string) => acc && acc[part], obj) || defaultValue;
+}
diff --git a/public/components/common/search/sql_search.tsx b/public/components/common/search/sql_search.tsx
@@ -18,11 +18,11 @@ import {
   EuiPopoverFooter,
   EuiToolTip,
 } from '@elastic/eui';
-import { isEqual, lowerCase } from 'lodash';
+import { isEqual } from 'lodash';
 import React, { useEffect, useState } from 'react';
 import { batch, useDispatch, useSelector } from 'react-redux';
 import { APP_ANALYTICS_TAB_ID_REGEX, RAW_QUERY } from '../../../../common/constants/explorer';
-import { DirectQueryLoadingStatus } from '../../../../common/types/explorer';
+import { DirectQueryLoadingStatus, DirectQueryRequest } from '../../../../common/types/explorer';
 import { PPL_NEWLINE_REGEX, PPL_SPAN_REGEX } from '../../../../common/constants/shared';
 import { uiSettingsService } from '../../../../common/utils';
 import { useFetchEvents } from '../../../components/event_analytics/hooks';
@@ -38,6 +38,8 @@ import { PPLReferenceFlyout } from '../helpers';
 import { Autocomplete } from './autocomplete';
 import { changeQuery } from '../../../components/event_analytics/redux/slices/query_slice';
 import { QUERY_LANGUAGE } from '../../../../common/constants/data_sources';
+import { getAsyncSessionId, setAsyncSessionId } from '../../../../common/utils/query_session_utils';
+import { get as getObjValue } from '../../../../common/utils/shared';
 export interface IQueryBarProps {
   query: string;
   tempQuery: string;
@@ -101,7 +103,7 @@ export const DirectSearch = (props: any) => {
     stopPolling,
   } = usePolling<any, any>((params) => {
     return sqlService.fetchWithJobId(params);
-  }, 5000);
+  }, 2000);
 
   const requestParams = { tabId };
   const { dispatchOnGettingHis } = useFetchEvents({
@@ -190,13 +192,21 @@ export const DirectSearch = (props: any) => {
       );
     });
     dispatch(updateSearchMetaData({ tabId, data: { isPolling: true, lang } }));
+    const sessionId = getAsyncSessionId();
+    const requestPayload = {
+      lang: lang.toLowerCase(),
+      query: tempQuery || query,
+      datasource: explorerSearchMetadata.datasources[0].label,
+    } as DirectQueryRequest;
+
+    if (sessionId) {
+      requestPayload.sessionId = sessionId;
+    }
+
     sqlService
-      .fetch({
-        lang: lowerCase(lang),
-        query: tempQuery || query,
-        datasource: explorerSearchMetadata.datasources[0].name,
-      })
+      .fetch(requestPayload)
       .then((result) => {
+        setAsyncSessionId(getObjValue(result, 'sessionId', null));
         if (result.queryId) {
           startPolling({
             queryId: result.queryId,
@@ -208,8 +218,7 @@ export const DirectSearch = (props: any) => {
       .catch((e) => {
         setIsQueryRunning(false);
         console.error(e);
-      })
-      .finally(() => {});
+      });
   };
 
   useEffect(() => {

diff --git a/public/components/event_analytics/explorer/datasources/datasources_selection.tsx b/public/components/event_analytics/explorer/datasources/datasources_selection.tsx
@@ -23,7 +23,7 @@ import { reset as resetQueryResults } from '../../redux/slices/query_result_slic
 import { reset as resetVisualization } from '../../redux/slices/visualization_slice';
 import { reset as resetVisConfig } from '../../redux/slices/viualization_config_slice';
 import { reset as resetQuery } from '../../redux/slices/query_slice';
-import { SelectedDataSource } from '../../../../../common/types/explorer';
+import { DirectQueryRequest, SelectedDataSource } from '../../../../../common/types/explorer';
 import { ObservabilityDefaultDataSource } from '../../../../framework/datasources/obs_opensearch_datasource';
 import {
   DATA_SOURCE_TYPE_URL_PARAM_KEY,
@@ -32,7 +32,16 @@ import {
   DEFAULT_DATA_SOURCE_TYPE,
   DEFAULT_DATA_SOURCE_TYPE_NAME,
   DEFAULT_DATA_SOURCE_OBSERVABILITY_DISPLAY_NAME,
+  DATA_SOURCE_TYPES,
+  QUERY_LANGUAGE,
 } from '../../../../../common/constants/data_sources';
+import { SQLService } from '../../../../services/requests/sql';
+import { get as getObjValue } from '../../../../../common/utils/shared';
+import {
+  setAsyncSessionId,
+  getAsyncSessionId,
+} from '../../../../../common/utils/query_session_utils';
+import { DIRECT_DUMMY_QUERY } from '../../../../../common/constants/shared';
 
 const getDataSourceState = (selectedSourceState: SelectedDataSource[]) => {
   if (selectedSourceState.length === 0) return [];
@@ -70,7 +79,8 @@ const removeDataSourceFromURLParams = (currURL: string) => {
 };
 
 export const DataSourceSelection = ({ tabId }: { tabId: string }) => {
-  const { dataSources } = coreRefs;
+  const { dataSources, http } = coreRefs;
+  const sqlService = new SQLService(http!);
   const dispatch = useDispatch();
   const routerContext = useContext(LogExplorerRouterContext);
   const explorerSearchMetadata = useSelector(selectSearchMetaData)[tabId];
@@ -109,6 +119,23 @@ export const DataSourceSelection = ({ tabId }: { tabId: string }) => {
     setSelectedSources(selectedSource);
   };
 
+  const runDummyQuery = (dataSource: string) => {
+    const requestPayload = {
+      lang: QUERY_LANGUAGE.SQL.toLowerCase(),
+      query: DIRECT_DUMMY_QUERY,
+      datasource: dataSource,
+    } as DirectQueryRequest;
+
+    sqlService
+      .fetch(requestPayload)
+      .then((result) => {
+        setAsyncSessionId(getObjValue(result, 'sessionId', null));
+      })
+      .catch((e) => {
+        console.error(e);
+      });
+  };
+
   useEffect(() => {
     setSelectedSources(getDataSourceState(explorerSearchMetadata.datasources));
   }, [explorerSearchMetadata.datasources]);
@@ -166,6 +193,19 @@ export const DataSourceSelection = ({ tabId }: { tabId: string }) => {
     }
   }, []);
 
+  useEffect(() => {
+    // Execute a dummy query to initialize the cluster and obtain a sessionId for subsequent queries.
+    const dsType = explorerSearchMetadata.datasources?.[0]?.type;
+    const dsName = explorerSearchMetadata.datasources?.[0]?.label;
+    if (
+      !getAsyncSessionId() &&
+      [DATA_SOURCE_TYPES.SPARK, DATA_SOURCE_TYPES.S3Glue].includes(dsType) &&
+      dsName
+    ) {
+      runDummyQuery(dsName);
+    }
+  }, [explorerSearchMetadata.datasources]);
+
   /**
    * Process the data source options to display different than discover's group names.
    * Temporary solution for version 2.11.

diff --git a/public/services/requests/sql.ts b/public/services/requests/sql.ts
@@ -4,22 +4,15 @@
  */
 
 import { CoreStart } from '../../../../../src/core/public';
-import { PPL_BASE, PPL_SEARCH } from '../../../common/constants/shared';
+import { DirectQueryRequest } from '../../../common/types/explorer';
 
 export class SQLService {
   private http;
   constructor(http: CoreStart['http']) {
     this.http = http;
   }
 
-  fetch = async (
-    params: {
-      query: string;
-      lang: string;
-      datasource: string;
-    },
-    errorHandler?: (error: any) => void
-  ) => {
+  fetch = async (params: DirectQueryRequest, errorHandler?: (error: any) => void) => {
     return this.http
       .post('/api/observability/query/jobs', {
         body: JSON.stringify(params),

diff --git a/server/routes/datasources/datasources_router.ts b/server/routes/datasources/datasources_router.ts
@@ -16,6 +16,7 @@ export function registerDatasourcesRoute(router: IRouter) {
           query: schema.string(),
           lang: schema.string(),
           datasource: schema.string(),
+          sessionId: schema.maybe(schema.string()),
         }),
       },
     },