take in dataset type to only send agg queries for index patterns and …

…indexes Signed-off-by: Paul Sebastian <[email protected]>
opensearch-project · Jan 9, 2025 · c4023ad · c4023ad
1 parent 524212c
commit c4023ad
Show file tree

Hide file tree

Showing 7 changed files with 111 additions and 30 deletions.
diff --git a/src/plugins/data/public/antlr/opensearch_ppl/code_completion.ts b/src/plugins/data/public/antlr/opensearch_ppl/code_completion.ts
@@ -26,6 +26,7 @@ export const getSuggestions = async ({
   selectionStart,
   selectionEnd,
   indexPattern,
+  datasetType,
   position,
   query,
   services,
@@ -51,7 +52,8 @@ export const getSuggestions = async ({
             indexPattern.title,
             suggestions.suggestValuesForColumn,
             services,
-            indexPattern.fields.find((field) => field.name === suggestions.suggestValuesForColumn)
+            indexPattern.fields.find((field) => field.name === suggestions.suggestValuesForColumn),
+            datasetType
           ).catch(() => []),
           (val: any) => (typeof val === 'string' ? `"${val}" ` : `${val} `)
         )

diff --git a/src/plugins/data/public/antlr/opensearch_sql/code_completion.ts b/src/plugins/data/public/antlr/opensearch_sql/code_completion.ts
@@ -35,6 +35,7 @@ export const getSuggestions = async ({
   selectionStart,
   selectionEnd,
   indexPattern,
+  datasetType,
   position,
   query,
   services,
@@ -109,7 +110,8 @@ export const getSuggestions = async ({
                   services,
                   indexPattern.fields.find(
                     (field) => field.name === suggestions.suggestValuesForColumn
-                  )
+                  ),
+                  datasetType
                 ).catch(() => []),
                 (val: any) => (typeof val === 'string' ? `'${val}' ` : `${val} `)
               )

diff --git a/src/plugins/data/public/antlr/opensearch_sql/opensearch_sql_autocomplete.ts b/src/plugins/data/public/antlr/opensearch_sql/opensearch_sql_autocomplete.ts
@@ -84,7 +84,6 @@ const rulesToVisit = new Set([
   OpenSearchSQLParser.RULE_predicate,
 ]);
 
-// TODO: double check symbol table
 class OpenSearchSqlSymbolTableVisitor
   extends OpenSearchSQLParserVisitor<{}>
   implements ISymbolTableVisitor {

diff --git a/src/plugins/data/public/antlr/shared/utils.test.ts b/src/plugins/data/public/antlr/shared/utils.test.ts
@@ -308,9 +308,15 @@ describe('fetchColumnValues', () => {
   });
 
   it('should return boolean values for boolean fields', async () => {
-    const result = await fetchColumnValues('test-table', 'boolean-column', mockServices, {
-      type: 'boolean',
-    } as IndexPatternField);
+    const result = await fetchColumnValues(
+      'test-table',
+      'boolean-column',
+      mockServices,
+      {
+        type: 'boolean',
+      } as IndexPatternField,
+      'INDEX_PATTERN'
+    );
 
     expect(result).toEqual(['true', 'false']);
   });
@@ -322,17 +328,29 @@ describe('fetchColumnValues', () => {
       }
     });
 
-    const result = await fetchColumnValues('test-table', 'test-column', mockServices, {
-      type: 'string',
-    } as IndexPatternField);
+    const result = await fetchColumnValues(
+      'test-table',
+      'test-column',
+      mockServices,
+      {
+        type: 'string',
+      } as IndexPatternField,
+      'INDEX_PATTERN'
+    );
 
     expect(result).toEqual([]);
   });
 
   it('should return empty array for unsupported field types', async () => {
-    const result = await fetchColumnValues('test-table', 'number-column', mockServices, {
-      type: 'number',
-    } as IndexPatternField);
+    const result = await fetchColumnValues(
+      'test-table',
+      'number-column',
+      mockServices,
+      {
+        type: 'number',
+      } as IndexPatternField,
+      'INDEX_PATTERN'
+    );
 
     expect(result).toEqual([]);
   });
@@ -350,9 +368,15 @@ describe('fetchColumnValues', () => {
 
     (mockHttp.fetch as jest.Mock).mockResolvedValue(mockResponse);
 
-    const result = await fetchColumnValues('test-table', 'string-column', mockServices, {
-      type: 'string',
-    } as IndexPatternField);
+    const result = await fetchColumnValues(
+      'test-table',
+      'string-column',
+      mockServices,
+      {
+        type: 'string',
+      } as IndexPatternField,
+      'INDEX_PATTERN'
+    );
 
     expect(mockHttp.fetch).toHaveBeenCalledWith({
       method: 'POST',
@@ -375,14 +399,66 @@ describe('fetchColumnValues', () => {
     (mockHttp.fetch as jest.Mock).mockRejectedValue(new Error('SQL API Error'));
 
     await expect(
-      fetchColumnValues('test-table', 'string-column', mockServices, {
-        type: 'string',
-      } as IndexPatternField)
+      fetchColumnValues(
+        'test-table',
+        'string-column',
+        mockServices,
+        {
+          type: 'string',
+        } as IndexPatternField,
+        'INDEX_PATTERN'
+      )
     ).rejects.toThrow('SQL API Error');
   });
 
   it('should return empty array when field is undefined', async () => {
-    const result = await fetchColumnValues('test-table', 'string-column', mockServices, undefined);
+    const result = await fetchColumnValues(
+      'test-table',
+      'string-column',
+      mockServices,
+      undefined,
+      'INDEX_PATTERN'
+    );
+
+    expect(result).toEqual([]);
+  });
+
+  it('should fetch values when datasetType is INDEXES', async () => {
+    const mockResponse = {
+      body: {
+        fields: [
+          {
+            values: ['value1', 'value2'],
+          },
+        ],
+      },
+    };
+
+    (mockHttp.fetch as jest.Mock).mockResolvedValue(mockResponse);
+
+    const result = await fetchColumnValues(
+      'test-table',
+      'string-column',
+      mockServices,
+      {
+        type: 'string',
+      } as IndexPatternField,
+      'INDEXES'
+    );
+
+    expect(result).toEqual(['value1', 'value2']);
+  });
+
+  it('should return empty array when datasetType is unsupported', async () => {
+    const result = await fetchColumnValues(
+      'test-table',
+      'string-column',
+      mockServices,
+      {
+        type: 'string',
+      } as IndexPatternField,
+      'S3'
+    );
 
     expect(result).toEqual([]);
   });

diff --git a/src/plugins/data/public/antlr/shared/utils.ts b/src/plugins/data/public/antlr/shared/utils.ts
@@ -18,7 +18,7 @@ import { ParsingSubject } from './types';
 import { quotesRegex, SuggestionItemDetailsTags } from './constants';
 import { IndexPattern, IndexPatternField } from '../../index_patterns';
 import { IDataPluginServices } from '../../types';
-import { UI_SETTINGS } from '../../../common';
+import { DEFAULT_DATA, UI_SETTINGS } from '../../../common';
 import { MonacoCompatibleQuerySuggestion } from '../../autocomplete/providers/query_suggestion_provider';
 
 export interface IDataSourceRequestHandlerParams {
@@ -100,8 +100,13 @@ export const fetchColumnValues = async (
   table: string,
   column: string,
   services: IDataPluginServices,
-  fieldInOsd: IndexPatternField | undefined
+  fieldInOsd: IndexPatternField | undefined,
+  datasetType: string | undefined
 ): Promise<any[]> => {
+  if (!datasetType || !Object.values(DEFAULT_DATA.SET_TYPES).includes(datasetType)) {
+    return [];
+  }
+
   // default to true/false values for type boolean
   if (fieldInOsd?.type === 'boolean') {
     return ['true', 'false'];

diff --git a/src/plugins/data/public/autocomplete/providers/query_suggestion_provider.ts b/src/plugins/data/public/autocomplete/providers/query_suggestion_provider.ts
@@ -48,6 +48,7 @@ export type QuerySuggestionGetFn = (
 export interface QuerySuggestionGetFnArgs {
   language: string;
   indexPattern: IndexPattern | undefined;
+  datasetType?: string;
   query: string;
   selectionStart: number;
   selectionEnd: number;

diff --git a/src/plugins/data/public/ui/query_editor/query_editor.tsx b/src/plugins/data/public/ui/query_editor/query_editor.tsx
@@ -218,24 +218,20 @@ export const QueryEditorUI: React.FC<Props> = (props) => {
     return <QueryControls queryControls={queryControls} />;
   };
 
-  const fetchIndexPattern = async () => {
-    const dataset = queryString.getQuery().dataset;
-    if (!dataset) return undefined;
-    const indexPattern = await getIndexPatterns().get(dataset.id);
-    return indexPattern;
-  };
-
   const provideCompletionItems = async (
     model: monaco.editor.ITextModel,
     position: monaco.Position
   ): Promise<monaco.languages.CompletionList> => {
-    const indexPattern = await fetchIndexPattern();
+    const dataset = queryString.getQuery().dataset;
+    const indexPattern = dataset ? await getIndexPatterns().get(dataset.id) : undefined;
+
     const suggestions = await services.data.autocomplete.getQuerySuggestions({
       query: inputRef.current?.getValue() ?? '',
-      selectionStart: model.getOffsetAt(position),
+      selectionStart: model.getOffsetAt(position), // not needed, position handles same thing. remove
       selectionEnd: model.getOffsetAt(position),
       language: queryRef.current.language,
       indexPattern,
+      datasetType: dataset?.type,
       position,
       services,
     });