From cdbf8f394abd31e888fd33af504678fc030a507a Mon Sep 17 00:00:00 2001
From: "JUST.in DO IT" <justin.park@airbnb.com>
Date: Fri, 26 Apr 2024 09:48:01 -0700
Subject: [PATCH] chore: Add custom keywords for SQL Lab autocomplete (#28153)

---
 .../src/ui-overrides/types.ts                 | 24 +++++++++++
 .../AceEditorWrapper/useKeywords.test.ts      | 41 +++++++++++++++++++
 .../AceEditorWrapper/useKeywords.ts           | 24 +++++++++--
 3 files changed, 86 insertions(+), 3 deletions(-)

diff --git a/superset-frontend/packages/superset-ui-core/src/ui-overrides/types.ts b/superset-frontend/packages/superset-ui-core/src/ui-overrides/types.ts
index cf03650e6..f8f10982b 100644
--- a/superset-frontend/packages/superset-ui-core/src/ui-overrides/types.ts
+++ b/superset-frontend/packages/superset-ui-core/src/ui-overrides/types.ts
@@ -17,6 +17,7 @@
  * under the License.
  */
 import React, { ReactNode, MouseEventHandler } from 'react';
+import type { Editor } from 'brace';
 
 /**
  * A function which returns text (or marked-up text)
@@ -164,6 +165,26 @@ export interface SubMenuProps {
   activeChild?: string;
 }
 
+export interface CustomAutoCompleteArgs {
+  queryEditorId: string;
+  dbId?: string | number;
+  schema?: string;
+}
+
+interface AutocompleteItem {
+  name: string;
+  value: string;
+  score: number;
+  meta: string;
+  label?: string;
+  docHTML?: string;
+  docText?: string;
+}
+
+export interface CustomAutocomplete extends AutocompleteItem {
+  insertMatch?: (editor: Editor, data: AutocompleteItem) => void;
+}
+
 export type Extensions = Partial<{
   'alertsreports.header.icon': React.ComponentType;
   'embedded.documentation.configuration_details': React.ComponentType<ConfigDetailsProps>;
@@ -187,4 +208,7 @@ export type Extensions = Partial<{
   'sqleditor.extension.form': React.ComponentType<SQLFormExtensionProps>;
   'sqleditor.extension.resultTable': React.ComponentType<SQLResultTableExtentionProps>;
   'dashboard.slice.header': React.ComponentType<SliceHeaderExtension>;
+  'sqleditor.extension.customAutocomplete': (
+    args: CustomAutoCompleteArgs,
+  ) => CustomAutocomplete[] | undefined;
 }>;
diff --git a/superset-frontend/src/SqlLab/components/AceEditorWrapper/useKeywords.test.ts b/superset-frontend/src/SqlLab/components/AceEditorWrapper/useKeywords.test.ts
index 7aa306d8b..adee7d681 100644
--- a/superset-frontend/src/SqlLab/components/AceEditorWrapper/useKeywords.test.ts
+++ b/superset-frontend/src/SqlLab/components/AceEditorWrapper/useKeywords.test.ts
@@ -18,6 +18,7 @@
  */
 import fetchMock from 'fetch-mock';
 import { act, renderHook } from '@testing-library/react-hooks';
+import { getExtensionsRegistry } from '@superset-ui/core';
 import {
   createWrapper,
   defaultStore as store,
@@ -313,3 +314,43 @@ test('returns long keywords with docText', async () => {
     ),
   );
 });
+
+test('Add custom keywords for autocomplete', () => {
+  const expected = [
+    {
+      name: 'Custom keyword 1',
+      label: 'Custom keyword 1',
+      meta: 'Custom',
+      value: 'custom1',
+      score: 100,
+    },
+    {
+      name: 'Custom keyword 2',
+      label: 'Custom keyword 2',
+      meta: 'Custom',
+      value: 'custom2',
+      score: 50,
+    },
+  ];
+  const extensionsRegistry = getExtensionsRegistry();
+  extensionsRegistry.set(
+    'sqleditor.extension.customAutocomplete',
+    () => expected,
+  );
+  const { result } = renderHook(
+    () =>
+      useKeywords({
+        queryEditorId: 'testqueryid',
+        dbId: expectDbId,
+        schema: expectSchema,
+      }),
+    {
+      wrapper: createWrapper({
+        useRedux: true,
+        store,
+      }),
+    },
+  );
+  expect(result.current).toContainEqual(expect.objectContaining(expected[0]));
+  expect(result.current).toContainEqual(expect.objectContaining(expected[1]));
+});
diff --git a/superset-frontend/src/SqlLab/components/AceEditorWrapper/useKeywords.ts b/superset-frontend/src/SqlLab/components/AceEditorWrapper/useKeywords.ts
index b42edc582..09a503526 100644
--- a/superset-frontend/src/SqlLab/components/AceEditorWrapper/useKeywords.ts
+++ b/superset-frontend/src/SqlLab/components/AceEditorWrapper/useKeywords.ts
@@ -18,7 +18,7 @@
  */
 import { useEffect, useMemo, useRef } from 'react';
 import { useSelector, useDispatch, shallowEqual, useStore } from 'react-redux';
-import { t } from '@superset-ui/core';
+import { getExtensionsRegistry, t } from '@superset-ui/core';
 
 import { Editor } from 'src/components/AsyncAceEditor';
 import sqlKeywords from 'src/SqlLab/utils/sqlKeywords';
@@ -55,10 +55,21 @@ const getHelperText = (value: string) =>
     docText: value,
   };
 
+const extensionsRegistry = getExtensionsRegistry();
+
 export function useKeywords(
   { queryEditorId, dbId, schema }: Params,
   skip = false,
 ) {
+  const useCustomKeywords = extensionsRegistry.get(
+    'sqleditor.extension.customAutocomplete',
+  );
+
+  const customKeywords = useCustomKeywords?.({
+    queryEditorId: String(queryEditorId),
+    dbId,
+    schema,
+  });
   const dispatch = useDispatch();
   const hasFetchedKeywords = useRef(false);
   // skipFetch is used to prevent re-evaluating memoized keywords
@@ -207,8 +218,15 @@ export function useKeywords(
         .concat(schemaKeywords)
         .concat(tableKeywords)
         .concat(functionKeywords)
-        .concat(sqlKeywords),
-    [schemaKeywords, tableKeywords, columnKeywords, functionKeywords],
+        .concat(sqlKeywords)
+        .concat(customKeywords ?? []),
+    [
+      schemaKeywords,
+      tableKeywords,
+      columnKeywords,
+      functionKeywords,
+      customKeywords,
+    ],
   );
 
   hasFetchedKeywords.current = !skip;