From 27a74abf44d3a4c80cc84137f30fbb9766836449 Mon Sep 17 00:00:00 2001
From: Paul Sebastian <paulstn@amazon.com>
Date: Thu, 25 Jul 2024 16:41:44 -0700
Subject: [PATCH] [Auto Suggest] DQL Updates (#7498)

* update code completion to not return for visualize

Signed-off-by: Paul Sebastian <paulstn@amazon.com>

* update types to match completionitemkind

Signed-off-by: Paul Sebastian <paulstn@amazon.com>

---------

Signed-off-by: Paul Sebastian <paulstn@amazon.com>
---
 .../data/public/antlr/dql/code_completion.ts  | 28 +++++++++++++------
 .../providers/query_suggestion_provider.ts    |  2 +-
 .../public/ui/query_editor/query_editor.tsx   |  2 +-
 3 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/src/plugins/data/public/antlr/dql/code_completion.ts b/src/plugins/data/public/antlr/dql/code_completion.ts
index 37ff62e2fdca..6c7bbac6d9e5 100644
--- a/src/plugins/data/public/antlr/dql/code_completion.ts
+++ b/src/plugins/data/public/antlr/dql/code_completion.ts
@@ -6,12 +6,14 @@
 import { CharStream, CommonTokenStream, TokenStream } from 'antlr4ng';
 import { CodeCompletionCore } from 'antlr4-c3';
 import { HttpSetup } from 'opensearch-dashboards/public';
+import { monaco } from '@osd/monaco';
 import { DQLLexer } from './.generated/DQLLexer';
 import { DQLParser, KeyValueExpressionContext } from './.generated/DQLParser';
 import { getTokenPosition } from '../shared/cursor';
-import { IndexPattern, IndexPatternField } from '../../index_patterns';
+import { IndexPatternField } from '../../index_patterns';
 import { QuerySuggestionGetFnArgs } from '../../autocomplete';
 import { DQLParserVisitor } from './.generated/DQLParserVisitor';
+import { getUiService } from '../../services';
 
 const findCursorIndex = (
   tokenStream: TokenStream,
@@ -43,11 +45,12 @@ const findFieldSuggestions = (indexPattern: any) => {
       return idxField.name;
     });
 
-  const fieldSuggestions: Array<{ text: string; type: string }> = fieldNames.map(
-    (field: string) => {
-      return { text: field, type: 'field' };
-    }
-  );
+  const fieldSuggestions: Array<{
+    text: string;
+    type: monaco.languages.CompletionItemKind;
+  }> = fieldNames.map((field: string) => {
+    return { text: field, type: monaco.languages.CompletionItemKind.Field };
+  });
 
   return fieldSuggestions;
 };
@@ -116,7 +119,16 @@ export const getSuggestions = async ({
   position,
   selectionEnd,
   core: coreSetup,
+  services,
 }: QuerySuggestionGetFnArgs) => {
+  if (
+    !services ||
+    !services.appName ||
+    !getUiService().Settings.supportsEnhancementsEnabled(services.appName)
+  ) {
+    return [];
+  }
+
   const http = coreSetup?.http;
   const currentIndexPattern = indexPatterns[0];
 
@@ -173,7 +185,7 @@ export const getSuggestions = async ({
     if (!!values) {
       completions.push(
         ...values?.map((val: any) => {
-          return { text: val, type: 'value' };
+          return { text: val, type: monaco.languages.CompletionItemKind.Value };
         })
       );
     }
@@ -187,7 +199,7 @@ export const getSuggestions = async ({
     }
 
     const tokenSymbolName = parser.vocabulary.getSymbolicName(token)?.toLowerCase();
-    completions.push({ text: tokenSymbolName, type: 'keyword' });
+    completions.push({ text: tokenSymbolName, type: monaco.languages.CompletionItemKind.Keyword });
   });
 
   return completions;
diff --git a/src/plugins/data/public/autocomplete/providers/query_suggestion_provider.ts b/src/plugins/data/public/autocomplete/providers/query_suggestion_provider.ts
index 0825e5b69211..4cc6d7da59a3 100644
--- a/src/plugins/data/public/autocomplete/providers/query_suggestion_provider.ts
+++ b/src/plugins/data/public/autocomplete/providers/query_suggestion_provider.ts
@@ -54,7 +54,7 @@ export interface QuerySuggestionGetFnArgs {
   signal?: AbortSignal;
   boolFilter?: any;
   position?: monaco.Position;
-  connectionService?: any; // will need to add type when ConnectionService is properly exposed from queryEnhancements
+  services?: any; // will need to add type when ConnectionService is properly exposed from queryEnhancements
   core?: CoreSetup;
 }
 
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 f4ef7c88035f..1f8086202348 100644
--- a/src/plugins/data/public/ui/query_editor/query_editor.tsx
+++ b/src/plugins/data/public/ui/query_editor/query_editor.tsx
@@ -316,7 +316,7 @@ export default class QueryEditorUI extends Component<Props, State> {
       language: this.props.query.language,
       indexPatterns,
       position,
-      openSearchDashboards: this.services,
+      services: this.services,
     });
 
     return {