From 365a2fdf990e756d7772bf79fcefb1672566dcfd Mon Sep 17 00:00:00 2001 From: Eric Wei Date: Tue, 24 Oct 2023 20:42:50 -0700 Subject: [PATCH] [Explorer] Supports session for s3 direct query (#1178) * remove unused files Signed-off-by: Eric Wei * missing snapshots Signed-off-by: Eric Wei * remove unused files Signed-off-by: Eric Wei * supports session id for direct query Signed-off-by: Eric * update snapshots Signed-off-by: Eric * remove logging Signed-off-by: Eric --------- Signed-off-by: Eric Wei Signed-off-by: Eric --- common/constants/data_sources.ts | 5 +++ common/constants/shared.ts | 4 ++ common/types/explorer.ts | 7 +++ common/utils/query_session_utils.ts | 16 +++++++ common/utils/shared.ts | 8 ++++ .../components/common/search/sql_search.tsx | 29 +++++++----- .../datasources/datasources_selection.tsx | 44 ++++++++++++++++++- public/services/requests/sql.ts | 11 +---- .../routes/datasources/datasources_router.ts | 1 + 9 files changed, 104 insertions(+), 21 deletions(-) create mode 100644 common/utils/query_session_utils.ts create mode 100644 common/utils/shared.ts diff --git a/common/constants/data_sources.ts b/common/constants/data_sources.ts index aa596ad10..ad5c0cfdf 100644 --- 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 index 198ce05a3..211645485 100644 --- 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 index a36dd81a8..9daf69074 100644 --- 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 new file mode 100644 index 000000000..ad0e0a1af --- /dev/null +++ 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 new file mode 100644 index 000000000..a6f226a6d --- /dev/null +++ b/common/utils/shared.ts @@ -0,0 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export function get(obj: Record, 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 index 5a0b18ae7..0ba84061d 100644 --- 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((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 index bbea0a628..3c6e34df1 100644 --- 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 index e87d4d726..2073804ad 100644 --- a/public/services/requests/sql.ts +++ b/public/services/requests/sql.ts @@ -4,7 +4,7 @@ */ 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; @@ -12,14 +12,7 @@ export class SQLService { 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 index eb8d6a635..45279ccf7 100644 --- 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()), }), }, },