diff --git a/changelogs/fragments/7212.yml b/changelogs/fragments/7212.yml new file mode 100644 index 00000000000..597428db472 --- /dev/null +++ b/changelogs/fragments/7212.yml @@ -0,0 +1,2 @@ +feat: +- Add query enhancements plugin as a core plugin ([#7212](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/7212)) \ No newline at end of file diff --git a/docs/_sidebar.md b/docs/_sidebar.md index 66a259b8c2f..e5136a6df4c 100644 --- a/docs/_sidebar.md +++ b/docs/_sidebar.md @@ -59,6 +59,7 @@ - forms - [Form_wizard](../src/plugins/opensearch_ui_shared/public/forms/form_wizard/README.md) - [Multi_content](../src/plugins/opensearch_ui_shared/public/forms/multi_content/README.md) + - [Query_enhancements](../src/plugins/query_enhancements/README.md) - [Saved_objects](../src/plugins/saved_objects/README.md) - [Saved_objects_management](../src/plugins/saved_objects_management/README.md) - [Share](../src/plugins/share/README.md) diff --git a/src/plugins/data/public/index.ts b/src/plugins/data/public/index.ts index b8200101888..f1ac419e9ec 100644 --- a/src/plugins/data/public/index.ts +++ b/src/plugins/data/public/index.ts @@ -437,6 +437,12 @@ export { IndexPatternSelectProps, QueryStringInput, QueryStringInputProps, + QueryEditor, + QueryEditorExtensionConfig, + QueryEditorExtensions, + QueryEditorExtensionDependencies, + QueryEditorProps, + QueryEditorTopRow, // for BWC, keeping the old name IUiStart as DataPublicPluginStartUi, } from './ui'; diff --git a/src/plugins/data/public/ui/index.ts b/src/plugins/data/public/ui/index.ts index 582470cc997..5483b540d5b 100644 --- a/src/plugins/data/public/ui/index.ts +++ b/src/plugins/data/public/ui/index.ts @@ -39,5 +39,13 @@ export { export { IndexPatternSelectProps } from './index_pattern_select'; export { FilterLabel } from './filter_bar'; export { QueryStringInput, QueryStringInputProps } from './query_string_input'; +export { + QueryEditorTopRow, + QueryEditor, + QueryEditorProps, + QueryEditorExtensions, + QueryEditorExtensionDependencies, + QueryEditorExtensionConfig, +} from './query_editor'; export { SearchBar, SearchBarProps, StatefulSearchBarProps } from './search_bar'; export { SuggestionsComponent } from './typeahead'; diff --git a/src/plugins/data/public/ui/query_editor/index.tsx b/src/plugins/data/public/ui/query_editor/index.tsx index bddef49af1d..52fba0ecea7 100644 --- a/src/plugins/data/public/ui/query_editor/index.tsx +++ b/src/plugins/data/public/ui/query_editor/index.tsx @@ -25,4 +25,8 @@ export const QueryEditor = (props: QueryEditorProps) => ( ); export type { QueryEditorProps }; -export { QueryEditorExtensions, QueryEditorExtensionConfig } from './query_editor_extensions'; +export { + QueryEditorExtensions, + QueryEditorExtensionDependencies, + QueryEditorExtensionConfig, +} from './query_editor_extensions'; 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 19c4c527038..59d5645fcf6 100644 --- a/src/plugins/data/public/ui/query_editor/query_editor.tsx +++ b/src/plugins/data/public/ui/query_editor/query_editor.tsx @@ -215,9 +215,9 @@ export default class QueryEditorUI extends Component { : undefined; this.onChange(newQuery, dateRange); this.onSubmit(newQuery, dateRange); - this.setState({ isDataSetsVisible: enhancement?.searchBar?.showDataSetsSelector ?? true }); this.setState({ isDataSourcesVisible: enhancement?.searchBar?.showDataSourcesSelector ?? true, + isDataSetsVisible: enhancement?.searchBar?.showDataSetsSelector ?? true, }); }; @@ -231,19 +231,15 @@ export default class QueryEditorUI extends Component { private initDataSourcesVisibility = () => { if (this.componentIsUnmounting) return; - const isDataSourcesVisible = - this.props.settings.getQueryEnhancements(this.props.query.language)?.searchBar - ?.showDataSourcesSelector ?? true; - this.setState({ isDataSourcesVisible }); + return this.props.settings.getQueryEnhancements(this.props.query.language)?.searchBar + ?.showDataSourcesSelector; }; private initDataSetsVisibility = () => { if (this.componentIsUnmounting) return; - const isDataSetsVisible = - this.props.settings.getQueryEnhancements(this.props.query.language)?.searchBar - ?.showDataSetsSelector ?? true; - this.setState({ isDataSetsVisible }); + return this.props.settings.getQueryEnhancements(this.props.query.language)?.searchBar + ?.showDataSetsSelector; }; public onMouseEnterSuggestion = (index: number) => { @@ -260,8 +256,10 @@ export default class QueryEditorUI extends Component { this.initPersistedLog(); // this.fetchIndexPatterns().then(this.updateSuggestions); - this.initDataSourcesVisibility(); - this.initDataSetsVisibility(); + this.setState({ + isDataSourcesVisible: this.initDataSourcesVisibility() || true, + isDataSetsVisible: this.initDataSetsVisibility() || true, + }); } public componentDidUpdate(prevProps: Props) { diff --git a/src/plugins/data/public/ui/query_editor/query_editor_extensions/index.tsx b/src/plugins/data/public/ui/query_editor/query_editor_extensions/index.tsx index f406423d616..8dcaaa459dd 100644 --- a/src/plugins/data/public/ui/query_editor/query_editor_extensions/index.tsx +++ b/src/plugins/data/public/ui/query_editor/query_editor_extensions/index.tsx @@ -14,4 +14,7 @@ export const QueryEditorExtensions = (props: ComponentProps ); -export { QueryEditorExtensionConfig } from './query_editor_extension'; +export { + QueryEditorExtensionDependencies, + QueryEditorExtensionConfig, +} from './query_editor_extension'; diff --git a/src/plugins/query_enhancements/.i18nrc.json b/src/plugins/query_enhancements/.i18nrc.json new file mode 100644 index 00000000000..bb9a3ef5e50 --- /dev/null +++ b/src/plugins/query_enhancements/.i18nrc.json @@ -0,0 +1,7 @@ +{ + "prefix": "queryEnhancements", + "paths": { + "queryEnhancements": "." + }, + "translations": ["translations/ja-JP.json"] +} diff --git a/src/plugins/query_enhancements/README.md b/src/plugins/query_enhancements/README.md new file mode 100755 index 00000000000..f7c6326fd09 --- /dev/null +++ b/src/plugins/query_enhancements/README.md @@ -0,0 +1,9 @@ +# Query Enhancements Plugin + +Optional plugin, that registers query enhancing capabilities within +the application. + +## List of enhancements + +* PPL within Discover +* SQL within Discover diff --git a/src/plugins/query_enhancements/common/config.ts b/src/plugins/query_enhancements/common/config.ts new file mode 100644 index 00000000000..b9ea4750e60 --- /dev/null +++ b/src/plugins/query_enhancements/common/config.ts @@ -0,0 +1,23 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { schema, TypeOf } from '@osd/config-schema'; + +export const configSchema = schema.object({ + enabled: schema.boolean({ defaultValue: true }), + queryAssist: schema.object({ + supportedLanguages: schema.arrayOf( + schema.object({ + language: schema.string(), + agentConfig: schema.string(), + }), + { + defaultValue: [{ language: 'PPL', agentConfig: 'os_query_assist_ppl' }], + } + ), + }), +}); + +export type ConfigSchema = TypeOf; diff --git a/src/plugins/query_enhancements/common/constants.ts b/src/plugins/query_enhancements/common/constants.ts new file mode 100644 index 00000000000..7e82677407d --- /dev/null +++ b/src/plugins/query_enhancements/common/constants.ts @@ -0,0 +1,47 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export const PLUGIN_ID = 'queryEnhancements'; +export const PLUGIN_NAME = 'queryEnhancements'; + +export const BASE_API = '/api/enhancements'; + +export const SEARCH_STRATEGY = { + PPL: 'ppl', + SQL: 'sql', + SQL_ASYNC: 'sqlasync', +}; + +export const API = { + SEARCH: `${BASE_API}/search`, + PPL_SEARCH: `${BASE_API}/search/${SEARCH_STRATEGY.PPL}`, + SQL_SEARCH: `${BASE_API}/search/${SEARCH_STRATEGY.SQL}`, + SQL_ASYNC_SEARCH: `${BASE_API}/search/${SEARCH_STRATEGY.SQL_ASYNC}`, + QUERY_ASSIST: { + LANGUAGES: `${BASE_API}/assist/languages`, + GENERATE: `${BASE_API}/assist/generate`, + }, + DATA_SOURCE: { + CONNECTIONS: `${BASE_API}/datasource/connections`, + }, +}; + +export const URI = { + PPL: '/_plugins/_ppl', + SQL: '/_plugins/_sql', + ASYNC_QUERY: '/_plugins/_async_query', + ML: '/_plugins/_ml', + OBSERVABILITY: '/_plugins/_observability', + DATA_CONNECTIONS: '/_plugins/_query/_datasources', +}; + +export const OPENSEARCH_API = { + PANELS: `${URI.OBSERVABILITY}/object`, + DATA_CONNECTIONS: URI.DATA_CONNECTIONS, +}; + +export const UI_SETTINGS = {}; + +export const ERROR_DETAILS = { GUARDRAILS_TRIGGERED: 'guardrails triggered' }; diff --git a/src/plugins/query_enhancements/common/index.ts b/src/plugins/query_enhancements/common/index.ts new file mode 100644 index 00000000000..f21e3c24507 --- /dev/null +++ b/src/plugins/query_enhancements/common/index.ts @@ -0,0 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export * from './constants'; +export * from './types'; +export * from './utils'; diff --git a/src/plugins/query_enhancements/common/query_assist/index.ts b/src/plugins/query_enhancements/common/query_assist/index.ts new file mode 100644 index 00000000000..9469a3a2771 --- /dev/null +++ b/src/plugins/query_enhancements/common/query_assist/index.ts @@ -0,0 +1,6 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export { QueryAssistParameters, QueryAssistResponse } from './types'; diff --git a/src/plugins/query_enhancements/common/query_assist/types.ts b/src/plugins/query_enhancements/common/query_assist/types.ts new file mode 100644 index 00000000000..057ff1e708d --- /dev/null +++ b/src/plugins/query_enhancements/common/query_assist/types.ts @@ -0,0 +1,19 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { TimeRange } from 'src/plugins/data/common'; + +export interface QueryAssistResponse { + query: string; + timeRange?: TimeRange; +} + +export interface QueryAssistParameters { + question: string; + index: string; + language: string; + // for MDS + dataSourceId?: string; +} diff --git a/src/plugins/query_enhancements/common/types.ts b/src/plugins/query_enhancements/common/types.ts new file mode 100644 index 00000000000..c98ca028496 --- /dev/null +++ b/src/plugins/query_enhancements/common/types.ts @@ -0,0 +1,15 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { CoreSetup } from 'opensearch-dashboards/public'; +import { Observable } from 'rxjs'; + +export interface FetchDataFrameContext { + http: CoreSetup['http']; + path: string; + signal?: AbortSignal; +} + +export type FetchFunction = (params?: P) => Observable; diff --git a/src/plugins/query_enhancements/common/utils.ts b/src/plugins/query_enhancements/common/utils.ts new file mode 100644 index 00000000000..f4bdde2a26e --- /dev/null +++ b/src/plugins/query_enhancements/common/utils.ts @@ -0,0 +1,156 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { IDataFrame } from 'src/plugins/data/common'; +import { Observable, Subscription, from, throwError, timer } from 'rxjs'; +import { catchError, concatMap, last, takeWhile, tap } from 'rxjs/operators'; +import { FetchDataFrameContext, FetchFunction } from './types'; + +export const formatDate = (dateString: string) => { + const date = new Date(dateString); + return ( + date.getFullYear() + + '-' + + ('0' + (date.getMonth() + 1)).slice(-2) + + '-' + + ('0' + date.getDate()).slice(-2) + + ' ' + + ('0' + date.getHours()).slice(-2) + + ':' + + ('0' + date.getMinutes()).slice(-2) + + ':' + + ('0' + date.getSeconds()).slice(-2) + ); +}; + +export const getFields = (rawResponse: any) => { + return rawResponse.data.schema?.map((field: any, index: any) => ({ + ...field, + values: rawResponse.data.datarows?.map((row: any) => row[index]), + })); +}; + +export const removeKeyword = (queryString: string | undefined) => { + return queryString?.replace(new RegExp('.keyword'), '') ?? ''; +}; + +export class DataFramePolling { + public data: T | null = null; + public error: Error | null = null; + public loading: boolean = true; + private shouldPoll: boolean = false; + private intervalRef?: NodeJS.Timeout; + private subscription?: Subscription; + + constructor( + private fetchFunction: FetchFunction, + private interval: number = 5000, + private onPollingSuccess: (data: T) => boolean, + private onPollingError: (error: Error) => boolean + ) {} + + fetch(): Observable { + return timer(0, this.interval).pipe( + concatMap(() => this.fetchFunction()), + takeWhile((resp) => this.onPollingSuccess(resp), true), + tap((resp: T) => { + this.data = resp; + }), + last(), + catchError((error: Error) => { + this.onPollingError(error); + return throwError(error); + }) + ); + } + + fetchData(params?: P) { + this.loading = true; + this.subscription = this.fetchFunction(params).subscribe({ + next: (result: any) => { + this.data = result; + this.loading = false; + + if (this.onPollingSuccess && this.onPollingSuccess(result)) { + this.stopPolling(); + } + }, + error: (err: any) => { + this.error = err as Error; + this.loading = false; + + if (this.onPollingError && this.onPollingError(this.error)) { + this.stopPolling(); + } + }, + }); + } + + startPolling(params?: P) { + this.shouldPoll = true; + if (!this.intervalRef) { + this.intervalRef = setInterval(() => { + if (this.shouldPoll) { + this.fetchData(params); + } + }, this.interval); + } + } + + stopPolling() { + this.shouldPoll = false; + if (this.intervalRef) { + clearInterval(this.intervalRef); + this.intervalRef = undefined; + } + if (this.subscription) { + this.subscription.unsubscribe(); + this.subscription = undefined; + } + } + + waitForPolling(): Promise { + return new Promise((resolve) => { + const checkLoading = () => { + if (!this.loading) { + resolve(this.data); + } else { + setTimeout(checkLoading, this.interval); + } + }; + checkLoading(); + }); + } +} + +export const fetchDataFrame = ( + context: FetchDataFrameContext, + queryString: string, + df: IDataFrame +) => { + const { http, path, signal } = context; + const body = JSON.stringify({ query: { qs: queryString, format: 'jdbc' }, df }); + return from( + http.fetch({ + method: 'POST', + path, + body, + signal, + }) + ); +}; + +export const fetchDataFramePolling = (context: FetchDataFrameContext, df: IDataFrame) => { + const { http, path, signal } = context; + const queryId = df.meta?.queryId; + const dataSourceId = df.meta?.queryConfig?.dataSourceId; + return from( + http.fetch({ + method: 'GET', + path: `${path}/${queryId}${dataSourceId ? `/${dataSourceId}` : ''}`, + signal, + }) + ); +}; diff --git a/src/plugins/query_enhancements/opensearch_dashboards.json b/src/plugins/query_enhancements/opensearch_dashboards.json new file mode 100644 index 00000000000..e6ed7e2e0a1 --- /dev/null +++ b/src/plugins/query_enhancements/opensearch_dashboards.json @@ -0,0 +1,9 @@ + { + "id": "queryEnhancements", + "version": "opensearchDashboards", + "server": true, + "ui": true, + "requiredPlugins": ["data", "opensearchDashboardsReact", "opensearchDashboardsUtils", "dataSource", "dataSourceManagement", "savedObjects", "uiActions"], + "optionalPlugins": [] +} + diff --git a/src/plugins/query_enhancements/public/assets/query_assist_mark.svg b/src/plugins/query_enhancements/public/assets/query_assist_mark.svg new file mode 100644 index 00000000000..b744e8c35e8 --- /dev/null +++ b/src/plugins/query_enhancements/public/assets/query_assist_mark.svg @@ -0,0 +1,18 @@ + + + + + + + + + + + + + + + + + + diff --git a/src/plugins/query_enhancements/public/data_source_connection/components/connections_bar.tsx b/src/plugins/query_enhancements/public/data_source_connection/components/connections_bar.tsx new file mode 100644 index 00000000000..d7590b27822 --- /dev/null +++ b/src/plugins/query_enhancements/public/data_source_connection/components/connections_bar.tsx @@ -0,0 +1,93 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import React, { useEffect, useRef, useState } from 'react'; +import { EuiPortal } from '@elastic/eui'; +import { distinctUntilChanged } from 'rxjs/operators'; +import { ToastsSetup } from 'opensearch-dashboards/public'; +import { DataPublicPluginStart, QueryEditorExtensionDependencies } from '../../../../data/public'; +import { DataSourceSelector } from '../../../../data_source_management/public'; +import { ConnectionsService } from '../services'; + +interface ConnectionsProps { + dependencies: QueryEditorExtensionDependencies; + toasts: ToastsSetup; + connectionsService: ConnectionsService; +} + +export const ConnectionsBar: React.FC = ({ connectionsService, toasts }) => { + const [isDataSourceEnabled, setIsDataSourceEnabled] = useState(false); + const [uiService, setUiService] = useState(undefined); + const containerRef = useRef(null); + + useEffect(() => { + const uiServiceSubscription = connectionsService.getUiService().subscribe(setUiService); + const dataSourceEnabledSubscription = connectionsService + .getIsDataSourceEnabled$() + .subscribe(setIsDataSourceEnabled); + + return () => { + uiServiceSubscription.unsubscribe(); + dataSourceEnabledSubscription.unsubscribe(); + }; + }, [connectionsService]); + + useEffect(() => { + if (!uiService || !isDataSourceEnabled || !containerRef.current) return; + const subscriptions = uiService.dataSourceContainer$.subscribe((container) => { + if (container && containerRef.current) { + container.append(containerRef.current); + } + }); + + return () => subscriptions.unsubscribe(); + }, [uiService, isDataSourceEnabled]); + + useEffect(() => { + const selectedConnectionSubscription = connectionsService + .getSelectedConnection$() + .pipe(distinctUntilChanged()) + .subscribe((connection) => { + if (connection) { + // Assuming setSelectedConnection$ is meant to update some state or perform an action outside this component + connectionsService.setSelectedConnection$(connection); + } + }); + + return () => selectedConnectionSubscription.unsubscribe(); + }, [connectionsService]); + + const handleSelectedConnection = (id: string | undefined) => { + if (!id) { + connectionsService.setSelectedConnection$(undefined); + return; + } + connectionsService.getConnectionById(id).subscribe((connection) => { + connectionsService.setSelectedConnection$(connection); + }); + }; + + return ( + { + containerRef.current = node; + }} + > +
+ + handleSelectedConnection(dataSource[0]?.id || undefined) + } + /> +
+ + ); +}; diff --git a/src/plugins/query_enhancements/public/data_source_connection/components/index.ts b/src/plugins/query_enhancements/public/data_source_connection/components/index.ts new file mode 100644 index 00000000000..1ee969a1d07 --- /dev/null +++ b/src/plugins/query_enhancements/public/data_source_connection/components/index.ts @@ -0,0 +1,6 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export { ConnectionsBar } from './connections_bar'; diff --git a/src/plugins/query_enhancements/public/data_source_connection/index.ts b/src/plugins/query_enhancements/public/data_source_connection/index.ts new file mode 100644 index 00000000000..e334163d91d --- /dev/null +++ b/src/plugins/query_enhancements/public/data_source_connection/index.ts @@ -0,0 +1,7 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export { createDataSourceConnectionExtension } from './utils'; +export * from './services'; diff --git a/src/plugins/query_enhancements/public/data_source_connection/services/connections_service.ts b/src/plugins/query_enhancements/public/data_source_connection/services/connections_service.ts new file mode 100644 index 00000000000..6afec4b51a9 --- /dev/null +++ b/src/plugins/query_enhancements/public/data_source_connection/services/connections_service.ts @@ -0,0 +1,84 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { BehaviorSubject, Observable, from } from 'rxjs'; +import { DataPublicPluginStart } from 'src/plugins/data/public'; +import { CoreStart } from 'opensearch-dashboards/public'; +import { API } from '../../../common'; +import { Connection, ConnectionsServiceDeps } from '../../types'; + +export class ConnectionsService { + protected http!: ConnectionsServiceDeps['http']; + protected savedObjects!: CoreStart['savedObjects']; + private uiService$ = new BehaviorSubject(undefined); + private isDataSourceEnabled = false; + private isDataSourceEnabled$ = new BehaviorSubject(this.isDataSourceEnabled); + private selectedConnection: Connection | undefined = undefined; + private selectedConnection$ = new BehaviorSubject( + this.selectedConnection + ); + + constructor(deps: ConnectionsServiceDeps) { + deps.startServices.then(([coreStart, depsStart]) => { + this.http = deps.http; + this.savedObjects = coreStart.savedObjects; + this.uiService$.next(depsStart.data.ui); + this.setIsDataSourceEnabled$(depsStart.dataSource?.dataSourceEnabled || false); + }); + } + + getSavedObjects = () => { + return this.savedObjects; + }; + + getIsDataSourceEnabled = () => { + return this.isDataSourceEnabled; + }; + + setIsDataSourceEnabled$ = (isDataSourceEnabled: boolean) => { + this.isDataSourceEnabled = isDataSourceEnabled; + this.isDataSourceEnabled$.next(this.isDataSourceEnabled); + }; + + getIsDataSourceEnabled$ = () => { + return this.isDataSourceEnabled$.asObservable(); + }; + + getUiService = () => { + return this.uiService$.asObservable(); + }; + + getConnections = (): Observable => { + return from( + this.http.fetch({ + method: 'GET', + path: API.DATA_SOURCE.CONNECTIONS, + }) + ); + }; + + getConnectionById = (id: string): Observable => { + const path = `${API.DATA_SOURCE.CONNECTIONS}/${id}`; + return from( + this.http.fetch({ + method: 'GET', + path, + }) + ); + }; + + getSelectedConnection = () => { + return this.selectedConnection; + }; + + setSelectedConnection$ = (connection: Connection | undefined) => { + this.selectedConnection = connection; + this.selectedConnection$.next(this.selectedConnection); + }; + + getSelectedConnection$ = () => { + return this.selectedConnection$.asObservable(); + }; +} diff --git a/src/plugins/query_enhancements/public/data_source_connection/services/index.ts b/src/plugins/query_enhancements/public/data_source_connection/services/index.ts new file mode 100644 index 00000000000..08eeda5a7aa --- /dev/null +++ b/src/plugins/query_enhancements/public/data_source_connection/services/index.ts @@ -0,0 +1,6 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export { ConnectionsService } from './connections_service'; diff --git a/src/plugins/query_enhancements/public/data_source_connection/utils/create_extension.tsx b/src/plugins/query_enhancements/public/data_source_connection/utils/create_extension.tsx new file mode 100644 index 00000000000..e5822c4b378 --- /dev/null +++ b/src/plugins/query_enhancements/public/data_source_connection/utils/create_extension.tsx @@ -0,0 +1,34 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import React from 'react'; +import { ToastsSetup } from 'opensearch-dashboards/public'; +import { QueryEditorExtensionConfig } from '../../../../data/public'; +import { ConfigSchema } from '../../../common/config'; +import { ConnectionsBar } from '../components'; +import { ConnectionsService } from '../services'; + +export const createDataSourceConnectionExtension = ( + connectionsService: ConnectionsService, + toasts: ToastsSetup, + config: ConfigSchema +): QueryEditorExtensionConfig => { + return { + id: 'data-source-connection', + order: 2000, + isEnabled$: (dependencies) => { + return connectionsService.getIsDataSourceEnabled$(); + }, + getComponent: (dependencies) => { + return ( + + ); + }, + }; +}; diff --git a/src/plugins/query_enhancements/public/data_source_connection/utils/index.ts b/src/plugins/query_enhancements/public/data_source_connection/utils/index.ts new file mode 100644 index 00000000000..9eccc9e6f35 --- /dev/null +++ b/src/plugins/query_enhancements/public/data_source_connection/utils/index.ts @@ -0,0 +1,6 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export * from './create_extension'; diff --git a/src/plugins/query_enhancements/public/index.scss b/src/plugins/query_enhancements/public/index.scss new file mode 100644 index 00000000000..ff7112406ea --- /dev/null +++ b/src/plugins/query_enhancements/public/index.scss @@ -0,0 +1 @@ +/* stylelint-disable no-empty-source */ diff --git a/src/plugins/query_enhancements/public/index.ts b/src/plugins/query_enhancements/public/index.ts new file mode 100644 index 00000000000..8f141c20e69 --- /dev/null +++ b/src/plugins/query_enhancements/public/index.ts @@ -0,0 +1,14 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { PluginInitializerContext } from '../../../core/public'; +import './index.scss'; +import { QueryEnhancementsPlugin } from './plugin'; + +export function plugin(initializerContext: PluginInitializerContext) { + return new QueryEnhancementsPlugin(initializerContext); +} + +export { QueryEnhancementsPluginSetup, QueryEnhancementsPluginStart } from './types'; diff --git a/src/plugins/query_enhancements/public/plugin.tsx b/src/plugins/query_enhancements/public/plugin.tsx new file mode 100644 index 00000000000..0ea557db8ce --- /dev/null +++ b/src/plugins/query_enhancements/public/plugin.tsx @@ -0,0 +1,181 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import moment from 'moment'; +import { CoreSetup, CoreStart, Plugin, PluginInitializerContext } from '../../../core/public'; +import { IStorageWrapper, Storage } from '../../opensearch_dashboards_utils/public'; +import { ConfigSchema } from '../common/config'; +import { ConnectionsService, createDataSourceConnectionExtension } from './data_source_connection'; +import { createQueryAssistExtension } from './query_assist'; +import { PPLSearchInterceptor, SQLAsyncSearchInterceptor, SQLSearchInterceptor } from './search'; +import { setData, setStorage } from './services'; +import { + QueryEnhancementsPluginSetup, + QueryEnhancementsPluginSetupDependencies, + QueryEnhancementsPluginStart, + QueryEnhancementsPluginStartDependencies, +} from './types'; + +export class QueryEnhancementsPlugin + implements + Plugin< + QueryEnhancementsPluginSetup, + QueryEnhancementsPluginStart, + QueryEnhancementsPluginSetupDependencies, + QueryEnhancementsPluginStartDependencies + > { + private readonly storage: IStorageWrapper; + private readonly config: ConfigSchema; + private connectionsService!: ConnectionsService; + + constructor(initializerContext: PluginInitializerContext) { + this.config = initializerContext.config.get(); + this.storage = new Storage(window.localStorage); + } + + public setup( + core: CoreSetup, + { data }: QueryEnhancementsPluginSetupDependencies + ): QueryEnhancementsPluginSetup { + this.connectionsService = new ConnectionsService({ + startServices: core.getStartServices(), + http: core.http, + }); + + const pplSearchInterceptor = new PPLSearchInterceptor( + { + toasts: core.notifications.toasts, + http: core.http, + uiSettings: core.uiSettings, + startServices: core.getStartServices(), + usageCollector: data.search.usageCollector, + }, + this.connectionsService + ); + + const sqlSearchInterceptor = new SQLSearchInterceptor( + { + toasts: core.notifications.toasts, + http: core.http, + uiSettings: core.uiSettings, + startServices: core.getStartServices(), + usageCollector: data.search.usageCollector, + }, + this.connectionsService + ); + + const sqlAsyncSearchInterceptor = new SQLAsyncSearchInterceptor( + { + toasts: core.notifications.toasts, + http: core.http, + uiSettings: core.uiSettings, + startServices: core.getStartServices(), + usageCollector: data.search.usageCollector, + }, + this.connectionsService + ); + + data.__enhance({ + ui: { + query: { + language: 'PPL', + search: pplSearchInterceptor, + searchBar: { + queryStringInput: { initialValue: 'source=' }, + dateRange: { + initialFrom: moment().subtract(2, 'days').toISOString(), + initialTo: moment().add(2, 'days').toISOString(), + }, + showFilterBar: false, + showDataSetsSelector: false, + showDataSourcesSelector: true, + }, + fields: { + filterable: false, + visualizable: false, + }, + supportedAppNames: ['discover'], + }, + }, + }); + + data.__enhance({ + ui: { + query: { + language: 'SQL', + search: sqlSearchInterceptor, + searchBar: { + showDatePicker: false, + showFilterBar: false, + showDataSetsSelector: false, + showDataSourcesSelector: true, + queryStringInput: { initialValue: 'SELECT * FROM ' }, + }, + fields: { + filterable: false, + visualizable: false, + }, + showDocLinks: false, + supportedAppNames: ['discover'], + }, + }, + }); + + data.__enhance({ + ui: { + query: { + language: 'SQLAsync', + search: sqlAsyncSearchInterceptor, + searchBar: { + showDatePicker: false, + showFilterBar: false, + showDataSetsSelector: false, + showDataSourcesSelector: true, + queryStringInput: { initialValue: 'SHOW DATABASES IN ::mys3::' }, + }, + fields: { + filterable: false, + visualizable: false, + }, + showDocLinks: false, + supportedAppNames: ['discover'], + }, + }, + }); + + data.__enhance({ + ui: { + queryEditorExtension: createQueryAssistExtension( + core.http, + this.connectionsService, + this.config.queryAssist + ), + }, + }); + + data.__enhance({ + ui: { + queryEditorExtension: createDataSourceConnectionExtension( + this.connectionsService, + core.notifications.toasts, + this.config + ), + }, + }); + + return {}; + } + + public start( + core: CoreStart, + deps: QueryEnhancementsPluginStartDependencies + ): QueryEnhancementsPluginStart { + setStorage(this.storage); + setData(deps.data); + return {}; + } + + public stop() {} +} diff --git a/src/plugins/query_enhancements/public/query_assist/components/__snapshots__/call_outs.test.tsx.snap b/src/plugins/query_enhancements/public/query_assist/components/__snapshots__/call_outs.test.tsx.snap new file mode 100644 index 00000000000..cfe754d3ae0 --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/components/__snapshots__/call_outs.test.tsx.snap @@ -0,0 +1,122 @@ +// Jest Snapshot v1, https://goo.gl/fbAQLP + +exports[`CallOuts spec should display empty_index call out 1`] = ` +
+
+
+
+
+
+`; + +exports[`CallOuts spec should display empty_query call out 1`] = ` +
+
+
+
+
+
+`; + +exports[`CallOuts spec should display invalid_query call out 1`] = ` +
+
+
+
+
+
+`; + +exports[`CallOuts spec should display query_generated call out 1`] = ` +
+
+
+
+ +
+
+`; diff --git a/src/plugins/query_enhancements/public/query_assist/components/call_outs.test.tsx b/src/plugins/query_enhancements/public/query_assist/components/call_outs.test.tsx new file mode 100644 index 00000000000..4ac107d49c6 --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/components/call_outs.test.tsx @@ -0,0 +1,58 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { render } from '@testing-library/react'; +import React, { ComponentProps } from 'react'; +import { IntlProvider } from 'react-intl'; +import { QueryAssistCallOut } from './call_outs'; + +type Props = ComponentProps; + +const IntlWrapper = ({ children }: { children: unknown }) => ( + {children} +); + +const renderCallOut = (overrideProps: Partial = {}) => { + const props: Props = Object.assign>( + { + type: 'empty_query', + language: 'test lang', + onDismiss: jest.fn(), + }, + overrideProps + ); + const component = render(, { + wrapper: IntlWrapper, + }); + return { component, props: props as jest.MockedObjectDeep }; +}; + +describe('CallOuts spec', () => { + it('should display nothing if type is invalid', () => { + // @ts-expect-error testing invalid type + const { component } = renderCallOut({ type: '' }); + expect(component.container).toBeEmptyDOMElement(); + }); + + it('should display empty_query call out', () => { + const { component } = renderCallOut({ type: 'empty_query' }); + expect(component.container).toMatchSnapshot(); + }); + + it('should display empty_index call out', () => { + const { component } = renderCallOut({ type: 'empty_index' }); + expect(component.container).toMatchSnapshot(); + }); + + it('should display invalid_query call out', () => { + const { component } = renderCallOut({ type: 'invalid_query' }); + expect(component.container).toMatchSnapshot(); + }); + + it('should display query_generated call out', () => { + const { component } = renderCallOut({ type: 'query_generated' }); + expect(component.container).toMatchSnapshot(); + }); +}); diff --git a/src/plugins/query_enhancements/public/query_assist/components/call_outs.tsx b/src/plugins/query_enhancements/public/query_assist/components/call_outs.tsx new file mode 100644 index 00000000000..16d218bed67 --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/components/call_outs.tsx @@ -0,0 +1,106 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { EuiCallOut, EuiCallOutProps } from '@elastic/eui'; +import { FormattedMessage } from '@osd/i18n/react'; +import React from 'react'; + +interface QueryAssistCallOutProps extends Required> { + language: string; + type: QueryAssistCallOutType; +} + +export type QueryAssistCallOutType = + | undefined + | 'invalid_query' + | 'prohibited_query' + | 'empty_query' + | 'empty_index' + | 'query_generated'; + +const EmptyIndexCallOut: React.FC = (props) => ( + + } + size="s" + color="warning" + iconType="iInCircle" + dismissible + onDismiss={props.onDismiss} + /> +); + +const ProhibitedQueryCallOut: React.FC = (props) => ( + + } + size="s" + color="danger" + iconType="alert" + dismissible + onDismiss={props.onDismiss} + /> +); + +const EmptyQueryCallOut: React.FC = (props) => ( + + } + size="s" + color="warning" + iconType="iInCircle" + dismissible + onDismiss={props.onDismiss} + /> +); + +const QueryGeneratedCallOut: React.FC = (props) => ( + + } + size="s" + color="success" + iconType="check" + dismissible + onDismiss={props.onDismiss} + /> +); + +export const QueryAssistCallOut: React.FC = (props) => { + switch (props.type) { + case 'empty_query': + return ; + case 'empty_index': + return ; + case 'invalid_query': + return ; + case 'query_generated': + return ; + default: + break; + } + return null; +}; diff --git a/src/plugins/query_enhancements/public/query_assist/components/index.ts b/src/plugins/query_enhancements/public/query_assist/components/index.ts new file mode 100644 index 00000000000..6301c474eeb --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/components/index.ts @@ -0,0 +1,7 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export { QueryAssistBar } from './query_assist_bar'; +export { QueryAssistBanner } from './query_assist_banner'; diff --git a/src/plugins/query_enhancements/public/query_assist/components/index_selector.tsx b/src/plugins/query_enhancements/public/query_assist/components/index_selector.tsx new file mode 100644 index 00000000000..4e591e3401c --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/components/index_selector.tsx @@ -0,0 +1,47 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { EuiComboBox, EuiComboBoxOptionOption, EuiText } from '@elastic/eui'; +import React from 'react'; +import { useIndexPatterns, useIndices } from '../hooks/use_indices'; + +interface IndexSelectorProps { + dataSourceId?: string; + selectedIndex?: string; + setSelectedIndex: React.Dispatch>; +} + +// TODO this is a temporary solution, there will be a dataset selector from discover +export const IndexSelector: React.FC = (props) => { + const { data: indices, loading: indicesLoading } = useIndices(props.dataSourceId); + const { data: indexPatterns, loading: indexPatternsLoading } = useIndexPatterns(); + const loading = indicesLoading || indexPatternsLoading; + const indicesAndIndexPatterns = + indexPatterns && indices + ? [...indexPatterns, ...indices].filter( + (v1, index, array) => array.findIndex((v2) => v1 === v2) === index + ) + : []; + const options: EuiComboBoxOptionOption[] = indicesAndIndexPatterns.map((index) => ({ + label: index, + })); + const selectedOptions = props.selectedIndex ? [{ label: props.selectedIndex }] : undefined; + + return ( + Index} + singleSelection={{ asPlainText: true }} + isLoading={loading} + options={options} + selectedOptions={selectedOptions} + onChange={(index) => { + props.setSelectedIndex(index[0].label); + }} + /> + ); +}; diff --git a/src/plugins/query_enhancements/public/query_assist/components/query_assist_banner.test.tsx b/src/plugins/query_enhancements/public/query_assist/components/query_assist_banner.test.tsx new file mode 100644 index 00000000000..03655e0e266 --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/components/query_assist_banner.test.tsx @@ -0,0 +1,50 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { I18nProvider } from '@osd/i18n/react'; +import { fireEvent, render } from '@testing-library/react'; +import React, { ComponentProps } from 'react'; +import { QueryAssistBanner } from './query_assist_banner'; + +jest.mock('../../services', () => ({ + getStorage: () => ({ + get: jest.fn(), + set: jest.fn(), + }), +})); + +type QueryAssistBannerProps = ComponentProps; + +const renderQueryAssistBanner = (overrideProps: Partial = {}) => { + const props: QueryAssistBannerProps = Object.assign< + QueryAssistBannerProps, + Partial + >( + { + languages: ['test-lang1', 'test-lang2'], + }, + overrideProps + ); + const component = render( + + + + ); + return { component, props: props as jest.MockedObjectDeep }; +}; + +describe(' spec', () => { + it('should dismiss callout', async () => { + const { component } = renderQueryAssistBanner(); + expect( + component.getByText('Natural Language Query Generation for test-lang1, test-lang2') + ).toBeInTheDocument(); + + fireEvent.click(component.getByTestId('closeCallOutButton')); + expect( + component.queryByText('Natural Language Query Generation for test-lang1, test-lang2') + ).toBeNull(); + }); +}); diff --git a/src/plugins/query_enhancements/public/query_assist/components/query_assist_banner.tsx b/src/plugins/query_enhancements/public/query_assist/components/query_assist_banner.tsx new file mode 100644 index 00000000000..68faac461a6 --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/components/query_assist_banner.tsx @@ -0,0 +1,72 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + EuiBadge, + EuiCallOut, + EuiFlexGroup, + EuiFlexItem, + EuiIcon, + EuiLink, + EuiTextColor, +} from '@elastic/eui'; +import { FormattedMessage } from '@osd/i18n/react'; +import React, { useState } from 'react'; +import assistantMark from '../../assets/query_assist_mark.svg'; +import { getStorage } from '../../services'; + +const BANNER_STORAGE_KEY = 'queryAssist:banner:show'; + +interface QueryAssistBannerProps { + languages: string[]; +} + +export const QueryAssistBanner: React.FC = (props) => { + const storage = getStorage(); + const [showCallOut, _setShowCallOut] = useState(true); + const setShowCallOut: typeof _setShowCallOut = (show) => { + if (!show) { + storage.set(BANNER_STORAGE_KEY, false); + } + _setShowCallOut(show); + }; + + if (!showCallOut || storage.get(BANNER_STORAGE_KEY) === false) return null; + + return ( + + + + + + + + + + + + + + + + + + + } + dismissible + onDismiss={() => setShowCallOut(false)} + /> + ); +}; diff --git a/src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.tsx b/src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.tsx new file mode 100644 index 00000000000..84ac854c980 --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.tsx @@ -0,0 +1,121 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { EuiFlexGroup, EuiFlexItem, EuiForm, EuiFormRow } from '@elastic/eui'; +import React, { SyntheticEvent, useEffect, useMemo, useRef, useState } from 'react'; +import { + IDataPluginServices, + PersistedLog, + QueryEditorExtensionDependencies, +} from '../../../../data/public'; +import { useOpenSearchDashboards } from '../../../../opensearch_dashboards_react/public'; +import { QueryAssistParameters } from '../../../common/query_assist'; +import { ConnectionsService } from '../../data_source_connection'; +import { getStorage } from '../../services'; +import { useGenerateQuery } from '../hooks'; +import { getPersistedLog, ProhibitedQueryError } from '../utils'; +import { QueryAssistCallOut, QueryAssistCallOutType } from './call_outs'; +import { IndexSelector } from './index_selector'; +import { QueryAssistInput } from './query_assist_input'; +import { QueryAssistSubmitButton } from './submit_button'; + +interface QueryAssistInputProps { + dependencies: QueryEditorExtensionDependencies; + connectionsService: ConnectionsService; +} + +export const QueryAssistBar: React.FC = (props) => { + const { services } = useOpenSearchDashboards(); + const inputRef = useRef(null); + const storage = getStorage(); + const persistedLog: PersistedLog = useMemo( + () => getPersistedLog(services.uiSettings, storage, 'query-assist'), + [services.uiSettings, storage] + ); + const { generateQuery, loading } = useGenerateQuery(); + const [callOutType, setCallOutType] = useState(); + const dismissCallout = () => setCallOutType(undefined); + const [selectedIndex, setSelectedIndex] = useState(''); + const dataSourceIdRef = useRef(); + const previousQuestionRef = useRef(); + + useEffect(() => { + const subscription = props.connectionsService + .getSelectedConnection$() + .subscribe((connection) => { + dataSourceIdRef.current = connection?.id; + }); + return () => subscription.unsubscribe(); + }, [props.connectionsService]); + + const onSubmit = async (e: SyntheticEvent) => { + e.preventDefault(); + if (!inputRef.current?.value) { + setCallOutType('empty_query'); + return; + } + if (!selectedIndex) { + setCallOutType('empty_index'); + return; + } + dismissCallout(); + previousQuestionRef.current = inputRef.current.value; + persistedLog.add(inputRef.current.value); + const params: QueryAssistParameters = { + question: inputRef.current.value, + index: selectedIndex, + language: props.dependencies.language, + dataSourceId: dataSourceIdRef.current, + }; + const { response, error } = await generateQuery(params); + if (error) { + if (error instanceof ProhibitedQueryError) { + setCallOutType('invalid_query'); + } else { + services.notifications.toasts.addError(error, { title: 'Failed to generate results' }); + } + } else if (response) { + services.data.query.queryString.setQuery({ + query: response.query, + language: params.language, + }); + if (response.timeRange) services.data.query.timefilter.timefilter.setTime(response.timeRange); + setCallOutType('query_generated'); + } + }; + + return ( + + + + + + + + + + + + + + + + + ); +}; diff --git a/src/plugins/query_enhancements/public/query_assist/components/query_assist_input.test.tsx b/src/plugins/query_enhancements/public/query_assist/components/query_assist_input.test.tsx new file mode 100644 index 00000000000..1e6eb89da1f --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/components/query_assist_input.test.tsx @@ -0,0 +1,76 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { I18nProvider } from '@osd/i18n/react'; +import { fireEvent, render } from '@testing-library/react'; +import React, { ComponentProps } from 'react'; +import { SuggestionsComponentProps } from '../../../../data/public/ui/typeahead/suggestions_component'; +import { QueryAssistInput } from './query_assist_input'; + +jest.mock('../../services', () => ({ + getData: () => ({ + ui: { + SuggestionsComponent: ({ show, suggestions, onClick }: SuggestionsComponentProps) => ( +
+ {show && + suggestions.map((s, i) => ( + + ))} +
+ ), + }, + }), +})); + +const mockPersistedLog = { + get: () => ['mock suggestion 1', 'mock suggestion 2'], +} as any; + +type QueryAssistInputProps = ComponentProps; + +const renderQueryAssistInput = (overrideProps: Partial = {}) => { + const props: QueryAssistInputProps = Object.assign< + QueryAssistInputProps, + Partial + >( + { inputRef: { current: null }, persistedLog: mockPersistedLog, isDisabled: false }, + overrideProps + ); + const component = render( + + + + ); + return { component, props: props as jest.MockedObjectDeep }; +}; + +describe(' spec', () => { + it('should display input', () => { + const { component } = renderQueryAssistInput(); + const inputElement = component.getByTestId('query-assist-input-field-text') as HTMLInputElement; + expect(inputElement).toBeInTheDocument(); + fireEvent.change(inputElement, { target: { value: 'new value' } }); + expect(inputElement.value).toBe('new value'); + }); + + it('should display suggestions on input click', () => { + const { component } = renderQueryAssistInput(); + const inputElement = component.getByTestId('query-assist-input-field-text') as HTMLInputElement; + fireEvent.click(inputElement); + const suggestionsComponent = component.getByTestId('suggestions-component'); + expect(suggestionsComponent).toBeInTheDocument(); + }); + + it('should update input value on suggestion click', () => { + const { component } = renderQueryAssistInput(); + const inputElement = component.getByTestId('query-assist-input-field-text') as HTMLInputElement; + fireEvent.click(inputElement); + const suggestionButton = component.getByText('mock suggestion 1'); + fireEvent.click(suggestionButton); + expect(inputElement.value).toBe('mock suggestion 1'); + }); +}); diff --git a/src/plugins/query_enhancements/public/query_assist/components/query_assist_input.tsx b/src/plugins/query_enhancements/public/query_assist/components/query_assist_input.tsx new file mode 100644 index 00000000000..5ea5caaa5f4 --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/components/query_assist_input.tsx @@ -0,0 +1,114 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { EuiFieldText, EuiIcon, EuiOutsideClickDetector, EuiPortal } from '@elastic/eui'; +import React, { useMemo, useState } from 'react'; +import { PersistedLog, QuerySuggestionTypes } from '../../../../data/public'; +import assistantMark from '../../assets/query_assist_mark.svg'; +import { getData } from '../../services'; + +interface QueryAssistInputProps { + inputRef: React.RefObject; + persistedLog: PersistedLog; + isDisabled: boolean; + initialValue?: string; + selectedIndex?: string; + previousQuestion?: string; +} + +export const QueryAssistInput: React.FC = (props) => { + const { + ui: { SuggestionsComponent }, + } = getData(); + const [isSuggestionsVisible, setIsSuggestionsVisible] = useState(false); + const [suggestionIndex, setSuggestionIndex] = useState(null); + const [value, setValue] = useState(props.initialValue ?? ''); + + const sampleDataSuggestions = useMemo(() => { + switch (props.selectedIndex) { + case 'opensearch_dashboards_sample_data_ecommerce': + return [ + 'How many unique customers placed orders this week?', + 'Count the number of orders grouped by manufacturer and category', + 'find customers with first names like Eddie', + ]; + + case 'opensearch_dashboards_sample_data_logs': + return [ + 'Are there any errors in my logs?', + 'How many requests were there grouped by response code last week?', + "What's the average request size by week?", + ]; + + case 'opensearch_dashboards_sample_data_flights': + return [ + 'how many flights were there this week grouped by destination country?', + 'what were the longest flight delays this week?', + 'what carriers have the furthest flights?', + ]; + + default: + return []; + } + }, [props.selectedIndex]); + + const suggestions = useMemo(() => { + if (!props.persistedLog) return []; + return props.persistedLog + .get() + .concat(sampleDataSuggestions) + .filter( + (suggestion, i, array) => array.indexOf(suggestion) === i && suggestion.includes(value) + ) + .map((suggestion) => ({ + type: QuerySuggestionTypes.RecentSearch, + text: suggestion, + start: 0, + end: value.length, + })); + }, [props.persistedLog, value, sampleDataSuggestions]); + + return ( + setIsSuggestionsVisible(false)}> +
+ setIsSuggestionsVisible(true)} + onChange={(e) => setValue(e.target.value)} + onKeyDown={() => setIsSuggestionsVisible(true)} + placeholder={ + props.previousQuestion || + (props.selectedIndex + ? `Ask a natural language question about ${props.selectedIndex} to generate a query` + : 'Select an index to ask a question') + } + prepend={} + fullWidth + /> + + { + if (!props.inputRef.current) return; + setValue(suggestion.text); + setIsSuggestionsVisible(false); + setSuggestionIndex(null); + props.inputRef.current.focus(); + }} + onMouseEnter={(i) => setSuggestionIndex(i)} + loadMore={() => {}} + queryBarRect={props.inputRef.current?.getBoundingClientRect()} + size="s" + /> + +
+ + ); +}; diff --git a/src/plugins/query_enhancements/public/query_assist/components/submit_button.test.tsx b/src/plugins/query_enhancements/public/query_assist/components/submit_button.test.tsx new file mode 100644 index 00000000000..73d359e9739 --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/components/submit_button.test.tsx @@ -0,0 +1,34 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { fireEvent, render } from '@testing-library/react'; +import React, { ComponentProps } from 'react'; +import { QueryAssistSubmitButton } from './submit_button'; + +type SubmitButtonProps = ComponentProps; + +const renderSubmitButton = (overrideProps: Partial = {}) => { + const props: SubmitButtonProps = Object.assign>( + { + isDisabled: false, + }, + overrideProps + ); + const onSubmit = jest.fn((e) => e.preventDefault()); + const component = render( +
+ + + ); + return { component, onSubmit, props: props as jest.MockedObjectDeep }; +}; + +describe(' spec', () => { + it('should trigger submit form', () => { + const { component, onSubmit } = renderSubmitButton(); + fireEvent.click(component.getByTestId('query-assist-submit-button')); + expect(onSubmit).toBeCalled(); + }); +}); diff --git a/src/plugins/query_enhancements/public/query_assist/components/submit_button.tsx b/src/plugins/query_enhancements/public/query_assist/components/submit_button.tsx new file mode 100644 index 00000000000..52d2c5a6301 --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/components/submit_button.tsx @@ -0,0 +1,25 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { EuiButtonIcon } from '@elastic/eui'; +import React from 'react'; + +interface SubmitButtonProps { + isDisabled: boolean; +} + +export const QueryAssistSubmitButton: React.FC = (props) => { + return ( + + ); +}; diff --git a/src/plugins/query_enhancements/public/query_assist/hooks/index.ts b/src/plugins/query_enhancements/public/query_assist/hooks/index.ts new file mode 100644 index 00000000000..a2076151efb --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/hooks/index.ts @@ -0,0 +1,6 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export * from './use_generate'; diff --git a/src/plugins/query_enhancements/public/query_assist/hooks/use_generate.test.ts b/src/plugins/query_enhancements/public/query_assist/hooks/use_generate.test.ts new file mode 100644 index 00000000000..a15035b5b10 --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/hooks/use_generate.test.ts @@ -0,0 +1,92 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { act, renderHook } from '@testing-library/react-hooks/dom'; +import { coreMock } from '../../../../../core/public/mocks'; +import { useOpenSearchDashboards } from '../../../../opensearch_dashboards_react/public'; +import { useGenerateQuery } from './use_generate'; + +const coreSetup = coreMock.createSetup(); +const mockHttp = coreSetup.http; + +jest.mock('../../../../opensearch_dashboards_react/public', () => ({ + useOpenSearchDashboards: jest.fn(), + withOpenSearchDashboards: jest.fn((component: React.Component) => component), +})); + +describe('useGenerateQuery', () => { + beforeEach(() => { + (useOpenSearchDashboards as jest.MockedFunction) + // @ts-ignore for this test we only need http implemented + .mockImplementation(() => ({ + services: { + http: mockHttp, + }, + })); + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + it('should generate results', async () => { + mockHttp.post.mockResolvedValueOnce({ query: 'test query' }); + const { result } = renderHook(() => useGenerateQuery()); + const { generateQuery } = result.current; + + await act(async () => { + const response = await generateQuery({ + index: 'test', + language: 'test-lang', + question: 'test question', + }); + + expect(response).toEqual({ response: { query: 'test query' } }); + }); + }); + + it('should handle errors', async () => { + const { result } = renderHook(() => useGenerateQuery()); + const { generateQuery } = result.current; + const mockError = new Error('mockError'); + mockHttp.post.mockRejectedValueOnce(mockError); + + await act(async () => { + const response = await generateQuery({ + index: 'test', + language: 'test-lang', + question: 'test question', + }); + + expect(response).toEqual({ error: mockError }); + expect(result.current.loading).toBe(false); + }); + }); + + it('should abort previous call', async () => { + const { result } = renderHook(() => useGenerateQuery()); + const { generateQuery, abortControllerRef } = result.current; + + await act(async () => { + await generateQuery({ index: 'test', language: 'test-lang', question: 'test question' }); + const controller = abortControllerRef.current; + await generateQuery({ index: 'test', language: 'test-lang', question: 'test question' }); + + expect(controller?.signal.aborted).toBe(true); + }); + }); + + it('should abort call with controller', async () => { + const { result } = renderHook(() => useGenerateQuery()); + const { generateQuery, abortControllerRef } = result.current; + + await act(async () => { + await generateQuery({ index: 'test', language: 'test-lang', question: 'test question' }); + abortControllerRef.current?.abort(); + + expect(abortControllerRef.current?.signal.aborted).toBe(true); + }); + }); +}); diff --git a/src/plugins/query_enhancements/public/query_assist/hooks/use_generate.ts b/src/plugins/query_enhancements/public/query_assist/hooks/use_generate.ts new file mode 100644 index 00000000000..091e9dd7c29 --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/hooks/use_generate.ts @@ -0,0 +1,51 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { useEffect, useRef, useState } from 'react'; +import { IDataPluginServices } from '../../../../data/public'; +import { useOpenSearchDashboards } from '../../../../opensearch_dashboards_react/public'; +import { API } from '../../../common'; +import { QueryAssistParameters, QueryAssistResponse } from '../../../common/query_assist'; +import { formatError } from '../utils'; + +export const useGenerateQuery = () => { + const mounted = useRef(false); + const [loading, setLoading] = useState(false); + const abortControllerRef = useRef(); + const { services } = useOpenSearchDashboards(); + + useEffect(() => { + mounted.current = true; + return () => { + mounted.current = false; + if (abortControllerRef.current) { + abortControllerRef.current.abort(); + abortControllerRef.current = undefined; + } + }; + }, []); + + const generateQuery = async ( + params: QueryAssistParameters + ): Promise<{ response?: QueryAssistResponse; error?: Error }> => { + abortControllerRef.current?.abort(); + abortControllerRef.current = new AbortController(); + setLoading(true); + try { + const response = await services.http.post(API.QUERY_ASSIST.GENERATE, { + body: JSON.stringify(params), + signal: abortControllerRef.current?.signal, + }); + if (mounted.current) return { response }; + } catch (error) { + if (mounted.current) return { error: formatError(error) }; + } finally { + if (mounted.current) setLoading(false); + } + return {}; + }; + + return { generateQuery, loading, abortControllerRef }; +}; diff --git a/src/plugins/query_enhancements/public/query_assist/hooks/use_indices.ts b/src/plugins/query_enhancements/public/query_assist/hooks/use_indices.ts new file mode 100644 index 00000000000..1c985463877 --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/hooks/use_indices.ts @@ -0,0 +1,97 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { CatIndicesResponse } from '@opensearch-project/opensearch/api/types'; +import { Reducer, useEffect, useReducer, useState } from 'react'; +import { IDataPluginServices } from '../../../../data/public'; +import { useOpenSearchDashboards } from '../../../../opensearch_dashboards_react/public'; + +interface State { + data?: T; + loading: boolean; + error?: Error; +} + +type Action = + | { type: 'request' } + | { type: 'success'; payload: State['data'] } + | { type: 'failure'; error: NonNullable['error']> }; + +// TODO use instantiation expressions when typescript is upgraded to >= 4.7 +type GenericReducer = Reducer, Action>; +export const genericReducer: GenericReducer = (state, action) => { + switch (action.type) { + case 'request': + return { data: state.data, loading: true }; + case 'success': + return { loading: false, data: action.payload }; + case 'failure': + return { loading: false, error: action.error }; + default: + return state; + } +}; + +export const useIndices = (dataSourceId: string | undefined) => { + const reducer: GenericReducer = genericReducer; + const [state, dispatch] = useReducer(reducer, { loading: false }); + const [refresh, setRefresh] = useState({}); + const { services } = useOpenSearchDashboards(); + + useEffect(() => { + const abortController = new AbortController(); + dispatch({ type: 'request' }); + services.http + .post('/api/console/proxy', { + query: { path: '_cat/indices?format=json', method: 'GET', dataSourceId }, + signal: abortController.signal, + }) + .then((payload: CatIndicesResponse) => + dispatch({ + type: 'success', + payload: payload + .filter((meta) => meta.index && !meta.index.startsWith('.')) + .map((meta) => meta.index!), + }) + ) + .catch((error) => dispatch({ type: 'failure', error })); + + return () => abortController.abort(); + }, [refresh, services.http, dataSourceId]); + + return { ...state, refresh: () => setRefresh({}) }; +}; + +export const useIndexPatterns = () => { + const reducer: GenericReducer = genericReducer; + const [state, dispatch] = useReducer(reducer, { loading: false }); + const [refresh, setRefresh] = useState({}); + const { services } = useOpenSearchDashboards(); + + useEffect(() => { + let abort = false; + dispatch({ type: 'request' }); + + services.data.indexPatterns + .getTitles() + .then((payload) => { + if (!abort) + dispatch({ + type: 'success', + // temporary solution does not support index patterns from other data sources + payload: payload.filter((title) => !title.includes('::')), + }); + }) + .catch((error) => { + if (!abort) dispatch({ type: 'failure', error }); + }); + + return () => { + abort = true; + }; + }, [refresh, services.data.indexPatterns]); + + return { ...state, refresh: () => setRefresh({}) }; +}; diff --git a/src/plugins/query_enhancements/public/query_assist/index.ts b/src/plugins/query_enhancements/public/query_assist/index.ts new file mode 100644 index 00000000000..51d08717b22 --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/index.ts @@ -0,0 +1,6 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export { createQueryAssistExtension } from './utils'; diff --git a/src/plugins/query_enhancements/public/query_assist/utils/create_extension.test.tsx b/src/plugins/query_enhancements/public/query_assist/utils/create_extension.test.tsx new file mode 100644 index 00000000000..ea568959152 --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/utils/create_extension.test.tsx @@ -0,0 +1,104 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { firstValueFrom } from '@osd/std'; +import { act, render, screen } from '@testing-library/react'; +import React from 'react'; +import { coreMock } from '../../../../../core/public/mocks'; +import { IIndexPattern } from '../../../../data/public'; +import { ConfigSchema } from '../../../common/config'; +import { ConnectionsService } from '../../data_source_connection'; +import { Connection } from '../../types'; +import { createQueryAssistExtension } from './create_extension'; + +const coreSetupMock = coreMock.createSetup({ + pluginStartDeps: { + data: { + ui: {}, + }, + }, +}); +const httpMock = coreSetupMock.http; + +jest.mock('../components', () => ({ + QueryAssistBar: jest.fn(() =>
QueryAssistBar
), +})); + +jest.mock('../components/query_assist_banner', () => ({ + QueryAssistBanner: jest.fn(() =>
QueryAssistBanner
), +})); + +describe.skip('CreateExtension', () => { + afterEach(() => { + jest.clearAllMocks(); + }); + + const config: ConfigSchema['queryAssist'] = { + supportedLanguages: [{ language: 'PPL', agentConfig: 'os_query_assist_ppl' }], + }; + const connectionsService = new ConnectionsService({ + startServices: coreSetupMock.getStartServices(), + http: httpMock, + }); + + // for these tests we only need id field in the connection + connectionsService.setSelectedConnection$({ + dataSource: { id: 'mock-data-source-id' }, + } as Connection); + + it('should be enabled if at least one language is configured', async () => { + httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] }); + const extension = createQueryAssistExtension(httpMock, connectionsService, config); + const isEnabled = await firstValueFrom(extension.isEnabled$({ language: 'PPL' })); + expect(isEnabled).toBeTruthy(); + expect(httpMock.get).toBeCalledWith('/api/enhancements/assist/languages', { + query: { dataSourceId: 'mock-data-source-id' }, + }); + }); + + it('should be disabled for unsupported language', async () => { + httpMock.get.mockRejectedValueOnce(new Error('network failure')); + const extension = createQueryAssistExtension(httpMock, connectionsService, config); + const isEnabled = await firstValueFrom(extension.isEnabled$({ language: 'PPL' })); + expect(isEnabled).toBeFalsy(); + expect(httpMock.get).toBeCalledWith('/api/enhancements/assist/languages', { + query: { dataSourceId: 'mock-data-source-id' }, + }); + }); + + it('should render the component if language is supported', async () => { + httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] }); + const extension = createQueryAssistExtension(httpMock, connectionsService, config); + const component = extension.getComponent?.({ + language: 'PPL', + indexPatterns: [{ id: 'test-pattern' }] as IIndexPattern[], + }); + + if (!component) throw new Error('QueryEditorExtensions Component is undefined'); + + await act(async () => { + render(component); + }); + + expect(screen.getByText('QueryAssistBar')).toBeInTheDocument(); + }); + + it('should render the banner if language is not supported', async () => { + httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] }); + const extension = createQueryAssistExtension(httpMock, connectionsService, config); + const banner = extension.getBanner?.({ + language: 'DQL', + indexPatterns: [{ id: 'test-pattern' }] as IIndexPattern[], + }); + + if (!banner) throw new Error('QueryEditorExtensions Banner is undefined'); + + await act(async () => { + render(banner); + }); + + expect(screen.getByText('QueryAssistBanner')).toBeInTheDocument(); + }); +}); diff --git a/src/plugins/query_enhancements/public/query_assist/utils/create_extension.tsx b/src/plugins/query_enhancements/public/query_assist/utils/create_extension.tsx new file mode 100644 index 00000000000..bd990b57f5a --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/utils/create_extension.tsx @@ -0,0 +1,126 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { HttpSetup } from 'opensearch-dashboards/public'; +import React, { useEffect, useState } from 'react'; +import { of } from 'rxjs'; +import { distinctUntilChanged, switchMap, map } from 'rxjs/operators'; +import { + QueryEditorExtensionConfig, + QueryEditorExtensionDependencies, +} from '../../../../data/public'; +import { API } from '../../../common'; +import { ConfigSchema } from '../../../common/config'; +import { ConnectionsService } from '../../data_source_connection'; +import { QueryAssistBar, QueryAssistBanner } from '../components'; + +/** + * @returns observable list of query assist agent configured languages in the + * selected data source. + */ +const getAvailableLanguages$ = ( + availableLanguagesByDataSource: Map, + connectionsService: ConnectionsService, + http: HttpSetup +) => + connectionsService.getSelectedConnection$().pipe( + distinctUntilChanged(), + switchMap(async (connection) => { + const dataSourceId = connection?.id; + const cached = availableLanguagesByDataSource.get(dataSourceId); + if (cached !== undefined) return cached; + const languages = await http + .get<{ configuredLanguages: string[] }>(API.QUERY_ASSIST.LANGUAGES, { + query: { dataSourceId }, + }) + .then((response) => response.configuredLanguages) + .catch(() => []); + availableLanguagesByDataSource.set(dataSourceId, languages); + return languages; + }) + ); + +export const createQueryAssistExtension = ( + http: HttpSetup, + connectionsService: ConnectionsService, + config: ConfigSchema['queryAssist'] +): QueryEditorExtensionConfig => { + const availableLanguagesByDataSource: Map = new Map(); + + return { + id: 'query-assist', + order: 1000, + isEnabled$: (dependencies) => { + // currently query assist tool relies on opensearch API to get index + // mappings, non-default data source types are not supported + if (dependencies.dataSource && dependencies.dataSource?.getType() !== 'default') + return of(false); + + return getAvailableLanguages$(availableLanguagesByDataSource, connectionsService, http).pipe( + map((languages) => languages.length > 0) + ); + }, + getComponent: (dependencies) => { + // only show the component if user is on a supported language. + return ( + + + + ); + }, + getBanner: (dependencies) => { + // advertise query assist if user is not on a supported language. + return ( + + conf.language)} /> + + ); + }, + }; +}; + +interface QueryAssistWrapperProps { + availableLanguagesByDataSource: Map; + dependencies: QueryEditorExtensionDependencies; + connectionsService: ConnectionsService; + http: HttpSetup; + invert?: boolean; +} + +const QueryAssistWrapper: React.FC = (props) => { + const [visible, setVisible] = useState(false); + + useEffect(() => { + let mounted = true; + + const subscription = getAvailableLanguages$( + props.availableLanguagesByDataSource, + props.connectionsService, + props.http + ).subscribe((languages) => { + const available = languages.includes(props.dependencies.language); + if (mounted) setVisible(props.invert ? !available : available); + }); + + return () => { + mounted = false; + subscription.unsubscribe(); + }; + }, [props]); + + if (!visible) return null; + return <>{props.children}; +}; diff --git a/src/plugins/query_enhancements/public/query_assist/utils/errors.test.ts b/src/plugins/query_enhancements/public/query_assist/utils/errors.test.ts new file mode 100644 index 00000000000..1adbae11b46 --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/utils/errors.test.ts @@ -0,0 +1,65 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { ResponseError } from '@opensearch-project/opensearch/lib/errors'; +import { ERROR_DETAILS } from '../../../common'; +import { formatError, ProhibitedQueryError } from './errors'; + +describe('formatError', () => { + it('should return an error with a custom message for status code 429', () => { + const error = new ResponseError({ + statusCode: 429, + body: { + statusCode: 429, + message: 'Too many requests', + }, + warnings: [], + headers: null, + meta: {} as any, + }); + + const formattedError = formatError(error); + expect(formattedError.message).toEqual( + 'Request is throttled. Try again later or contact your administrator' + ); + }); + + it('should return a ProhibitedQueryError for guardrails triggered', () => { + const error = new ResponseError({ + statusCode: 400, + body: { + statusCode: 400, + message: ERROR_DETAILS.GUARDRAILS_TRIGGERED, + }, + warnings: [], + headers: null, + meta: {} as any, + }); + const formattedError = formatError(error); + expect(formattedError).toBeInstanceOf(ProhibitedQueryError); + expect(formattedError.message).toEqual(error.body.message); + }); + + it('should return the original error body for other errors', () => { + const error = new ResponseError({ + statusCode: 500, + body: { + statusCode: 500, + message: 'Internal server error', + }, + warnings: [], + headers: null, + meta: {} as any, + }); + const formattedError = formatError(error); + expect(formattedError).toEqual(error.body); + }); + + it('should return the original error if no body property', () => { + const error = new Error('Some error'); + const formattedError = formatError(error); + expect(formattedError).toEqual(error); + }); +}); diff --git a/src/plugins/query_enhancements/public/query_assist/utils/errors.ts b/src/plugins/query_enhancements/public/query_assist/utils/errors.ts new file mode 100644 index 00000000000..dbd0a9ef8fc --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/utils/errors.ts @@ -0,0 +1,30 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { ResponseError } from '@opensearch-project/opensearch/lib/errors'; +import { ERROR_DETAILS } from '../../../common'; + +export class ProhibitedQueryError extends Error { + constructor(message?: string) { + super(message); + } +} + +export const formatError = (error: ResponseError | Error): Error => { + if ('body' in error) { + if (error.body.statusCode === 429) + return { + ...error.body, + message: 'Request is throttled. Try again later or contact your administrator', + } as Error; + if ( + error.body.statusCode === 400 && + error.body.message.includes(ERROR_DETAILS.GUARDRAILS_TRIGGERED) + ) + return new ProhibitedQueryError(error.body.message); + return error.body as Error; + } + return error; +}; diff --git a/src/plugins/query_enhancements/public/query_assist/utils/get_persisted_log.ts b/src/plugins/query_enhancements/public/query_assist/utils/get_persisted_log.ts new file mode 100644 index 00000000000..c86edcf125c --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/utils/get_persisted_log.ts @@ -0,0 +1,24 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { IUiSettingsClient } from 'opensearch-dashboards/public'; +import { UI_SETTINGS } from '../../../../data/common'; +import { PersistedLog } from '../../../../data/public'; +import { IStorageWrapper } from '../../../../opensearch_dashboards_utils/public'; + +export function getPersistedLog( + uiSettings: IUiSettingsClient, + storage: IStorageWrapper, + language: string +) { + return new PersistedLog( + `typeahead:${language}`, + { + maxLength: uiSettings.get(UI_SETTINGS.HISTORY_LIMIT), + filterDuplicates: true, + }, + storage + ); +} diff --git a/src/plugins/query_enhancements/public/query_assist/utils/index.ts b/src/plugins/query_enhancements/public/query_assist/utils/index.ts new file mode 100644 index 00000000000..517be64e9e9 --- /dev/null +++ b/src/plugins/query_enhancements/public/query_assist/utils/index.ts @@ -0,0 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export * from './create_extension'; +export * from './errors'; +export * from './get_persisted_log'; diff --git a/src/plugins/query_enhancements/public/search/index.ts b/src/plugins/query_enhancements/public/search/index.ts new file mode 100644 index 00000000000..9835c1345f0 --- /dev/null +++ b/src/plugins/query_enhancements/public/search/index.ts @@ -0,0 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export { PPLSearchInterceptor } from './ppl_search_interceptor'; +export { SQLSearchInterceptor } from './sql_search_interceptor'; +export { SQLAsyncSearchInterceptor } from './sql_async_search_interceptor'; diff --git a/src/plugins/query_enhancements/public/search/ppl_search_interceptor.ts b/src/plugins/query_enhancements/public/search/ppl_search_interceptor.ts new file mode 100644 index 00000000000..aac50de3bb9 --- /dev/null +++ b/src/plugins/query_enhancements/public/search/ppl_search_interceptor.ts @@ -0,0 +1,216 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { trimEnd } from 'lodash'; +import { Observable, throwError } from 'rxjs'; +import { concatMap } from 'rxjs/operators'; +import { + DataFrameAggConfig, + getAggConfig, + getRawDataFrame, + getRawQueryString, + getTimeField, + formatTimePickerDate, + getUniqueValuesForRawAggs, + updateDataFrameMeta, +} from '../../../data/common'; +import { + DataPublicPluginStart, + IOpenSearchDashboardsSearchRequest, + IOpenSearchDashboardsSearchResponse, + ISearchOptions, + SearchInterceptor, + SearchInterceptorDeps, +} from '../../../data/public'; +import { + formatDate, + SEARCH_STRATEGY, + removeKeyword, + API, + FetchDataFrameContext, + fetchDataFrame, +} from '../../common'; +import { QueryEnhancementsPluginStartDependencies } from '../types'; +import { ConnectionsService } from '../data_source_connection'; + +export class PPLSearchInterceptor extends SearchInterceptor { + protected queryService!: DataPublicPluginStart['query']; + protected aggsService!: DataPublicPluginStart['search']['aggs']; + + constructor( + deps: SearchInterceptorDeps, + private readonly connectionsService: ConnectionsService + ) { + super(deps); + + deps.startServices.then(([coreStart, depsStart]) => { + this.queryService = (depsStart as QueryEnhancementsPluginStartDependencies).data.query; + this.aggsService = (depsStart as QueryEnhancementsPluginStartDependencies).data.search.aggs; + }); + } + + protected runSearch( + request: IOpenSearchDashboardsSearchRequest, + signal?: AbortSignal, + strategy?: string + ): Observable { + const { id, ...searchRequest } = request; + const dfContext: FetchDataFrameContext = { + http: this.deps.http, + path: trimEnd(API.PPL_SEARCH), + signal, + }; + const { timefilter } = this.queryService; + const dateRange = timefilter.timefilter.getTime(); + const { fromDate, toDate } = formatTimePickerDate(dateRange, 'YYYY-MM-DD HH:mm:ss.SSS'); + + const getTimeFilter = (timeField: any) => { + return ` | where ${timeField?.name} >= '${formatDate(fromDate)}' and ${ + timeField?.name + } <= '${formatDate(toDate)}'`; + }; + + const insertTimeFilter = (query: string, filter: string) => { + const pipes = query.split('|'); + return pipes + .slice(0, 1) + .concat(filter.substring(filter.indexOf('where')), pipes.slice(1)) + .join(' | '); + }; + + const getAggQsFn = ({ + qs, + aggConfig, + timeField, + timeFilter, + }: { + qs: string; + aggConfig: DataFrameAggConfig; + timeField: any; + timeFilter: string; + }) => { + return removeKeyword(`${qs} ${getAggString(timeField, aggConfig)} ${timeFilter}`); + }; + + const getAggString = (timeField: any, aggsConfig?: DataFrameAggConfig) => { + if (!aggsConfig) { + return ` | stats count() by span(${ + timeField?.name + }, ${this.aggsService.calculateAutoTimeExpression({ + from: fromDate, + to: toDate, + mode: 'absolute', + })})`; + } + if (aggsConfig.date_histogram) { + return ` | stats count() by span(${timeField?.name}, ${ + aggsConfig.date_histogram.fixed_interval ?? + aggsConfig.date_histogram.calendar_interval ?? + this.aggsService.calculateAutoTimeExpression({ + from: fromDate, + to: toDate, + mode: 'absolute', + }) + })`; + } + if (aggsConfig.avg) { + return ` | stats avg(${aggsConfig.avg.field})`; + } + if (aggsConfig.cardinality) { + return ` | dedup ${aggsConfig.cardinality.field} | stats count()`; + } + if (aggsConfig.terms) { + return ` | stats count() by ${aggsConfig.terms.field}`; + } + if (aggsConfig.id === 'other-filter') { + const uniqueConfig = getUniqueValuesForRawAggs(aggsConfig); + if ( + !uniqueConfig || + !uniqueConfig.field || + !uniqueConfig.values || + uniqueConfig.values.length === 0 + ) { + return ''; + } + + let otherQueryString = ` | stats count() by ${uniqueConfig.field}`; + uniqueConfig.values.forEach((value, index) => { + otherQueryString += ` ${index === 0 ? '| where' : 'and'} ${ + uniqueConfig.field + }<>'${value}'`; + }); + return otherQueryString; + } + }; + + const dataFrame = getRawDataFrame(searchRequest); + if (!dataFrame) { + return throwError(this.handleSearchError('DataFrame is not defined', request, signal!)); + } + + let queryString = dataFrame.meta?.queryConfig?.qs ?? getRawQueryString(searchRequest) ?? ''; + + dataFrame.meta = { + ...dataFrame.meta, + queryConfig: { + ...dataFrame.meta.queryConfig, + ...(this.connectionsService.getSelectedConnection() && { + dataSourceId: this.connectionsService.getSelectedConnection()?.id, + }), + }, + }; + const aggConfig = getAggConfig( + searchRequest, + {}, + this.aggsService.types.get.bind(this) + ) as DataFrameAggConfig; + + if (!dataFrame.schema) { + return fetchDataFrame(dfContext, queryString, dataFrame).pipe( + concatMap((response) => { + const df = response.body; + const timeField = getTimeField(df, aggConfig); + if (timeField) { + const timeFilter = getTimeFilter(timeField); + const newQuery = insertTimeFilter(queryString, timeFilter); + updateDataFrameMeta({ + dataFrame: df, + qs: newQuery, + aggConfig, + timeField, + timeFilter, + getAggQsFn: getAggQsFn.bind(this), + }); + return fetchDataFrame(dfContext, newQuery, df); + } + return fetchDataFrame(dfContext, queryString, df); + }) + ); + } + + if (dataFrame.schema) { + const timeField = getTimeField(dataFrame, aggConfig); + if (timeField) { + const timeFilter = getTimeFilter(timeField); + const newQuery = insertTimeFilter(queryString, timeFilter); + updateDataFrameMeta({ + dataFrame, + qs: newQuery, + aggConfig, + timeField, + timeFilter, + getAggQsFn: getAggQsFn.bind(this), + }); + queryString += timeFilter; + } + } + + return fetchDataFrame(dfContext, queryString, dataFrame); + } + + public search(request: IOpenSearchDashboardsSearchRequest, options: ISearchOptions) { + return this.runSearch(request, options.abortSignal, SEARCH_STRATEGY.PPL); + } +} diff --git a/src/plugins/query_enhancements/public/search/sql_async_search_interceptor.ts b/src/plugins/query_enhancements/public/search/sql_async_search_interceptor.ts new file mode 100644 index 00000000000..9232ef146cd --- /dev/null +++ b/src/plugins/query_enhancements/public/search/sql_async_search_interceptor.ts @@ -0,0 +1,137 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { trimEnd } from 'lodash'; +import { BehaviorSubject, Observable, throwError } from 'rxjs'; +import { i18n } from '@osd/i18n'; +import { concatMap, map } from 'rxjs/operators'; +import { + DATA_FRAME_TYPES, + DataPublicPluginStart, + IOpenSearchDashboardsSearchRequest, + IOpenSearchDashboardsSearchResponse, + ISearchOptions, + SearchInterceptor, + SearchInterceptorDeps, +} from '../../../data/public'; +import { getRawDataFrame, getRawQueryString, IDataFrameResponse } from '../../../data/common'; +import { + API, + DataFramePolling, + FetchDataFrameContext, + SEARCH_STRATEGY, + fetchDataFrame, + fetchDataFramePolling, +} from '../../common'; +import { QueryEnhancementsPluginStartDependencies } from '../types'; +import { ConnectionsService } from '../data_source_connection'; + +export class SQLAsyncSearchInterceptor extends SearchInterceptor { + protected queryService!: DataPublicPluginStart['query']; + protected aggsService!: DataPublicPluginStart['search']['aggs']; + protected indexPatterns!: DataPublicPluginStart['indexPatterns']; + protected dataFrame$ = new BehaviorSubject(undefined); + + constructor( + deps: SearchInterceptorDeps, + private readonly connectionsService: ConnectionsService + ) { + super(deps); + + deps.startServices.then(([coreStart, depsStart]) => { + this.queryService = (depsStart as QueryEnhancementsPluginStartDependencies).data.query; + this.aggsService = (depsStart as QueryEnhancementsPluginStartDependencies).data.search.aggs; + }); + } + + protected runSearch( + request: IOpenSearchDashboardsSearchRequest, + signal?: AbortSignal, + strategy?: string + ): Observable { + const { id, ...searchRequest } = request; + const path = trimEnd(API.SQL_ASYNC_SEARCH); + const dfContext: FetchDataFrameContext = { + http: this.deps.http, + path, + signal, + }; + + const dataFrame = getRawDataFrame(searchRequest); + if (!dataFrame) { + return throwError(this.handleSearchError('DataFrame is not defined', request, signal!)); + } + + const queryString = + dataFrame.meta?.queryConfig?.formattedQs() ?? getRawQueryString(searchRequest) ?? ''; + + dataFrame.meta = { + ...dataFrame.meta, + queryConfig: { + ...dataFrame.meta.queryConfig, + ...(this.connectionsService.getSelectedConnection() && + this.connectionsService.getSelectedConnection()?.dataSource && { + dataSourceId: this.connectionsService.getSelectedConnection()?.dataSource.id, + }), + }, + }; + + const onPollingSuccess = (pollingResult: any) => { + if (pollingResult && pollingResult.body.meta.status === 'SUCCESS') { + return false; + } + if (pollingResult && pollingResult.body.meta.status === 'FAILED') { + const jsError = new Error(pollingResult.data.error.response); + this.deps.toasts.addError(jsError, { + title: i18n.translate('queryEnhancements.sqlQueryError', { + defaultMessage: 'Could not complete the SQL async query', + }), + toastMessage: pollingResult.data.error.response, + }); + return false; + } + + this.deps.toasts.addInfo({ + title: i18n.translate('queryEnhancements.sqlQueryPolling', { + defaultMessage: 'Polling query job results...', + }), + }); + + return true; + }; + + const onPollingError = (error: Error) => { + throw new Error(error.message); + }; + + this.deps.toasts.addInfo({ + title: i18n.translate('queryEnhancements.sqlQueryInfo', { + defaultMessage: 'Starting query job...', + }), + }); + return fetchDataFrame(dfContext, queryString, dataFrame).pipe( + concatMap((jobResponse) => { + const df = jobResponse.body; + const dataFramePolling = new DataFramePolling( + () => fetchDataFramePolling(dfContext, df), + 5000, + onPollingSuccess, + onPollingError + ); + return dataFramePolling.fetch().pipe( + map(() => { + const dfPolling = dataFramePolling.data; + dfPolling.type = DATA_FRAME_TYPES.DEFAULT; + return dfPolling; + }) + ); + }) + ); + } + + public search(request: IOpenSearchDashboardsSearchRequest, options: ISearchOptions) { + return this.runSearch(request, options.abortSignal, SEARCH_STRATEGY.SQL_ASYNC); + } +} diff --git a/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts b/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts new file mode 100644 index 00000000000..c4dd7409faf --- /dev/null +++ b/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts @@ -0,0 +1,92 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { trimEnd } from 'lodash'; +import { Observable, throwError } from 'rxjs'; +import { i18n } from '@osd/i18n'; +import { concatMap } from 'rxjs/operators'; +import { getRawDataFrame, getRawQueryString } from '../../../data/common'; +import { + DataPublicPluginStart, + IOpenSearchDashboardsSearchRequest, + IOpenSearchDashboardsSearchResponse, + ISearchOptions, + SearchInterceptor, + SearchInterceptorDeps, +} from '../../../data/public'; +import { API, FetchDataFrameContext, SEARCH_STRATEGY, fetchDataFrame } from '../../common'; +import { QueryEnhancementsPluginStartDependencies } from '../types'; +import { ConnectionsService } from '../data_source_connection'; + +export class SQLSearchInterceptor extends SearchInterceptor { + protected queryService!: DataPublicPluginStart['query']; + protected aggsService!: DataPublicPluginStart['search']['aggs']; + + constructor( + deps: SearchInterceptorDeps, + private readonly connectionsService: ConnectionsService + ) { + super(deps); + + deps.startServices.then(([coreStart, depsStart]) => { + this.queryService = (depsStart as QueryEnhancementsPluginStartDependencies).data.query; + this.aggsService = (depsStart as QueryEnhancementsPluginStartDependencies).data.search.aggs; + }); + } + + protected runSearch( + request: IOpenSearchDashboardsSearchRequest, + signal?: AbortSignal, + strategy?: string + ): Observable { + const { id, ...searchRequest } = request; + const dfContext: FetchDataFrameContext = { + http: this.deps.http, + path: trimEnd(API.SQL_SEARCH), + signal, + }; + + const dataFrame = getRawDataFrame(searchRequest); + if (!dataFrame) { + return throwError(this.handleSearchError('DataFrame is not defined', request, signal!)); + } + + const queryString = dataFrame.meta?.queryConfig?.qs ?? getRawQueryString(searchRequest) ?? ''; + + dataFrame.meta = { + ...dataFrame.meta, + queryConfig: { + ...dataFrame.meta.queryConfig, + ...(this.connectionsService.getSelectedConnection() && { + dataSourceId: this.connectionsService.getSelectedConnection()?.id, + }), + }, + }; + + if (!dataFrame.schema) { + return fetchDataFrame(dfContext, queryString, dataFrame).pipe( + concatMap((response) => { + const df = response.body; + if (df.error) { + const jsError = new Error(df.error.response); + this.deps.toasts.addError(jsError, { + title: i18n.translate('queryEnhancements.sqlQueryError', { + defaultMessage: 'Could not complete the SQL query', + }), + toastMessage: df.error.msg, + }); + } + return fetchDataFrame(dfContext, queryString, df); + }) + ); + } + + return fetchDataFrame(dfContext, queryString, dataFrame); + } + + public search(request: IOpenSearchDashboardsSearchRequest, options: ISearchOptions) { + return this.runSearch(request, options.abortSignal, SEARCH_STRATEGY.SQL); + } +} diff --git a/src/plugins/query_enhancements/public/services.ts b/src/plugins/query_enhancements/public/services.ts new file mode 100644 index 00000000000..d11233be2dc --- /dev/null +++ b/src/plugins/query_enhancements/public/services.ts @@ -0,0 +1,11 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { createGetterSetter } from '../../opensearch_dashboards_utils/common'; +import { IStorageWrapper } from '../../opensearch_dashboards_utils/public'; +import { DataPublicPluginStart } from '../../data/public'; + +export const [getStorage, setStorage] = createGetterSetter('storage'); +export const [getData, setData] = createGetterSetter('data'); diff --git a/src/plugins/query_enhancements/public/types.ts b/src/plugins/query_enhancements/public/types.ts new file mode 100644 index 00000000000..c31da11e6b1 --- /dev/null +++ b/src/plugins/query_enhancements/public/types.ts @@ -0,0 +1,38 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { CoreSetup, CoreStart } from 'opensearch-dashboards/public'; +import { DataSourcePluginStart } from 'src/plugins/data_source/public'; +import { DataPublicPluginSetup, DataPublicPluginStart } from '../../data/public'; + +// eslint-disable-next-line @typescript-eslint/no-empty-interface +export interface QueryEnhancementsPluginSetup {} + +// eslint-disable-next-line @typescript-eslint/no-empty-interface +export interface QueryEnhancementsPluginStart {} + +export interface QueryEnhancementsPluginSetupDependencies { + data: DataPublicPluginSetup; +} + +export interface QueryEnhancementsPluginStartDependencies { + data: DataPublicPluginStart; + dataSource?: DataSourcePluginStart; +} + +export interface Connection { + dataSource: { + id: string; + title: string; + endpoint?: string; + installedPlugins?: string[]; + auth?: any; + }; +} + +export interface ConnectionsServiceDeps { + http: CoreSetup['http']; + startServices: Promise<[CoreStart, QueryEnhancementsPluginStartDependencies, unknown]>; +} diff --git a/src/plugins/query_enhancements/server/index.ts b/src/plugins/query_enhancements/server/index.ts new file mode 100644 index 00000000000..4d72ff3eb27 --- /dev/null +++ b/src/plugins/query_enhancements/server/index.ts @@ -0,0 +1,29 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { PluginConfigDescriptor, PluginInitializerContext } from '../../../core/server'; +import { QueryEnhancementsPlugin } from './plugin'; +import { configSchema, ConfigSchema } from '../common/config'; + +export const config: PluginConfigDescriptor = { + exposeToBrowser: { + queryAssist: true, + }, + schema: configSchema, +}; + +export function plugin(initializerContext: PluginInitializerContext) { + return new QueryEnhancementsPlugin(initializerContext); +} + +export { + Facet, + FacetProps, + OpenSearchPPLPlugin, + OpenSearchObservabilityPlugin, + shimStats, + shimSchemaRow, +} from './utils'; +export { QueryEnhancementsPluginSetup, QueryEnhancementsPluginStart } from './types'; diff --git a/src/plugins/query_enhancements/server/plugin.ts b/src/plugins/query_enhancements/server/plugin.ts new file mode 100644 index 00000000000..db105849694 --- /dev/null +++ b/src/plugins/query_enhancements/server/plugin.ts @@ -0,0 +1,99 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { Observable } from 'rxjs'; +import { first } from 'rxjs/operators'; +import { + CoreSetup, + CoreStart, + Logger, + Plugin, + PluginInitializerContext, + SharedGlobalConfig, +} from '../../../core/server'; +import { SEARCH_STRATEGY } from '../common'; +import { ConfigSchema } from '../common/config'; +import { defineRoutes } from './routes'; +import { + pplSearchStrategyProvider, + sqlSearchStrategyProvider, + sqlAsyncSearchStrategyProvider, +} from './search'; +import { + QueryEnhancementsPluginSetup, + QueryEnhancementsPluginSetupDependencies, + QueryEnhancementsPluginStart, +} from './types'; +import { OpenSearchObservabilityPlugin, OpenSearchPPLPlugin } from './utils'; + +export class QueryEnhancementsPlugin + implements Plugin { + private readonly logger: Logger; + private readonly config$: Observable; + constructor(private initializerContext: PluginInitializerContext) { + this.logger = initializerContext.logger.get(); + this.config$ = initializerContext.config.legacy.globalConfig$; + } + + public setup(core: CoreSetup, { data, dataSource }: QueryEnhancementsPluginSetupDependencies) { + this.logger.debug('queryEnhancements: Setup'); + const router = core.http.createRouter(); + // Register server side APIs + const client = core.opensearch.legacy.createClient('opensearch_observability', { + plugins: [OpenSearchPPLPlugin, OpenSearchObservabilityPlugin], + }); + + if (dataSource) { + dataSource.registerCustomApiSchema(OpenSearchPPLPlugin); + dataSource.registerCustomApiSchema(OpenSearchObservabilityPlugin); + } + + const pplSearchStrategy = pplSearchStrategyProvider(this.config$, this.logger, client); + const sqlSearchStrategy = sqlSearchStrategyProvider(this.config$, this.logger, client); + const sqlAsyncSearchStrategy = sqlAsyncSearchStrategyProvider( + this.config$, + this.logger, + client + ); + + data.search.registerSearchStrategy(SEARCH_STRATEGY.PPL, pplSearchStrategy); + data.search.registerSearchStrategy(SEARCH_STRATEGY.SQL, sqlSearchStrategy); + data.search.registerSearchStrategy(SEARCH_STRATEGY.SQL_ASYNC, sqlAsyncSearchStrategy); + + core.http.registerRouteHandlerContext('query_assist', () => ({ + logger: this.logger, + configPromise: this.initializerContext.config + .create() + .pipe(first()) + .toPromise(), + dataSourceEnabled: !!dataSource, + })); + + core.http.registerRouteHandlerContext('data_source_connection', () => ({ + logger: this.logger, + configPromise: this.initializerContext.config + .create() + .pipe(first()) + .toPromise(), + dataSourceEnabled: !!dataSource, + })); + + defineRoutes(this.logger, router, { + ppl: pplSearchStrategy, + sql: sqlSearchStrategy, + sqlasync: sqlAsyncSearchStrategy, + }); + + this.logger.info('queryEnhancements: Setup complete'); + return {}; + } + + public start(core: CoreStart) { + this.logger.debug('queryEnhancements: Started'); + return {}; + } + + public stop() {} +} diff --git a/src/plugins/query_enhancements/server/routes/data_source_connection/index.ts b/src/plugins/query_enhancements/server/routes/data_source_connection/index.ts new file mode 100644 index 00000000000..0e6c700b4d4 --- /dev/null +++ b/src/plugins/query_enhancements/server/routes/data_source_connection/index.ts @@ -0,0 +1,6 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export { registerDataSourceConnectionsRoutes } from './routes'; diff --git a/src/plugins/query_enhancements/server/routes/data_source_connection/routes.ts b/src/plugins/query_enhancements/server/routes/data_source_connection/routes.ts new file mode 100644 index 00000000000..f4fe42779da --- /dev/null +++ b/src/plugins/query_enhancements/server/routes/data_source_connection/routes.ts @@ -0,0 +1,48 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { schema } from '@osd/config-schema'; +import { IRouter } from 'opensearch-dashboards/server'; +import { DataSourceAttributes } from '../../../../data_source/common/data_sources'; +import { API } from '../../../common'; + +export function registerDataSourceConnectionsRoutes(router: IRouter) { + router.get( + { + path: API.DATA_SOURCE.CONNECTIONS, + validate: { + params: schema.object({}, { unknowns: 'allow' }), + }, + }, + async (context, request, response) => { + const fields = ['id', 'title', 'auth.type']; + const resp = await context.core.savedObjects.client.find({ + type: 'data-source', + fields, + perPage: 10000, + }); + + return response.ok({ body: { savedObjects: resp.saved_objects } }); + } + ); + + router.get( + { + path: `${API.DATA_SOURCE.CONNECTIONS}/{dataSourceId}`, + validate: { + params: schema.object({ + dataSourceId: schema.string(), + }), + }, + }, + async (context, request, response) => { + const resp = await context.core.savedObjects.client.get( + 'data-source', + request.params.dataSourceId + ); + return response.ok({ body: resp }); + } + ); +} diff --git a/src/plugins/query_enhancements/server/routes/index.ts b/src/plugins/query_enhancements/server/routes/index.ts new file mode 100644 index 00000000000..8feaecc7b28 --- /dev/null +++ b/src/plugins/query_enhancements/server/routes/index.ts @@ -0,0 +1,139 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { schema } from '@osd/config-schema'; +import { + IOpenSearchDashboardsResponse, + IRouter, + Logger, + ResponseError, +} from '../../../../core/server'; +import { IDataFrameResponse, IOpenSearchDashboardsSearchRequest } from '../../../data/common'; +import { ISearchStrategy } from '../../../data/server'; +import { API, SEARCH_STRATEGY } from '../../common'; +import { registerQueryAssistRoutes } from './query_assist'; +import { registerDataSourceConnectionsRoutes } from './data_source_connection'; + +function defineRoute( + logger: Logger, + router: IRouter, + searchStrategies: Record< + string, + ISearchStrategy + >, + searchStrategyId: string +) { + const path = `${API.SEARCH}/${searchStrategyId}`; + router.post( + { + path, + validate: { + body: schema.object({ + query: schema.object({ + qs: schema.string(), + format: schema.string(), + }), + df: schema.nullable(schema.object({}, { unknowns: 'allow' })), + }), + }, + }, + async (context, req, res): Promise> => { + try { + const queryRes: IDataFrameResponse = await searchStrategies[searchStrategyId].search( + context, + req as any, + {} + ); + return res.ok({ body: { ...queryRes } }); + } catch (err) { + logger.error(err); + return res.custom({ + statusCode: 500, + body: err, + }); + } + } + ); + + router.get( + { + path: `${path}/{queryId}`, + validate: { + params: schema.object({ + queryId: schema.string(), + }), + }, + }, + async (context, req, res): Promise => { + try { + const queryRes: IDataFrameResponse = await searchStrategies[searchStrategyId].search( + context, + req as any, + {} + ); + const result: any = { + body: { + ...queryRes, + }, + }; + return res.ok(result); + } catch (err) { + logger.error(err); + return res.custom({ + statusCode: 500, + body: err, + }); + } + } + ); + + router.get( + { + path: `${path}/{queryId}/{dataSourceId}`, + validate: { + params: schema.object({ + queryId: schema.string(), + dataSourceId: schema.string(), + }), + }, + }, + async (context, req, res): Promise => { + try { + const queryRes: IDataFrameResponse = await searchStrategies[searchStrategyId].search( + context, + req as any, + {} + ); + const result: any = { + body: { + ...queryRes, + }, + }; + return res.ok(result); + } catch (err) { + logger.error(err); + return res.custom({ + statusCode: 500, + body: err, + }); + } + } + ); +} + +export function defineRoutes( + logger: Logger, + router: IRouter, + searchStrategies: Record< + string, + ISearchStrategy + > +) { + defineRoute(logger, router, searchStrategies, SEARCH_STRATEGY.PPL); + defineRoute(logger, router, searchStrategies, SEARCH_STRATEGY.SQL); + defineRoute(logger, router, searchStrategies, SEARCH_STRATEGY.SQL_ASYNC); + registerDataSourceConnectionsRoutes(router); + registerQueryAssistRoutes(router); +} diff --git a/src/plugins/query_enhancements/server/routes/query_assist/agents.test.ts b/src/plugins/query_enhancements/server/routes/query_assist/agents.test.ts new file mode 100644 index 00000000000..f22bdb14836 --- /dev/null +++ b/src/plugins/query_enhancements/server/routes/query_assist/agents.test.ts @@ -0,0 +1,109 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { ApiResponse } from '@opensearch-project/opensearch'; +import { ResponseError } from '@opensearch-project/opensearch/lib/errors'; +import { RequestHandlerContext } from 'src/core/server'; +// // eslint-disable-next-line @osd/eslint/no-restricted-paths +// import { CoreRouteHandlerContext } from 'src/core/server/core_route_handler_context'; +import { coreMock } from '../../../../../core/server/mocks'; +import { loggerMock } from '@osd/logging/target/mocks'; +import { getAgentIdByConfig, requestAgentByConfig } from './agents'; + +describe.skip('Agents helper functions', () => { + // const coreContext = new CoreRouteHandlerContext( + // coreMock.createInternalStart(), + // httpServerMock.createOpenSearchDashboardsRequest() + // ); + const coreContext = coreMock.createRequestHandlerContext(); + const client = coreContext.opensearch.client.asCurrentUser; + const mockedTransport = client.transport.request as jest.Mock; + const context: RequestHandlerContext = { + core: coreContext, + dataSource: jest.fn(), + query_assist: { dataSourceEnabled: false, logger: loggerMock.create() }, + }; + + afterEach(() => { + jest.clearAllMocks(); + }); + + it('searches agent id by name', async () => { + mockedTransport.mockResolvedValueOnce({ + body: { + type: 'agent', + configuration: { agent_id: 'agentId' }, + }, + }); + const id = await getAgentIdByConfig(client, 'test_agent'); + expect(id).toEqual('agentId'); + expect(mockedTransport.mock.calls[0]).toMatchInlineSnapshot(` + Array [ + Object { + "method": "GET", + "path": "/_plugins/_ml/config/test_agent", + }, + ] + `); + }); + + it('handles not found errors', async () => { + mockedTransport.mockRejectedValueOnce( + new ResponseError(({ + body: { + error: { + root_cause: [ + { + type: 'status_exception', + reason: 'Failed to find config with the provided config id: test_agent', + }, + ], + type: 'status_exception', + reason: 'Failed to find config with the provided config id: test_agent', + }, + status: 404, + }, + statusCode: 404, + } as unknown) as ApiResponse) + ); + await expect( + getAgentIdByConfig(client, 'test agent') + ).rejects.toThrowErrorMatchingInlineSnapshot( + `"Get agent 'test agent' failed, reason: {\\"error\\":{\\"root_cause\\":[{\\"type\\":\\"status_exception\\",\\"reason\\":\\"Failed to find config with the provided config id: test_agent\\"}],\\"type\\":\\"status_exception\\",\\"reason\\":\\"Failed to find config with the provided config id: test_agent\\"},\\"status\\":404}"` + ); + }); + + it('handles search errors', async () => { + mockedTransport.mockRejectedValueOnce('request failed'); + await expect( + getAgentIdByConfig(client, 'test agent') + ).rejects.toThrowErrorMatchingInlineSnapshot( + `"Get agent 'test agent' failed, reason: request failed"` + ); + }); + + it('searches for agent id and sends request', async () => { + mockedTransport + .mockResolvedValueOnce({ + body: { + type: 'agent', + configuration: { agent_id: 'new-id' }, + }, + }) + .mockResolvedValueOnce({ + body: { inference_results: [{ output: [{ result: 'test response' }] }] }, + }); + const response = await requestAgentByConfig({ + context, + configName: 'new_agent', + body: { parameters: { param1: 'value1' } }, + }); + expect(mockedTransport).toBeCalledWith( + expect.objectContaining({ path: '/_plugins/_ml/agents/new-id/_execute' }), + expect.anything() + ); + expect(response.body.inference_results[0].output[0].result).toEqual('test response'); + }); +}); diff --git a/src/plugins/query_enhancements/server/routes/query_assist/agents.ts b/src/plugins/query_enhancements/server/routes/query_assist/agents.ts new file mode 100644 index 00000000000..d065254c7f0 --- /dev/null +++ b/src/plugins/query_enhancements/server/routes/query_assist/agents.ts @@ -0,0 +1,71 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { ApiResponse } from '@opensearch-project/opensearch'; +import { RequestBody, TransportRequestPromise } from '@opensearch-project/opensearch/lib/Transport'; +import { RequestHandlerContext } from 'src/core/server'; +import { URI } from '../../../common'; + +const AGENT_REQUEST_OPTIONS = { + /** + * It is time-consuming for LLM to generate final answer + * Give it a large timeout window + */ + requestTimeout: 5 * 60 * 1000, + /** + * Do not retry + */ + maxRetries: 0, +}; + +export type AgentResponse = ApiResponse<{ + inference_results: Array<{ + output: Array<{ name: string; result?: string }>; + }>; +}>; + +type OpenSearchClient = RequestHandlerContext['core']['opensearch']['client']['asCurrentUser']; + +export const getAgentIdByConfig = async ( + client: OpenSearchClient, + configName: string +): Promise => { + try { + const response = (await client.transport.request({ + method: 'GET', + path: `${URI.ML}/config/${configName}`, + })) as ApiResponse<{ type: string; configuration: { agent_id?: string } }>; + + if (!response || response.body.configuration.agent_id === undefined) { + throw new Error('cannot find any agent by configuration: ' + configName); + } + return response.body.configuration.agent_id; + } catch (error) { + const errorMessage = JSON.stringify(error.meta?.body) || error; + throw new Error(`Get agent '${configName}' failed, reason: ` + errorMessage); + } +}; + +export const requestAgentByConfig = async (options: { + context: RequestHandlerContext; + configName: string; + body: RequestBody; + dataSourceId?: string; +}): Promise => { + const { context, configName, body, dataSourceId } = options; + const client = + context.query_assist.dataSourceEnabled && dataSourceId + ? await context.dataSource.opensearch.getClient(dataSourceId) + : context.core.opensearch.client.asCurrentUser; + const agentId = await getAgentIdByConfig(client, configName); + return client.transport.request( + { + method: 'POST', + path: `${URI.ML}/agents/${agentId}/_execute`, + body, + }, + AGENT_REQUEST_OPTIONS + ) as TransportRequestPromise; +}; diff --git a/src/plugins/query_enhancements/server/routes/query_assist/createResponse.ts b/src/plugins/query_enhancements/server/routes/query_assist/createResponse.ts new file mode 100644 index 00000000000..c88cd678a97 --- /dev/null +++ b/src/plugins/query_enhancements/server/routes/query_assist/createResponse.ts @@ -0,0 +1,28 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { QueryAssistResponse } from '../../../common/query_assist'; +import { AgentResponse } from './agents'; +import { createPPLResponseBody } from './ppl/create_response'; + +export const createResponseBody = ( + language: string, + agentResponse: AgentResponse +): QueryAssistResponse => { + switch (language) { + case 'PPL': + return createPPLResponseBody(agentResponse); + + default: + if (!agentResponse.body.inference_results[0].output[0].result) + throw new Error('Generated query not found.'); + const result = JSON.parse( + agentResponse.body.inference_results[0].output[0].result! + ) as Record; + const query = Object.values(result).at(0); + if (typeof query !== 'string') throw new Error('Generated query not found.'); + return { query }; + } +}; diff --git a/src/plugins/query_enhancements/server/routes/query_assist/index.ts b/src/plugins/query_enhancements/server/routes/query_assist/index.ts new file mode 100644 index 00000000000..1e28ddc34b2 --- /dev/null +++ b/src/plugins/query_enhancements/server/routes/query_assist/index.ts @@ -0,0 +1,6 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export { registerQueryAssistRoutes } from './routes'; diff --git a/src/plugins/query_enhancements/server/routes/query_assist/ppl/create_response.ts b/src/plugins/query_enhancements/server/routes/query_assist/ppl/create_response.ts new file mode 100644 index 00000000000..6519eb00ac1 --- /dev/null +++ b/src/plugins/query_enhancements/server/routes/query_assist/ppl/create_response.ts @@ -0,0 +1,21 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { QueryAssistResponse } from '../../../../common/query_assist'; +import { AgentResponse } from '../agents'; + +export const createPPLResponseBody = (agentResponse: AgentResponse): QueryAssistResponse => { + if (!agentResponse.body.inference_results[0].output[0].result) + throw new Error('Generated query not found.'); + const result = JSON.parse(agentResponse.body.inference_results[0].output[0].result!) as { + ppl: string; + }; + const ppl = result.ppl + .replace(/[\r\n]/g, ' ') + .trim() + .replace(/ISNOTNULL/g, 'isnotnull') // https://github.com/opensearch-project/sql/issues/2431 + .replace(/\bSPAN\(/g, 'span('); // https://github.com/opensearch-project/dashboards-observability/issues/759 + return { query: ppl }; +}; diff --git a/src/plugins/query_enhancements/server/routes/query_assist/routes.ts b/src/plugins/query_enhancements/server/routes/query_assist/routes.ts new file mode 100644 index 00000000000..d3039671953 --- /dev/null +++ b/src/plugins/query_enhancements/server/routes/query_assist/routes.ts @@ -0,0 +1,94 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { schema } from '@osd/config-schema'; +import { IRouter } from 'opensearch-dashboards/server'; +// eslint-disable-next-line @osd/eslint/no-restricted-paths +import { isResponseError } from '../../..../../../../../core/server/opensearch/client/errors'; +import { API, ERROR_DETAILS } from '../../../common'; +import { getAgentIdByConfig, requestAgentByConfig } from './agents'; +import { createResponseBody } from './createResponse'; + +export function registerQueryAssistRoutes(router: IRouter) { + router.get( + { + path: API.QUERY_ASSIST.LANGUAGES, + validate: { + query: schema.object({ + dataSourceId: schema.maybe(schema.string()), + }), + }, + }, + async (context, request, response) => { + const config = await context.query_assist.configPromise; + const client = + context.query_assist.dataSourceEnabled && request.query.dataSourceId + ? await context.dataSource.opensearch.getClient(request.query.dataSourceId) + : context.core.opensearch.client.asCurrentUser; + const configuredLanguages: string[] = []; + try { + await Promise.allSettled( + config.queryAssist.supportedLanguages.map((languageConfig) => + // if the call does not throw any error, then the agent is properly configured + getAgentIdByConfig(client, languageConfig.agentConfig).then(() => + configuredLanguages.push(languageConfig.language) + ) + ) + ); + return response.ok({ body: { configuredLanguages } }); + } catch (error) { + return response.ok({ body: { configuredLanguages, error: error.message } }); + } + } + ); + + router.post( + { + path: API.QUERY_ASSIST.GENERATE, + validate: { + body: schema.object({ + index: schema.string(), + question: schema.string(), + language: schema.string(), + dataSourceId: schema.maybe(schema.string()), + }), + }, + }, + async (context, request, response) => { + const config = await context.query_assist.configPromise; + const languageConfig = config.queryAssist.supportedLanguages.find( + (c) => c.language === request.body.language + ); + if (!languageConfig) return response.badRequest({ body: 'Unsupported language' }); + try { + const agentResponse = await requestAgentByConfig({ + context, + configName: languageConfig.agentConfig, + body: { + parameters: { + index: request.body.index, + question: request.body.question, + }, + }, + dataSourceId: request.body.dataSourceId, + }); + const responseBody = createResponseBody(languageConfig.language, agentResponse); + return response.ok({ body: responseBody }); + } catch (error) { + if (isResponseError(error)) { + if (error.statusCode === 400 && error.body.includes(ERROR_DETAILS.GUARDRAILS_TRIGGERED)) + return response.badRequest({ body: ERROR_DETAILS.GUARDRAILS_TRIGGERED }); + return response.badRequest({ + body: + typeof error.meta.body === 'string' + ? error.meta.body + : JSON.stringify(error.meta.body), + }); + } + return response.custom({ statusCode: error.statusCode || 500, body: error.message }); + } + } + ); +} diff --git a/src/plugins/query_enhancements/server/search/index.ts b/src/plugins/query_enhancements/server/search/index.ts new file mode 100644 index 00000000000..b3a528c5705 --- /dev/null +++ b/src/plugins/query_enhancements/server/search/index.ts @@ -0,0 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export { pplSearchStrategyProvider } from './ppl_search_strategy'; +export { sqlSearchStrategyProvider } from './sql_search_strategy'; +export { sqlAsyncSearchStrategyProvider } from './sql_async_search_strategy'; diff --git a/src/plugins/query_enhancements/server/search/ppl_search_strategy.ts b/src/plugins/query_enhancements/server/search/ppl_search_strategy.ts new file mode 100644 index 00000000000..7f2af4d4182 --- /dev/null +++ b/src/plugins/query_enhancements/server/search/ppl_search_strategy.ts @@ -0,0 +1,131 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { first } from 'rxjs/operators'; +import { SharedGlobalConfig, Logger, ILegacyClusterClient } from 'opensearch-dashboards/server'; +import { Observable } from 'rxjs'; +import { ISearchStrategy, getDefaultSearchParams, SearchUsage } from '../../../data/server'; +import { + DATA_FRAME_TYPES, + IDataFrameError, + IDataFrameResponse, + IDataFrameWithAggs, + IOpenSearchDashboardsSearchRequest, + createDataFrame, +} from '../../../data/common'; +import { getFields } from '../../common/utils'; +import { Facet } from '../utils'; + +export const pplSearchStrategyProvider = ( + config$: Observable, + logger: Logger, + client: ILegacyClusterClient, + usage?: SearchUsage +): ISearchStrategy => { + const pplFacet = new Facet({ + client, + logger, + endpoint: 'ppl.pplQuery', + useJobs: false, + shimResponse: true, + }); + + const parseRequest = (query: string) => { + const pipeMap = new Map(); + const pipeArray = query.split('|'); + pipeArray.forEach((pipe, index) => { + const splitChar = index === 0 ? '=' : ' '; + const split = pipe.trim().split(splitChar); + const key = split[0]; + const value = pipe.replace(index === 0 ? `${key}=` : key, '').trim(); + pipeMap.set(key, value); + }); + + const source = pipeMap.get('source'); + + const searchQuery = query; + + const filters = pipeMap.get('where'); + + const stats = pipeMap.get('stats'); + const aggsQuery = stats + ? `source=${source} ${filters ? `| where ${filters}` : ''} | stats ${stats}` + : undefined; + + return { + map: pipeMap, + search: searchQuery, + aggs: aggsQuery, + }; + }; + + return { + search: async (context, request: any, options) => { + const config = await config$.pipe(first()).toPromise(); + const uiSettingsClient = await context.core.uiSettings.client; + + const { dataFrameHydrationStrategy, ...defaultParams } = await getDefaultSearchParams( + uiSettingsClient + ); + + try { + const requestParams = parseRequest(request.body.query.qs); + const source = requestParams?.map.get('source'); + const { schema, meta } = request.body.df; + + request.body.query = + !schema || dataFrameHydrationStrategy === 'perQuery' + ? `source=${source} | head` + : requestParams.search; + const rawResponse: any = await pplFacet.describeQuery(context, request); + + if (!rawResponse.success) { + return { + type: DATA_FRAME_TYPES.DEFAULT, + body: { error: rawResponse.data }, + took: rawResponse.took, + } as IDataFrameError; + } + + const dataFrame = createDataFrame({ + name: source, + schema: schema ?? rawResponse.data.schema, + meta, + fields: getFields(rawResponse), + }); + + dataFrame.size = rawResponse.data.datarows.length; + + if (usage) usage.trackSuccess(rawResponse.took); + + if (dataFrame.meta?.aggsQs) { + for (const [key, aggQueryString] of Object.entries(dataFrame.meta.aggsQs)) { + const aggRequest = parseRequest(aggQueryString as string); + const query = aggRequest.aggs; + request.body.query = query; + const rawAggs: any = await pplFacet.describeQuery(context, request); + (dataFrame as IDataFrameWithAggs).aggs = {}; + (dataFrame as IDataFrameWithAggs).aggs[key] = rawAggs.data.datarows?.map((hit: any) => { + return { + key: hit[1], + value: hit[0], + }; + }); + } + } + + return { + type: DATA_FRAME_TYPES.DEFAULT, + body: dataFrame, + took: rawResponse.took, + } as IDataFrameResponse; + } catch (e) { + logger.error(`pplSearchStrategy: ${e.message}`); + if (usage) usage.trackError(); + throw e; + } + }, + }; +}; diff --git a/src/plugins/query_enhancements/server/search/sql_async_search_strategy.ts b/src/plugins/query_enhancements/server/search/sql_async_search_strategy.ts new file mode 100644 index 00000000000..acd0027d0bc --- /dev/null +++ b/src/plugins/query_enhancements/server/search/sql_async_search_strategy.ts @@ -0,0 +1,112 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { SharedGlobalConfig, Logger, ILegacyClusterClient } from 'opensearch-dashboards/server'; +import { Observable } from 'rxjs'; +import { ISearchStrategy, SearchUsage } from '../../../data/server'; +import { + DATA_FRAME_TYPES, + IDataFrameError, + IDataFrameResponse, + IOpenSearchDashboardsSearchRequest, + PartialDataFrame, + createDataFrame, +} from '../../../data/common'; +import { Facet } from '../utils'; + +export const sqlAsyncSearchStrategyProvider = ( + config$: Observable, + logger: Logger, + client: ILegacyClusterClient, + usage?: SearchUsage +): ISearchStrategy => { + const sqlAsyncFacet = new Facet({ client, logger, endpoint: 'observability.runDirectQuery' }); + const sqlAsyncJobsFacet = new Facet({ + client, + logger, + endpoint: 'observability.getJobStatus', + useJobs: true, + }); + + return { + search: async (context, request: any, options) => { + try { + // Create job: this should return a queryId and sessionId + if (request?.body?.query?.qs) { + const df = request.body?.df; + request.body = { + query: request.body.query.qs, + datasource: df?.meta?.queryConfig?.dataSource, + lang: 'sql', + sessionId: df?.meta?.sessionId, + }; + const rawResponse: any = await sqlAsyncFacet.describeQuery(context, request); + // handles failure + if (!rawResponse.success) { + return { + type: DATA_FRAME_TYPES.POLLING, + body: { error: rawResponse.data }, + took: rawResponse.took, + } as IDataFrameError; + } + const queryId = rawResponse.data?.queryId; + const sessionId = rawResponse.data?.sessionId; + + const partial: PartialDataFrame = { + name: '', + fields: rawResponse?.data?.schema || [], + }; + const dataFrame = createDataFrame(partial); + dataFrame.meta = { + query: request.body.query, + queryId, + sessionId, + }; + dataFrame.name = request.body?.datasource; + return { + type: DATA_FRAME_TYPES.POLLING, + body: dataFrame, + took: rawResponse.took, + } as IDataFrameResponse; + } else { + const queryId = request.params.queryId; + request.params = { queryId }; + const asyncResponse: any = await sqlAsyncJobsFacet.describeQuery(context, request); + const status = asyncResponse.data.status; + const partial: PartialDataFrame = { + name: '', + fields: asyncResponse?.data?.schema || [], + }; + const dataFrame = createDataFrame(partial); + dataFrame.fields?.forEach((field, index) => { + field.values = asyncResponse?.data.datarows.map((row: any) => row[index]); + }); + + dataFrame.size = asyncResponse?.data?.datarows?.length || 0; + + dataFrame.meta = { + status, + queryId, + error: status === 'FAILED' && asyncResponse.data?.error, + }; + dataFrame.name = request.body?.datasource; + + // TODO: MQL should this be the time for polling or the time for job creation? + if (usage) usage.trackSuccess(asyncResponse.took); + + return { + type: DATA_FRAME_TYPES.POLLING, + body: dataFrame, + took: asyncResponse.took, + } as IDataFrameResponse; + } + } catch (e) { + logger.error(`sqlAsyncSearchStrategy: ${e.message}`); + if (usage) usage.trackError(); + throw e; + } + }, + }; +}; diff --git a/src/plugins/query_enhancements/server/search/sql_search_strategy.test.ts b/src/plugins/query_enhancements/server/search/sql_search_strategy.test.ts new file mode 100644 index 00000000000..71fea93c211 --- /dev/null +++ b/src/plugins/query_enhancements/server/search/sql_search_strategy.test.ts @@ -0,0 +1,136 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { sqlSearchStrategyProvider } from './sql_search_strategy'; +import { Observable, of } from 'rxjs'; +import { + SharedGlobalConfig, + Logger, + ILegacyClusterClient, + RequestHandlerContext, +} from 'opensearch-dashboards/server'; +import { SearchUsage } from '../../../data/server'; +import { + DATA_FRAME_TYPES, + IDataFrameError, + IDataFrameResponse, + IOpenSearchDashboardsSearchRequest, +} from '../../../data/common'; +import * as facet from '../utils/facet'; + +describe('sqlSearchStrategyProvider', () => { + let config$: Observable; + let logger: Logger; + let client: ILegacyClusterClient; + let usage: SearchUsage; + const emptyRequestHandlerContext = ({} as unknown) as RequestHandlerContext; + + beforeEach(() => { + config$ = of({} as SharedGlobalConfig); + logger = ({ + error: jest.fn(), + } as unknown) as Logger; + client = {} as ILegacyClusterClient; + usage = { + trackSuccess: jest.fn(), + trackError: jest.fn(), + } as SearchUsage; + }); + + it('should return an object with a search method', () => { + const strategy = sqlSearchStrategyProvider(config$, logger, client, usage); + expect(strategy).toHaveProperty('search'); + expect(typeof strategy.search).toBe('function'); + }); + + it('should handle successful search response', async () => { + const mockResponse = { + success: true, + data: { + schema: [{ name: 'field1' }, { name: 'field2' }], + datarows: [ + [1, 'value1'], + [2, 'value2'], + ], + }, + took: 100, + }; + const mockFacet = ({ + describeQuery: jest.fn().mockResolvedValue(mockResponse), + } as unknown) as facet.Facet; + jest.spyOn(facet, 'Facet').mockImplementation(() => mockFacet); + + const strategy = sqlSearchStrategyProvider(config$, logger, client, usage); + const result = await strategy.search( + emptyRequestHandlerContext, + ({ + body: { query: { qs: 'SELECT * FROM table' } }, + } as unknown) as IOpenSearchDashboardsSearchRequest, + {} + ); + + expect(result).toEqual({ + type: DATA_FRAME_TYPES.DEFAULT, + body: { + name: '', + fields: [ + { name: 'field1', values: [1, 2] }, + { name: 'field2', values: ['value1', 'value2'] }, + ], + size: 2, + }, + took: 100, + } as IDataFrameResponse); + expect(usage.trackSuccess).toHaveBeenCalledWith(100); + }); + + it('should handle failed search response', async () => { + const mockResponse = { + success: false, + data: { cause: 'Query failed' }, + took: 50, + }; + const mockFacet = ({ + describeQuery: jest.fn().mockResolvedValue(mockResponse), + } as unknown) as facet.Facet; + jest.spyOn(facet, 'Facet').mockImplementation(() => mockFacet); + + const strategy = sqlSearchStrategyProvider(config$, logger, client, usage); + const result = await strategy.search( + emptyRequestHandlerContext, + ({ + body: { query: { qs: 'SELECT * FROM table' } }, + } as unknown) as IOpenSearchDashboardsSearchRequest, + {} + ); + + expect(result).toEqual(({ + type: DATA_FRAME_TYPES.DEFAULT, + body: { error: { cause: 'Query failed' } }, + took: 50, + } as unknown) as IDataFrameError); + }); + + it('should handle exceptions', async () => { + const mockError = new Error('Something went wrong'); + const mockFacet = ({ + describeQuery: jest.fn().mockRejectedValue(mockError), + } as unknown) as facet.Facet; + jest.spyOn(facet, 'Facet').mockImplementation(() => mockFacet); + + const strategy = sqlSearchStrategyProvider(config$, logger, client, usage); + await expect( + strategy.search( + emptyRequestHandlerContext, + ({ + body: { query: { qs: 'SELECT * FROM table' } }, + } as unknown) as IOpenSearchDashboardsSearchRequest, + {} + ) + ).rejects.toThrow(mockError); + expect(logger.error).toHaveBeenCalledWith(`sqlSearchStrategy: ${mockError.message}`); + expect(usage.trackError).toHaveBeenCalled(); + }); +}); diff --git a/src/plugins/query_enhancements/server/search/sql_search_strategy.ts b/src/plugins/query_enhancements/server/search/sql_search_strategy.ts new file mode 100644 index 00000000000..c5ebb40f882 --- /dev/null +++ b/src/plugins/query_enhancements/server/search/sql_search_strategy.ts @@ -0,0 +1,66 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { SharedGlobalConfig, Logger, ILegacyClusterClient } from 'opensearch-dashboards/server'; +import { Observable } from 'rxjs'; +import { ISearchStrategy, SearchUsage } from '../../../data/server'; +import { + DATA_FRAME_TYPES, + IDataFrameError, + IDataFrameResponse, + IOpenSearchDashboardsSearchRequest, + PartialDataFrame, + createDataFrame, +} from '../../../data/common'; +import { Facet } from '../utils'; + +export const sqlSearchStrategyProvider = ( + _config$: Observable, + logger: Logger, + client: ILegacyClusterClient, + usage?: SearchUsage +): ISearchStrategy => { + const sqlFacet = new Facet({ client, logger, endpoint: 'ppl.sqlQuery' }); + + return { + search: async (context, request: any, _options) => { + try { + request.body.query = request.body.query.qs; + const rawResponse: any = await sqlFacet.describeQuery(context, request); + + if (!rawResponse.success) { + return { + type: DATA_FRAME_TYPES.DEFAULT, + body: { error: rawResponse.data }, + took: rawResponse.took, + } as IDataFrameError; + } + + const partial: PartialDataFrame = { + name: '', + fields: rawResponse.data?.schema || [], + }; + const dataFrame = createDataFrame(partial); + dataFrame.fields?.forEach((field, index) => { + field.values = rawResponse.data.datarows.map((row: any) => row[index]); + }); + + dataFrame.size = rawResponse.data.datarows?.length || 0; + + if (usage) usage.trackSuccess(rawResponse.took); + + return { + type: DATA_FRAME_TYPES.DEFAULT, + body: dataFrame, + took: rawResponse.took, + } as IDataFrameResponse; + } catch (e) { + logger.error(`sqlSearchStrategy: ${e.message}`); + if (usage) usage.trackError(); + throw e; + } + }, + }; +}; diff --git a/src/plugins/query_enhancements/server/types.ts b/src/plugins/query_enhancements/server/types.ts new file mode 100644 index 00000000000..2ab716b9892 --- /dev/null +++ b/src/plugins/query_enhancements/server/types.ts @@ -0,0 +1,67 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { PluginSetup } from 'src/plugins/data/server'; +import { DataSourcePluginSetup } from '../../data_source/server'; +import { Logger } from '../../../core/server'; +import { ConfigSchema } from '../common/config'; + +// eslint-disable-next-line @typescript-eslint/no-empty-interface +export interface QueryEnhancementsPluginSetup {} +// eslint-disable-next-line @typescript-eslint/no-empty-interface +export interface QueryEnhancementsPluginStart {} + +export interface QueryEnhancementsPluginSetupDependencies { + data: PluginSetup; + dataSource?: DataSourcePluginSetup; +} + +export interface ISchema { + name: string; + type: string; +} + +export interface IPPLDataSource { + jsonData?: any[]; +} + +export interface IPPLVisualizationDataSource extends IPPLDataSource { + data: any; + metadata: any; + size: number; + status: number; +} + +export interface IPPLEventsDataSource extends IPPLDataSource { + schema: ISchema[]; + datarows: any[]; +} + +export interface FacetResponse { + success: boolean; + data: any; +} + +export interface FacetRequest { + body: { + query: string; + format?: string; + }; +} + +declare module '../../../core/server' { + interface RequestHandlerContext { + query_assist: { + logger: Logger; + configPromise: Promise; + dataSourceEnabled: boolean; + }; + data_source_connection: { + logger: Logger; + configPromise: Promise; + dataSourceEnabled: boolean; + }; + } +} diff --git a/src/plugins/query_enhancements/server/utils/facet.ts b/src/plugins/query_enhancements/server/utils/facet.ts new file mode 100644 index 00000000000..a6f23efba2a --- /dev/null +++ b/src/plugins/query_enhancements/server/utils/facet.ts @@ -0,0 +1,104 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { Logger } from 'opensearch-dashboards/server'; +import { FacetResponse, IPPLEventsDataSource, IPPLVisualizationDataSource } from '../types'; +import { shimSchemaRow, shimStats } from '.'; + +export interface FacetProps { + client: any; + logger: Logger; + endpoint: string; + useJobs?: boolean; + shimResponse?: boolean; +} + +export class Facet { + private defaultClient: any; + private logger: Logger; + private endpoint: string; + private useJobs: boolean; + private shimResponse: boolean; + + constructor({ client, logger, endpoint, useJobs = false, shimResponse = false }: FacetProps) { + this.defaultClient = client; + this.logger = logger; + this.endpoint = endpoint; + this.useJobs = useJobs; + this.shimResponse = shimResponse; + } + + protected fetch = async ( + context: any, + request: any, + endpoint: string + ): Promise => { + try { + const { format, df, ...query } = request.body; + const params = { + body: { ...query }, + ...(format !== 'jdbc' && { format }), + }; + const dataSourceId = df?.meta?.queryConfig?.dataSourceId; + const client = dataSourceId + ? context.dataSource.opensearch.legacy.getClient(dataSourceId).callAPI + : this.defaultClient.asScoped(request).callAsCurrentUser; + const queryRes = await client(endpoint, params); + return { + success: true, + data: queryRes, + }; + } catch (err: any) { + this.logger.error(`Facet fetch: ${endpoint}: ${err}`); + return { + success: false, + data: err, + }; + } + }; + + protected fetchJobs = async ( + context: any, + request: any, + endpoint: string + ): Promise => { + try { + const params = request.params; + const { df } = request.body; + const dataSourceId = df?.meta?.queryConfig?.dataSourceId; + const client = dataSourceId + ? context.dataSource.opensearch.legacy.getClient(dataSourceId).callAPI + : this.defaultClient.asScoped(request).callAsCurrentUser; + const queryRes = await client(endpoint, params); + return { + success: true, + data: queryRes, + }; + } catch (err: any) { + this.logger.error(`Facet fetch: ${endpoint}: ${err}`); + return { + success: false, + data: err, + }; + } + }; + + public describeQuery = async (context: any, request: any): Promise => { + const response = this.useJobs + ? await this.fetchJobs(context, request, this.endpoint) + : await this.fetch(context, request, this.endpoint); + if (!this.shimResponse) return response; + + const { format: dataType } = request.body; + const shimFunctions: { [key: string]: (data: any) => any } = { + jdbc: (data: any) => shimSchemaRow(data as IPPLEventsDataSource), + viz: (data: any) => shimStats(data as IPPLVisualizationDataSource), + }; + + return shimFunctions[dataType] + ? { ...response, data: shimFunctions[dataType](response.data) } + : response; + }; +} diff --git a/src/plugins/query_enhancements/server/utils/index.ts b/src/plugins/query_enhancements/server/utils/index.ts new file mode 100644 index 00000000000..86f21697377 --- /dev/null +++ b/src/plugins/query_enhancements/server/utils/index.ts @@ -0,0 +1,9 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export { Facet, FacetProps } from './facet'; +export { OpenSearchPPLPlugin, OpenSearchObservabilityPlugin } from './plugins'; +export { shimStats } from './shim_stats'; +export { shimSchemaRow } from './shim_schema_row'; diff --git a/src/plugins/query_enhancements/server/utils/plugins.ts b/src/plugins/query_enhancements/server/utils/plugins.ts new file mode 100644 index 00000000000..e8d581bc42f --- /dev/null +++ b/src/plugins/query_enhancements/server/utils/plugins.ts @@ -0,0 +1,159 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { OPENSEARCH_API, URI } from '../../common'; + +const createAction = ( + client: any, + components: any, + options: { + endpoint: string; + method: 'GET' | 'POST' | 'DELETE' | 'PATCH' | 'PUT'; + needBody?: boolean; + paramKey?: string; + params?: any; + } +) => { + const { endpoint, method, needBody = false, paramKey, params } = options; + let urlConfig; + + if (paramKey) { + urlConfig = { + fmt: `${endpoint}/<%=${paramKey}%>`, + req: { + [paramKey]: { + type: 'string', + required: true, + }, + }, + }; + } else if (params) { + urlConfig = { + fmt: endpoint, + params, + }; + } else { + urlConfig = { fmt: endpoint }; + } + + return components.clientAction.factory({ + url: urlConfig, + needBody, + method, + }); +}; + +export const OpenSearchPPLPlugin = (client: any, config: any, components: any) => { + client.prototype.ppl = components.clientAction.namespaceFactory(); + const ppl = client.prototype.ppl.prototype; + + ppl.pplQuery = createAction(client, components, { + endpoint: URI.PPL, + method: 'POST', + needBody: true, + }); + ppl.sqlQuery = createAction(client, components, { + endpoint: URI.SQL, + method: 'POST', + needBody: true, + }); + ppl.getDataConnectionById = createAction(client, components, { + endpoint: OPENSEARCH_API.DATA_CONNECTIONS, + method: 'GET', + paramKey: 'dataconnection', + }); + ppl.deleteDataConnection = createAction(client, components, { + endpoint: OPENSEARCH_API.DATA_CONNECTIONS, + method: 'DELETE', + paramKey: 'dataconnection', + }); + ppl.createDataSource = createAction(client, components, { + endpoint: OPENSEARCH_API.DATA_CONNECTIONS, + method: 'POST', + needBody: true, + }); + ppl.modifyDataConnection = createAction(client, components, { + endpoint: OPENSEARCH_API.DATA_CONNECTIONS, + method: 'PATCH', + needBody: true, + }); + ppl.getDataConnections = createAction(client, components, { + endpoint: OPENSEARCH_API.DATA_CONNECTIONS, + method: 'GET', + }); +}; + +export const OpenSearchObservabilityPlugin = (client: any, config: any, components: any) => { + client.prototype.observability = components.clientAction.namespaceFactory(); + const observability = client.prototype.observability.prototype; + + observability.getObject = createAction(client, components, { + endpoint: OPENSEARCH_API.PANELS, + method: 'GET', + params: { + objectId: { type: 'string' }, + objectIdList: { type: 'string' }, + objectType: { type: 'string' }, + sortField: { type: 'string' }, + sortOrder: { type: 'string' }, + fromIndex: { type: 'number' }, + maxItems: { type: 'number' }, + name: { type: 'string' }, + lastUpdatedTimeMs: { type: 'string' }, + createdTimeMs: { type: 'string' }, + }, + }); + + observability.getObjectById = createAction(client, components, { + endpoint: `${OPENSEARCH_API.PANELS}/<%=objectId%>`, + method: 'GET', + paramKey: 'objectId', + }); + + observability.createObject = createAction(client, components, { + endpoint: OPENSEARCH_API.PANELS, + method: 'POST', + needBody: true, + }); + + observability.updateObjectById = createAction(client, components, { + endpoint: `${OPENSEARCH_API.PANELS}/<%=objectId%>`, + method: 'PUT', + paramKey: 'objectId', + needBody: true, + }); + + observability.deleteObjectById = createAction(client, components, { + endpoint: `${OPENSEARCH_API.PANELS}/<%=objectId%>`, + method: 'DELETE', + paramKey: 'objectId', + }); + + observability.deleteObjectByIdList = createAction(client, components, { + endpoint: OPENSEARCH_API.PANELS, + method: 'DELETE', + params: { + objectIdList: { type: 'string', required: true }, + }, + }); + + observability.getJobStatus = createAction(client, components, { + endpoint: `${URI.ASYNC_QUERY}`, + method: 'GET', + paramKey: 'queryId', + }); + + observability.deleteJob = createAction(client, components, { + endpoint: `${URI.ASYNC_QUERY}/<%=queryId%>`, + method: 'DELETE', + paramKey: 'queryId', + }); + + observability.runDirectQuery = createAction(client, components, { + endpoint: URI.ASYNC_QUERY, + method: 'POST', + needBody: true, + }); +}; diff --git a/src/plugins/query_enhancements/server/utils/shim_schema_row.ts b/src/plugins/query_enhancements/server/utils/shim_schema_row.ts new file mode 100644 index 00000000000..a38d77553e7 --- /dev/null +++ b/src/plugins/query_enhancements/server/utils/shim_schema_row.ts @@ -0,0 +1,28 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { IPPLEventsDataSource } from '../types'; + +export function shimSchemaRow(response: IPPLEventsDataSource) { + const schemaLength = response.schema.length; + + const data = response.datarows.map((row) => { + return row.reduce((record: any, item: any, index: number) => { + if (index < schemaLength) { + const cur = response.schema[index]; + const value = + typeof item === 'object' + ? JSON.stringify(item) + : typeof item === 'boolean' + ? item.toString() + : item; + record[cur.name] = value; + } + return record; + }, {}); + }); + + return { ...response, jsonData: data }; +} diff --git a/src/plugins/query_enhancements/server/utils/shim_stats.ts b/src/plugins/query_enhancements/server/utils/shim_stats.ts new file mode 100644 index 00000000000..e616996e750 --- /dev/null +++ b/src/plugins/query_enhancements/server/utils/shim_stats.ts @@ -0,0 +1,48 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { IPPLVisualizationDataSource } from '../types'; + +/** + * Add vis mapping for runtime fields + * json data structure added to response will be + * [{ + * agent: "mozilla", + * avg(bytes): 5756 + * ... + * }, { + * agent: "MSIE", + * avg(bytes): 5605 + * ... + * }, { + * agent: "chrome", + * avg(bytes): 5648 + * ... + * }] + * + * @internal + */ +export function shimStats(response: IPPLVisualizationDataSource) { + if (!response?.metadata?.fields || !response?.data) { + return { ...response }; + } + + const { + data: statsDataSet, + metadata: { fields: queriedFields }, + size, + } = response; + const data = new Array(size).fill(null).map((_, i) => { + const entry: Record = {}; + queriedFields.forEach(({ name }: { name: string }) => { + if (statsDataSet[name] && i < statsDataSet[name].length) { + entry[name] = statsDataSet[name][i]; + } + }); + return entry; + }); + + return { ...response, jsonData: data }; +} diff --git a/src/plugins/query_enhancements/test/__mocks__/fileMock.js b/src/plugins/query_enhancements/test/__mocks__/fileMock.js new file mode 100644 index 00000000000..03b436dba6c --- /dev/null +++ b/src/plugins/query_enhancements/test/__mocks__/fileMock.js @@ -0,0 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +// some svg files can be used as EuiIcon types. exporting a valid EuiIcon type +// avoids unexpected Eui behaviors +module.exports = 'power'; diff --git a/src/plugins/query_enhancements/test/__mocks__/styleMock.js b/src/plugins/query_enhancements/test/__mocks__/styleMock.js new file mode 100644 index 00000000000..28de3c8b1ec --- /dev/null +++ b/src/plugins/query_enhancements/test/__mocks__/styleMock.js @@ -0,0 +1,6 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +module.exports = {}; diff --git a/src/plugins/query_enhancements/test/jest.config.js b/src/plugins/query_enhancements/test/jest.config.js new file mode 100644 index 00000000000..f5a1419f3ba --- /dev/null +++ b/src/plugins/query_enhancements/test/jest.config.js @@ -0,0 +1,31 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +process.env.TZ = 'UTC'; + +module.exports = { + rootDir: '../', + setupFiles: ['/test/setupTests.ts'], + setupFilesAfterEnv: ['/test/setup.jest.ts'], + roots: [''], + testMatch: ['**/*.test.js', '**/*.test.jsx', '**/*.test.ts', '**/*.test.tsx'], + clearMocks: true, + modulePathIgnorePatterns: ['/offline-module-cache/'], + testPathIgnorePatterns: ['/build/', '/node_modules/', '/__utils__/'], + snapshotSerializers: ['enzyme-to-json/serializer'], + coveragePathIgnorePatterns: [ + '/build/', + '/node_modules/', + '/test/', + '/public/requests/', + '/__utils__/', + ], + moduleNameMapper: { + '\\.(css|less|sass|scss)$': '/test/__mocks__/styleMock.js', + '\\.(gif|ttf|eot|svg|png)$': '/test/__mocks__/fileMock.js', + '^!!raw-loader!.*': 'jest-raw-loader', + }, + testEnvironment: 'jsdom', +}; diff --git a/src/plugins/query_enhancements/test/setup.jest.ts b/src/plugins/query_enhancements/test/setup.jest.ts new file mode 100644 index 00000000000..43627882b00 --- /dev/null +++ b/src/plugins/query_enhancements/test/setup.jest.ts @@ -0,0 +1,47 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { configure } from '@testing-library/react'; +import { TextDecoder, TextEncoder } from 'util'; +import '@testing-library/jest-dom'; + +configure({ testIdAttribute: 'data-test-subj' }); + +// https://github.com/inrupt/solid-client-authn-js/issues/1676#issuecomment-917016646 +global.TextEncoder = TextEncoder; +global.TextDecoder = TextDecoder as typeof global.TextDecoder; + +window.URL.createObjectURL = () => ''; +HTMLCanvasElement.prototype.getContext = () => '' as any; +Element.prototype.scrollIntoView = jest.fn(); +window.IntersectionObserver = (class IntersectionObserver { + constructor() {} + + disconnect() { + return null; + } + + observe() { + return null; + } + + takeRecords() { + return null; + } + + unobserve() { + return null; + } +} as unknown) as typeof window.IntersectionObserver; + +jest.mock('@elastic/eui/lib/components/form/form_row/make_id', () => () => 'random-id'); + +jest.mock('@elastic/eui/lib/services/accessibility/html_id_generator', () => ({ + htmlIdGenerator: () => { + return () => 'random_html_id'; + }, +})); + +jest.setTimeout(30000); diff --git a/src/plugins/query_enhancements/test/setupTests.ts b/src/plugins/query_enhancements/test/setupTests.ts new file mode 100644 index 00000000000..5a996f6f8c9 --- /dev/null +++ b/src/plugins/query_enhancements/test/setupTests.ts @@ -0,0 +1,7 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +require('babel-polyfill'); +require('core-js/stable');