Skip to content

Commit

Permalink
[query assist] use badge to show agent errors (#7998) (#8146)
Browse files Browse the repository at this point in the history
- enable query assist for index patterns without data source
- add agent error parsing utils
- update ml-commons response schema processing
   - previously ml-commons returns error.body as a string, now it is a JSON object. Ideally frontend should keep it as is to reduce serializing/deserializing, but since older version of ml-commons can be used through MDS, we'll keep it as a string for consistency
- use badge to show query assist error if possible
- add unit tests


(cherry picked from commit 461a395)

Signed-off-by: Joshua Li <[email protected]>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
  • Loading branch information
1 parent 94c41cc commit 5cce5a0
Show file tree
Hide file tree
Showing 18 changed files with 613 additions and 31 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
* under the License.
*/

import { of } from 'rxjs';
import { QueryStringContract } from '.';
import { Query, Dataset } from '../../../common';
import { datasetServiceMock } from './dataset_service/dataset_service.mock';
Expand All @@ -44,10 +45,10 @@ const createSetupContractMock = (isEnhancementsEnabled: boolean = false) => {
};

const queryStringManagerMock: jest.Mocked<QueryStringContract> = {
getQuery: jest.fn(),
getQuery: jest.fn().mockReturnValue(defaultQuery),
setQuery: jest.fn(),
getUpdates$: jest.fn(),
getDefaultQuery: jest.fn(),
getUpdates$: jest.fn().mockReturnValue(of(defaultQuery)),
getDefaultQuery: jest.fn().mockReturnValue(defaultQuery),
formatQuery: jest.fn(),
clearQuery: jest.fn(),
addToQueryHistory: jest.fn(),
Expand Down
6 changes: 3 additions & 3 deletions src/plugins/query_enhancements/public/plugin.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
*/

import { CoreSetup, CoreStart, Plugin, PluginInitializerContext } from '../../../core/public';
import { IStorageWrapper, Storage } from '../../opensearch_dashboards_utils/public';
import { ConfigSchema } from '../common/config';
import { setData, setStorage } from './services';
import { createQueryAssistExtension } from './query_assist';
Expand All @@ -19,6 +18,7 @@ import { LanguageConfig, Query } from '../../data/public';
import { s3TypeConfig } from './datasets';
import { createEditor, DefaultInput, SingleLineInput } from '../../data/public';
import { QueryLanguageReference } from './query_editor/query_language_reference';
import { DataStorage } from '../../data/common';

export class QueryEnhancementsPlugin
implements
Expand All @@ -28,12 +28,12 @@ export class QueryEnhancementsPlugin
QueryEnhancementsPluginSetupDependencies,
QueryEnhancementsPluginStartDependencies
> {
private readonly storage: IStorageWrapper;
private readonly storage: DataStorage;
private readonly config: ConfigSchema;

constructor(initializerContext: PluginInitializerContext) {
this.config = initializerContext.config.get<ConfigSchema>();
this.storage = new Storage(window.localStorage);
this.storage = new DataStorage(window.localStorage, 'discover.queryAssist.');
}

public setup(
Expand Down
25 changes: 25 additions & 0 deletions src/plugins/query_enhancements/public/query_assist/_index.scss
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,31 @@
&.queryAssist__callout {
margin-top: $euiSizeXS;
}

.queryAssist__popoverWrapper {
display: flex;
align-items: center;
padding-right: $euiSizeS;
}

// Hide the error badge when the input is focused
.queryAssist__inputWrapper:has(.queryAssist__input:focus) .queryAssist__popoverWrapper {
display: none;
}

.queryAssist__input {
text-overflow: ellipsis;
}
}

.queryAssist__popoverText {
padding: 0 $euiSizeM;
}

#queryAssistErrorTitle {
// there is a medium size icon with two extra-small gutter between the icon
// and the title pushing to the left so title is aligned with other content
margin-left: -($euiSizeM + $euiSizeS);
}

.osdQueryEditorExtensionComponent__query-assist {
Expand Down

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,13 @@
*/

import { render } from '@testing-library/react';
import React, { ComponentProps } from 'react';
import React, { ComponentProps, PropsWithChildren } from 'react';
import { IntlProvider } from 'react-intl';
import { QueryAssistCallOut } from './call_outs';

type Props = ComponentProps<typeof QueryAssistCallOut>;

const IntlWrapper = ({ children }: { children: unknown }) => (
const IntlWrapper = ({ children }: PropsWithChildren<unknown>) => (
<IntlProvider locale="en">{children}</IntlProvider>
);

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,218 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

import { fireEvent, render, screen, waitFor } from '@testing-library/react';
import React, { ComponentProps, PropsWithChildren } from 'react';
import { IntlProvider } from 'react-intl';
import { of } from 'rxjs';
import { QueryAssistBar } from '.';
import { notificationServiceMock, uiSettingsServiceMock } from '../../../../../core/public/mocks';
import { DataStorage } from '../../../../data/common';
import { QueryEditorExtensionDependencies, QueryStringContract } from '../../../../data/public';
import { dataPluginMock } from '../../../../data/public/mocks';
import { useOpenSearchDashboards } from '../../../../opensearch_dashboards_react/public';
import { setData, setStorage } from '../../services';
import { useGenerateQuery } from '../hooks';
import { AgentError, ProhibitedQueryError } from '../utils';
import { QueryAssistInput } from './query_assist_input';

jest.mock('../../../../opensearch_dashboards_react/public', () => ({
useOpenSearchDashboards: jest.fn(),
withOpenSearchDashboards: jest.fn((component: React.Component) => component),
}));

jest.mock('../hooks', () => ({
useGenerateQuery: jest.fn().mockReturnValue({ generateQuery: jest.fn(), loading: false }),
}));

jest.mock('./query_assist_input', () => ({
QueryAssistInput: ({ inputRef, error }: ComponentProps<typeof QueryAssistInput>) => (
<>
<input ref={inputRef} />
<div>{JSON.stringify(error)}</div>
</>
),
}));

const dataMock = dataPluginMock.createStartContract(true);
const queryStringMock = dataMock.query.queryString as jest.Mocked<QueryStringContract>;
const uiSettingsMock = uiSettingsServiceMock.createStartContract();
const notificationsMock = notificationServiceMock.createStartContract();

setData(dataMock);
setStorage(new DataStorage(window.localStorage, 'mock-prefix'));

const dependencies: QueryEditorExtensionDependencies = {
language: 'PPL',
onSelectLanguage: jest.fn(),
isCollapsed: false,
setIsCollapsed: jest.fn(),
};

type Props = ComponentProps<typeof QueryAssistBar>;

const IntlWrapper = ({ children }: PropsWithChildren<unknown>) => (
<IntlProvider locale="en">{children}</IntlProvider>
);

const renderQueryAssistBar = (overrideProps: Partial<Props> = {}) => {
const props: Props = Object.assign<Props, Partial<Props>>({ dependencies }, overrideProps);
const component = render(<QueryAssistBar {...props} />, {
wrapper: IntlWrapper,
});
return { component, props: props as jest.MockedObjectDeep<Props> };
};

describe('QueryAssistBar', () => {
beforeEach(() => {
(useOpenSearchDashboards as jest.Mock).mockReturnValue({
services: {
data: dataMock,
uiSettings: uiSettingsMock,
notifications: notificationsMock,
},
});
});
afterEach(() => {
jest.clearAllMocks();
});

it('renders null if collapsed', () => {
const { component } = renderQueryAssistBar({
dependencies: { ...dependencies, isCollapsed: true },
});
expect(component.container).toBeEmptyDOMElement();
});

it('matches snapshot', () => {
const { component } = renderQueryAssistBar();
expect(component.container).toMatchSnapshot();
});

it('displays callout when query input is empty on submit', async () => {
renderQueryAssistBar();

fireEvent.click(screen.getByRole('button'));

await waitFor(() => {
expect(screen.getByTestId('query-assist-empty-query-callout')).toBeInTheDocument();
});
});

it('displays callout when dataset is not selected on submit', async () => {
queryStringMock.getQuery.mockReturnValueOnce({ query: '', language: 'kuery' });
queryStringMock.getUpdates$.mockReturnValueOnce(of({ query: '', language: 'kuery' }));
renderQueryAssistBar();

fireEvent.change(screen.getByRole('textbox'), { target: { value: 'test query' } });
fireEvent.click(screen.getByRole('button'));

await waitFor(() => {
expect(screen.getByTestId('query-assist-empty-index-callout')).toBeInTheDocument();
});
});

it('displays callout for guardrail errors', async () => {
const generateQueryMock = jest.fn().mockResolvedValue({ error: new ProhibitedQueryError() });
(useGenerateQuery as jest.Mock).mockReturnValue({
generateQuery: generateQueryMock,
loading: false,
});

renderQueryAssistBar();

fireEvent.change(screen.getByRole('textbox'), { target: { value: 'test query' } });
fireEvent.click(screen.getByRole('button'));

await waitFor(() => {
expect(screen.getByTestId('query-assist-guard-callout')).toBeInTheDocument();
});
});

it('passes agent errors to input', async () => {
const generateQueryMock = jest.fn().mockResolvedValue({
error: new AgentError({
error: { type: 'mock-type', reason: 'mock-reason', details: 'mock-details' },
status: 303,
}),
});
(useGenerateQuery as jest.Mock).mockReturnValue({
generateQuery: generateQueryMock,
loading: false,
});

renderQueryAssistBar();

fireEvent.change(screen.getByRole('textbox'), { target: { value: 'test query' } });
fireEvent.click(screen.getByRole('button'));

await waitFor(() => {
expect(screen.getByText(/mock-reason/)).toBeInTheDocument();
});
});

it('displays toast for other unknown errors', async () => {
const mockError = new Error('mock-error');
const generateQueryMock = jest.fn().mockResolvedValue({
error: mockError,
});
(useGenerateQuery as jest.Mock).mockReturnValue({
generateQuery: generateQueryMock,
loading: false,
});

renderQueryAssistBar();

fireEvent.change(screen.getByRole('textbox'), { target: { value: 'test query' } });
fireEvent.click(screen.getByRole('button'));

await waitFor(() => {
expect(notificationsMock.toasts.addError).toHaveBeenCalledWith(mockError, {
title: 'Failed to generate results',
});
});
});

it('submits a valid query and updates services', async () => {
const generateQueryMock = jest
.fn()
.mockResolvedValue({ response: { query: 'generated query' } });
(useGenerateQuery as jest.Mock).mockReturnValue({
generateQuery: generateQueryMock,
loading: false,
});

renderQueryAssistBar();

fireEvent.change(screen.getByRole('textbox'), { target: { value: 'test query' } });
fireEvent.click(screen.getByRole('button'));

await waitFor(() => {
expect(generateQueryMock).toHaveBeenCalledWith({
question: 'test query',
index: 'Default Index Pattern',
language: 'PPL',
dataSourceId: 'mock-data-source-id',
});
});

expect(queryStringMock.setQuery).toHaveBeenCalledWith({
dataset: {
dataSource: {
id: 'mock-data-source-id',
title: 'Default Data Source',
type: 'OpenSearch',
},
id: 'default-index-pattern',
timeFieldName: '@timestamp',
title: 'Default Index Pattern',
type: 'INDEX_PATTERN',
},
language: 'PPL',
query: 'generated query',
});
expect(screen.getByTestId('query-assist-query-generated-callout')).toBeInTheDocument();
});
});
Loading

0 comments on commit 5cce5a0

Please sign in to comment.