From 9194d83303879b224a0d416b03be4cadfe855f7f Mon Sep 17 00:00:00 2001 From: Kawika Avilla Date: Tue, 16 Jul 2024 16:49:00 -0700 Subject: [PATCH] [Discover-next] add query enhancements plugin as a core plugin (#7212) ### Description Introduces the addition of a new core plugin called `Query Enhancements` This plugin seems to be focused on enhancing query capabilities, particularly in areas like PPL (Piped Processing Language), SQL, and natural language query assistance. #### Plugin Configuration and Setup The plugin is set up with the necessary configuration files: - src/plugins/query_enhancements/opensearch_dashboards.json: Defines the plugin's metadata and dependencies. - src/plugins/query_enhancements/common/config.ts: Contains the configuration schema for the plugin. - src/plugins/query_enhancements/public/plugin.tsx: The main plugin file that sets up the public-facing part of the plugin. - src/plugins/query_enhancements/server/plugin.ts: The server-side setup for the plugin. #### Query Assist Feature A significant part of the plugin is dedicated to a "Query Assist" feature: - src/plugins/query_enhancements/public/query_assist/: This directory contains components, hooks, and utilities for the Query Assist feature. - src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.tsx: A React component for the Query Assist bar. - src/plugins/query_enhancements/server/routes/query_assist/: Server-side routing for Query Assist functionality. #### Search Interceptors The plugin introduces several search interceptors: - src/plugins/query_enhancements/public/search/ppl_search_interceptor.ts - src/plugins/query_enhancements/public/search/sql_search_interceptor.ts - src/plugins/query_enhancements/public/search/sql_async_search_interceptor.ts These interceptors likely modify or enhance the search functionality for different query languages. #### Data Source Connection There's a new feature for data source connections: - src/plugins/query_enhancements/public/data_source_connection/: This directory contains components and services for managing data source connections. #### Server-side Search Strategies The plugin implements server-side search strategies for different query types: - src/plugins/query_enhancements/server/search/ppl_search_strategy.ts - src/plugins/query_enhancements/server/search/sql_search_strategy.ts - src/plugins/query_enhancements/server/search/sql_async_search_strategy.ts #### Modifications to Existing Files There are some modifications to existing files outside the new plugin directory: - src/plugins/data/public/index.ts: Exports new types and interfaces. - src/plugins/data/public/ui/index.ts: Exports new UI components. - src/plugins/data/public/ui/query_editor/index.tsx: Modifies the QueryEditor component. #### Testing The plugin includes a comprehensive test setup: - src/plugins/query_enhancements/test/: Contains Jest configuration and setup files for testing. #### Acknowledgements Contains work by @sejli @joshuali925 @Swiddis @paulstn @mengweieric Thx @ashwin-pc for the description summary ### Issues Related https://github.com/opensearch-project/OpenSearch-Dashboards/issues/6072 https://github.com/opensearch-project/OpenSearch-Dashboards/issues/6074 https://github.com/opensearch-project/OpenSearch-Dashboards/issues/6075 Commits include: * init Signed-off-by: Kawika Avilla * fixing bugs and removing target Signed-off-by: Sean Li * add target and build to .gitignore Signed-off-by: Sean Li * Remove SQL Async and Query Assist Going for the light weight solution for 2.15. So took out what wasn't required deliverable. However, I created a branch `feature/discover-next` which has the state where all the features where added so we can port it over post 2.15. Signed-off-by: Kawika Avilla * adding fix for PPL queries Signed-off-by: Sean Li * Clean up files and add helper functions Signed-off-by: Kawika Avilla * final touches Signed-off-by: Kawika Avilla * [Discover-next] add query assist to query enhancements plugin (#6895) it adds query assist specific logic in query enhancements plugin to show a UI above the PPL search bar if user has configured PPL agent. Issues Resolved: #6820 * add query assist to query enhancements Signed-off-by: Joshua Li * align language to uppercase Signed-off-by: Joshua Li * pick PR 6167 Signed-off-by: Joshua Li * use useState hooks for query assist There is a bug in data explorer `AppContainer` where memorized `DiscoverCanvas` gets unmounted after `setQuery`. PR 6167 works around it by memorizing `AppContainer`. As query assist is no longer being unmounted, we can use proper hooks to persist state now. Signed-off-by: Joshua Li * Revert "pick PR 6167" This reverts commit acb0d41937e30bd76c666a225407743243692d11. Wait for official 6167 to merge to avoid conflict Signed-off-by: Joshua Li * address comments for PR 6894 Signed-off-by: Joshua Li --------- Signed-off-by: Joshua Li (cherry picked from commit 016e0f2f73efd8bb0649151908c67dd7ac09d174) * [Discover-next] Address comments for search bar extensions and query assist (#6933) * pass dependencies to isEnabled func Signed-off-by: Joshua Li * add lazy and memo to search bar extensions Signed-off-by: Joshua Li * move ppl specific string out from query assist Signed-off-by: Joshua Li * prevent setstate after hook unmounts Signed-off-by: Joshua Li * add max-height to search bar extensions Signed-off-by: Joshua Li * prevent setstate after component unmounts Signed-off-by: Joshua Li * move ml-commons API to common/index.ts Signed-off-by: Joshua Li * improve i18n and accessibility usages Signed-off-by: Joshua Li * add hard-coded suggestions for sample data indices Signed-off-by: Joshua Li --------- Signed-off-by: Joshua Li (cherry picked from commit 4aade0f993559b0bae9cbcee8e889868afa88547) * [Discover-next] Support data sources for query assist (#6972) * disable query assist for non-default datasource Signed-off-by: Joshua Li * disable query assist input when loading Signed-off-by: Joshua Li * support MDS for query assist Signed-off-by: Joshua Li * add unit tests for agents Signed-off-by: Joshua Li * Revert "add unit tests for agents" This reverts commit 983514ee11362c5efe4cdb59802b3ff402b61ef2. The test configs are not yet setup in query_enhancements plugins. Signed-off-by: Joshua Li --------- Signed-off-by: Joshua Li (cherry picked from commit 328e08e688c39de1f47fee1c357e9928c0576390) * add query assist banner Signed-off-by: Joshua Li * implement banner rendering logic Signed-off-by: Joshua Li * pick opensearch-project/dashboards-observability/pull/1896 Signed-off-by: Joshua Li * Conditionally render the data source select config Signed-off-by: Kawika Avilla * restore no space Signed-off-by: Kawika Avilla * add basic tests Signed-off-by: Joshua Li * add index selector This is a temporary solution given that in discover the index pattern selector will be removed. Before datasource and dataset selectors are added, query assist will rely on this index pattern selector to determine which index user wants to query. Signed-off-by: Joshua Li * MDS working with PPL and SQL Utilizing the work from MDS to make PPL and SQL calls to a remote cluster. Signed-off-by: Kawika Avilla * rename logo to mark Signed-off-by: Joshua Li * extract supported languages to config Signed-off-by: Joshua Li * Add missing license headers Signed-off-by: Kawika Avilla * initial implementation for async queries Signed-off-by: Sean Li * remove queryId and sessionId from facet Signed-off-by: Sean Li * fix: Update call out tests with intl wrapper Signed-off-by: Simeon Widdis * maintenance: Add and apply lint command Signed-off-by: Simeon Widdis * changing to observables Signed-off-by: Sean Li * search interceptor and cleanup Signed-off-by: Sean Li * moving polling class into plugin Signed-off-by: Sean Li * Some minor clean ups plus lint Signed-off-by: Kawika Avilla * test: Add mock-heavy tests for sql search strategy Signed-off-by: Simeon Widdis * test: add GH workflows Signed-off-by: Simeon Widdis * More clean ups Add some props interfaces and consolidated some of the facets Signed-off-by: Kawika Avilla * chore: apply auto-lints to plugin Signed-off-by: Simeon Widdis * Not done Signed-off-by: Kawika Avilla * working Signed-off-by: Kawika Avilla * working and styled Signed-off-by: Kawika Avilla * add connection service Signed-off-by: Kawika Avilla * Update public/types.ts Co-authored-by: Joshua Li * Update public/plugin.tsx Co-authored-by: Joshua Li * Update public/plugin.tsx Co-authored-by: Joshua Li * fix linter Signed-off-by: Kawika Avilla * add more tests for query assist (#31) add more tests for query assist update eslint config to make it work increase index selector width Signed-off-by: Joshua Li * fix!: remove dataSource key in get connection response The `Connection` type does not have the extra `dataSource` object around response, and most of the time code references connection over dataSource. remove the redundent object to make response consistent with type. Signed-off-by: Joshua Li * change isEnabled to an observable based on selected connection Signed-off-by: Joshua Li * protect duplicated calls Signed-off-by: Joshua Li * [build] import via start services Signed-off-by: Kawika Avilla * Changeset file for PR #7212 created/updated * foobar Signed-off-by: Kawika Avilla * updating imports still in progress Signed-off-by: Kawika Avilla * Update more imports Signed-off-by: Kawika Avilla * fix setting of state Signed-off-by: Kawika Avilla * fixes bugs Signed-off-by: Kawika Avilla * Update test Signed-off-by: Kawika Avilla --------- Signed-off-by: Kawika Avilla Signed-off-by: Sean Li Signed-off-by: Joshua Li Signed-off-by: Simeon Widdis Co-authored-by: Sean Li Co-authored-by: Joshua Li Co-authored-by: Simeon Widdis Co-authored-by: opensearch-changeset-bot[bot] <154024398+opensearch-changeset-bot[bot]@users.noreply.github.com> Signed-off-by: Kawika Avilla --- changelogs/fragments/7212.yml | 2 + docs/_sidebar.md | 1 + src/plugins/data/public/index.ts | 6 + src/plugins/data/public/ui/index.ts | 8 + .../data/public/ui/query_editor/index.tsx | 6 +- .../public/ui/query_editor/query_editor.tsx | 20 +- .../query_editor_extensions/index.tsx | 5 +- src/plugins/query_enhancements/.i18nrc.json | 7 + src/plugins/query_enhancements/README.md | 9 + .../query_enhancements/common/config.ts | 23 ++ .../query_enhancements/common/constants.ts | 47 ++++ .../query_enhancements/common/index.ts | 8 + .../common/query_assist/index.ts | 6 + .../common/query_assist/types.ts | 19 ++ .../query_enhancements/common/types.ts | 15 ++ .../query_enhancements/common/utils.ts | 156 +++++++++++++ .../opensearch_dashboards.json | 9 + .../public/assets/query_assist_mark.svg | 18 ++ .../components/connections_bar.tsx | 93 ++++++++ .../components/index.ts | 6 + .../public/data_source_connection/index.ts | 7 + .../services/connections_service.ts | 84 +++++++ .../data_source_connection/services/index.ts | 6 + .../utils/create_extension.tsx | 34 +++ .../data_source_connection/utils/index.ts | 6 + .../query_enhancements/public/index.scss | 1 + .../query_enhancements/public/index.ts | 14 ++ .../query_enhancements/public/plugin.tsx | 181 +++++++++++++++ .../__snapshots__/call_outs.test.tsx.snap | 122 ++++++++++ .../components/call_outs.test.tsx | 58 +++++ .../query_assist/components/call_outs.tsx | 106 +++++++++ .../public/query_assist/components/index.ts | 7 + .../components/index_selector.tsx | 47 ++++ .../components/query_assist_banner.test.tsx | 50 ++++ .../components/query_assist_banner.tsx | 72 ++++++ .../components/query_assist_bar.tsx | 121 ++++++++++ .../components/query_assist_input.test.tsx | 76 ++++++ .../components/query_assist_input.tsx | 114 +++++++++ .../components/submit_button.test.tsx | 34 +++ .../query_assist/components/submit_button.tsx | 25 ++ .../public/query_assist/hooks/index.ts | 6 + .../query_assist/hooks/use_generate.test.ts | 92 ++++++++ .../public/query_assist/hooks/use_generate.ts | 51 +++++ .../public/query_assist/hooks/use_indices.ts | 97 ++++++++ .../public/query_assist/index.ts | 6 + .../utils/create_extension.test.tsx | 104 +++++++++ .../query_assist/utils/create_extension.tsx | 126 ++++++++++ .../public/query_assist/utils/errors.test.ts | 65 ++++++ .../public/query_assist/utils/errors.ts | 30 +++ .../query_assist/utils/get_persisted_log.ts | 24 ++ .../public/query_assist/utils/index.ts | 8 + .../query_enhancements/public/search/index.ts | 8 + .../public/search/ppl_search_interceptor.ts | 216 ++++++++++++++++++ .../search/sql_async_search_interceptor.ts | 137 +++++++++++ .../public/search/sql_search_interceptor.ts | 92 ++++++++ .../query_enhancements/public/services.ts | 11 + .../query_enhancements/public/types.ts | 38 +++ .../query_enhancements/server/index.ts | 29 +++ .../query_enhancements/server/plugin.ts | 99 ++++++++ .../routes/data_source_connection/index.ts | 6 + .../routes/data_source_connection/routes.ts | 48 ++++ .../query_enhancements/server/routes/index.ts | 139 +++++++++++ .../server/routes/query_assist/agents.test.ts | 109 +++++++++ .../server/routes/query_assist/agents.ts | 71 ++++++ .../routes/query_assist/createResponse.ts | 28 +++ .../server/routes/query_assist/index.ts | 6 + .../query_assist/ppl/create_response.ts | 21 ++ .../server/routes/query_assist/routes.ts | 94 ++++++++ .../query_enhancements/server/search/index.ts | 8 + .../server/search/ppl_search_strategy.ts | 131 +++++++++++ .../search/sql_async_search_strategy.ts | 112 +++++++++ .../server/search/sql_search_strategy.test.ts | 136 +++++++++++ .../server/search/sql_search_strategy.ts | 66 ++++++ .../query_enhancements/server/types.ts | 67 ++++++ .../query_enhancements/server/utils/facet.ts | 104 +++++++++ .../query_enhancements/server/utils/index.ts | 9 + .../server/utils/plugins.ts | 159 +++++++++++++ .../server/utils/shim_schema_row.ts | 28 +++ .../server/utils/shim_stats.ts | 48 ++++ .../test/__mocks__/fileMock.js | 8 + .../test/__mocks__/styleMock.js | 6 + .../query_enhancements/test/jest.config.js | 31 +++ .../query_enhancements/test/setup.jest.ts | 47 ++++ .../query_enhancements/test/setupTests.ts | 7 + 84 files changed, 4344 insertions(+), 13 deletions(-) create mode 100644 changelogs/fragments/7212.yml create mode 100644 src/plugins/query_enhancements/.i18nrc.json create mode 100755 src/plugins/query_enhancements/README.md create mode 100644 src/plugins/query_enhancements/common/config.ts create mode 100644 src/plugins/query_enhancements/common/constants.ts create mode 100644 src/plugins/query_enhancements/common/index.ts create mode 100644 src/plugins/query_enhancements/common/query_assist/index.ts create mode 100644 src/plugins/query_enhancements/common/query_assist/types.ts create mode 100644 src/plugins/query_enhancements/common/types.ts create mode 100644 src/plugins/query_enhancements/common/utils.ts create mode 100644 src/plugins/query_enhancements/opensearch_dashboards.json create mode 100644 src/plugins/query_enhancements/public/assets/query_assist_mark.svg create mode 100644 src/plugins/query_enhancements/public/data_source_connection/components/connections_bar.tsx create mode 100644 src/plugins/query_enhancements/public/data_source_connection/components/index.ts create mode 100644 src/plugins/query_enhancements/public/data_source_connection/index.ts create mode 100644 src/plugins/query_enhancements/public/data_source_connection/services/connections_service.ts create mode 100644 src/plugins/query_enhancements/public/data_source_connection/services/index.ts create mode 100644 src/plugins/query_enhancements/public/data_source_connection/utils/create_extension.tsx create mode 100644 src/plugins/query_enhancements/public/data_source_connection/utils/index.ts create mode 100644 src/plugins/query_enhancements/public/index.scss create mode 100644 src/plugins/query_enhancements/public/index.ts create mode 100644 src/plugins/query_enhancements/public/plugin.tsx create mode 100644 src/plugins/query_enhancements/public/query_assist/components/__snapshots__/call_outs.test.tsx.snap create mode 100644 src/plugins/query_enhancements/public/query_assist/components/call_outs.test.tsx create mode 100644 src/plugins/query_enhancements/public/query_assist/components/call_outs.tsx create mode 100644 src/plugins/query_enhancements/public/query_assist/components/index.ts create mode 100644 src/plugins/query_enhancements/public/query_assist/components/index_selector.tsx create mode 100644 src/plugins/query_enhancements/public/query_assist/components/query_assist_banner.test.tsx create mode 100644 src/plugins/query_enhancements/public/query_assist/components/query_assist_banner.tsx create mode 100644 src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.tsx create mode 100644 src/plugins/query_enhancements/public/query_assist/components/query_assist_input.test.tsx create mode 100644 src/plugins/query_enhancements/public/query_assist/components/query_assist_input.tsx create mode 100644 src/plugins/query_enhancements/public/query_assist/components/submit_button.test.tsx create mode 100644 src/plugins/query_enhancements/public/query_assist/components/submit_button.tsx create mode 100644 src/plugins/query_enhancements/public/query_assist/hooks/index.ts create mode 100644 src/plugins/query_enhancements/public/query_assist/hooks/use_generate.test.ts create mode 100644 src/plugins/query_enhancements/public/query_assist/hooks/use_generate.ts create mode 100644 src/plugins/query_enhancements/public/query_assist/hooks/use_indices.ts create mode 100644 src/plugins/query_enhancements/public/query_assist/index.ts create mode 100644 src/plugins/query_enhancements/public/query_assist/utils/create_extension.test.tsx create mode 100644 src/plugins/query_enhancements/public/query_assist/utils/create_extension.tsx create mode 100644 src/plugins/query_enhancements/public/query_assist/utils/errors.test.ts create mode 100644 src/plugins/query_enhancements/public/query_assist/utils/errors.ts create mode 100644 src/plugins/query_enhancements/public/query_assist/utils/get_persisted_log.ts create mode 100644 src/plugins/query_enhancements/public/query_assist/utils/index.ts create mode 100644 src/plugins/query_enhancements/public/search/index.ts create mode 100644 src/plugins/query_enhancements/public/search/ppl_search_interceptor.ts create mode 100644 src/plugins/query_enhancements/public/search/sql_async_search_interceptor.ts create mode 100644 src/plugins/query_enhancements/public/search/sql_search_interceptor.ts create mode 100644 src/plugins/query_enhancements/public/services.ts create mode 100644 src/plugins/query_enhancements/public/types.ts create mode 100644 src/plugins/query_enhancements/server/index.ts create mode 100644 src/plugins/query_enhancements/server/plugin.ts create mode 100644 src/plugins/query_enhancements/server/routes/data_source_connection/index.ts create mode 100644 src/plugins/query_enhancements/server/routes/data_source_connection/routes.ts create mode 100644 src/plugins/query_enhancements/server/routes/index.ts create mode 100644 src/plugins/query_enhancements/server/routes/query_assist/agents.test.ts create mode 100644 src/plugins/query_enhancements/server/routes/query_assist/agents.ts create mode 100644 src/plugins/query_enhancements/server/routes/query_assist/createResponse.ts create mode 100644 src/plugins/query_enhancements/server/routes/query_assist/index.ts create mode 100644 src/plugins/query_enhancements/server/routes/query_assist/ppl/create_response.ts create mode 100644 src/plugins/query_enhancements/server/routes/query_assist/routes.ts create mode 100644 src/plugins/query_enhancements/server/search/index.ts create mode 100644 src/plugins/query_enhancements/server/search/ppl_search_strategy.ts create mode 100644 src/plugins/query_enhancements/server/search/sql_async_search_strategy.ts create mode 100644 src/plugins/query_enhancements/server/search/sql_search_strategy.test.ts create mode 100644 src/plugins/query_enhancements/server/search/sql_search_strategy.ts create mode 100644 src/plugins/query_enhancements/server/types.ts create mode 100644 src/plugins/query_enhancements/server/utils/facet.ts create mode 100644 src/plugins/query_enhancements/server/utils/index.ts create mode 100644 src/plugins/query_enhancements/server/utils/plugins.ts create mode 100644 src/plugins/query_enhancements/server/utils/shim_schema_row.ts create mode 100644 src/plugins/query_enhancements/server/utils/shim_stats.ts create mode 100644 src/plugins/query_enhancements/test/__mocks__/fileMock.js create mode 100644 src/plugins/query_enhancements/test/__mocks__/styleMock.js create mode 100644 src/plugins/query_enhancements/test/jest.config.js create mode 100644 src/plugins/query_enhancements/test/setup.jest.ts create mode 100644 src/plugins/query_enhancements/test/setupTests.ts diff --git a/changelogs/fragments/7212.yml b/changelogs/fragments/7212.yml new file mode 100644 index 000000000000..597428db4724 --- /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 66a259b8c2f4..e5136a6df4c5 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 b82001018886..f1ac419e9ec1 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 582470cc997c..5483b540d5bf 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 bddef49af1d4..52fba0ecea78 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 19c4c527038c..59d5645fcf6c 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 f406423d616e..8dcaaa459dd9 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 000000000000..bb9a3ef5e506 --- /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 000000000000..f7c6326fd095 --- /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 000000000000..b9ea4750e601 --- /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 000000000000..7e82677407de --- /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 000000000000..f21e3c24507f --- /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 000000000000..9469a3a2771e --- /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 000000000000..057ff1e708d8 --- /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 000000000000..c98ca0284969 --- /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 000000000000..f4bdde2a26e1 --- /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 000000000000..e6ed7e2e0a17 --- /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 000000000000..b744e8c35e8f --- /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 000000000000..d7590b278220 --- /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 000000000000..1ee969a1d079 --- /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 000000000000..e334163d91d4 --- /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 000000000000..6afec4b51a99 --- /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 000000000000..08eeda5a7aa1 --- /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 000000000000..e5822c4b378e --- /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 000000000000..9eccc9e6f35a --- /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 000000000000..ff7112406eac --- /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 000000000000..8f141c20e69c --- /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 000000000000..0ea557db8ce2 --- /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 000000000000..cfe754d3ae0c --- /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 000000000000..4ac107d49c6b --- /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 000000000000..16d218bed67e --- /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 000000000000..6301c474eeb2 --- /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 000000000000..4e591e3401c1 --- /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 000000000000..03655e0e266e --- /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 000000000000..68faac461a6b --- /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 000000000000..84ac854c980f --- /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 000000000000..1e6eb89da1f8 --- /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 000000000000..5ea5caaa5f49 --- /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 000000000000..73d359e97397 --- /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 000000000000..52d2c5a63019 --- /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 000000000000..a2076151efb3 --- /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 000000000000..a15035b5b10f --- /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 000000000000..091e9dd7c295 --- /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 000000000000..1c985463877a --- /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 000000000000..51d08717b228 --- /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 000000000000..ea568959152d --- /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 000000000000..bd990b57f5a7 --- /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 000000000000..1adbae11b462 --- /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 000000000000..dbd0a9ef8fc7 --- /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 000000000000..c86edcf125c0 --- /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 000000000000..517be64e9e90 --- /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 000000000000..9835c1345f02 --- /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 000000000000..aac50de3bb98 --- /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 000000000000..9232ef146cdb --- /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 000000000000..c4dd7409faf6 --- /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 000000000000..d11233be2dca --- /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 000000000000..c31da11e6b15 --- /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 000000000000..4d72ff3eb278 --- /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 000000000000..db105849694e --- /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 000000000000..0e6c700b4d47 --- /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 000000000000..f4fe42779dae --- /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 000000000000..8feaecc7b282 --- /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 000000000000..f22bdb148368 --- /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 000000000000..d065254c7f05 --- /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 000000000000..c88cd678a971 --- /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 000000000000..1e28ddc34b2a --- /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 000000000000..6519eb00ac1a --- /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 000000000000..d30396719533 --- /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 000000000000..b3a528c57055 --- /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 000000000000..7f2af4d4182a --- /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 000000000000..acd0027d0bc1 --- /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 000000000000..71fea93c2113 --- /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 000000000000..c5ebb40f882b --- /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 000000000000..2ab716b98928 --- /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 000000000000..a6f23efba2a1 --- /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 000000000000..86f216973771 --- /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 000000000000..e8d581bc42f2 --- /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 000000000000..a38d77553e7b --- /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 000000000000..e616996e7509 --- /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 000000000000..03b436dba6cc --- /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 000000000000..28de3c8b1ecb --- /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 000000000000..f5a1419f3ba8 --- /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 000000000000..43627882b00d --- /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 000000000000..5a996f6f8c98 --- /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');