From 1d9a761de5410fa1bd208bca4c78614779cf3064 Mon Sep 17 00:00:00 2001
From: Antonio Rivero <38889534+Antonio-RiveroMartnez@users.noreply.github.com>
Date: Thu, 1 Jun 2023 15:30:26 -0400
Subject: [PATCH] feat(sqllab): SQLEditor Extension (#24205)

---
 .../src/ui-overrides/types.ts                 | 15 +++++++++++++
 .../components/SqlEditor/SqlEditor.test.jsx   | 16 ++++++++++++++
 .../src/SqlLab/components/SqlEditor/index.jsx | 22 ++++++++++++++++++-
 3 files changed, 52 insertions(+), 1 deletion(-)

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 df1cd3abae868..3f0a5592973ee 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
@@ -104,6 +104,20 @@ export interface DatabaseConnectionExtension {
   onDelete?: (db: any) => void;
 }
 
+/**
+ * Interface for extensions SQL Form.
+ * These will be passed from the SQLEditor
+ *
+ * @param queryEditorId the queryEditor's id
+ * @param setQueryEditorAndSaveSqlWithDebounce Debounced function that saves a query into the query editor
+ * @param startQuery Callback that starts a query from the query editor
+ */
+export interface SQLFormExtensionProps {
+  queryEditorId: string;
+  setQueryEditorAndSaveSqlWithDebounce: (sql: string) => void;
+  startQuery: (ctasArg?: any, ctas_method?: any) => void;
+}
+
 export type Extensions = Partial<{
   'alertsreports.header.icon': React.ComponentType;
   'embedded.documentation.configuration_details': React.ComponentType<ConfigDetailsProps>;
@@ -122,4 +136,5 @@ export type Extensions = Partial<{
   /* Custom components to show in the database and dataset delete modals */
   'database.delete.related': React.ComponentType<DatabaseDeleteRelatedExtensionProps>;
   'dataset.delete.related': React.ComponentType<DatasetDeleteRelatedExtensionProps>;
+  'sqleditor.extension.form': React.ComponentType<SQLFormExtensionProps>;
 }>;
diff --git a/superset-frontend/src/SqlLab/components/SqlEditor/SqlEditor.test.jsx b/superset-frontend/src/SqlLab/components/SqlEditor/SqlEditor.test.jsx
index d8a1aa260a7d6..464dbf9e8ca8b 100644
--- a/superset-frontend/src/SqlLab/components/SqlEditor/SqlEditor.test.jsx
+++ b/superset-frontend/src/SqlLab/components/SqlEditor/SqlEditor.test.jsx
@@ -31,6 +31,8 @@ import {
 } from 'src/SqlLab/fixtures';
 import SqlEditorLeftBar from 'src/SqlLab/components/SqlEditorLeftBar';
 import { api } from 'src/hooks/apiResources/queryApi';
+import { getExtensionsRegistry } from '@superset-ui/core';
+import setupExtensions from 'src/setup/setupExtensions';
 
 jest.mock('src/components/AsyncAceEditor', () => ({
   ...jest.requireActual('src/components/AsyncAceEditor'),
@@ -246,4 +248,18 @@ describe('SqlEditor', () => {
     fireEvent.click(await findByText('LIMIT:'));
     expect(await findByText('10 000')).toBeInTheDocument();
   });
+
+  it('renders an Extension if provided', async () => {
+    const extensionsRegistry = getExtensionsRegistry();
+
+    extensionsRegistry.set('sqleditor.extension.form', () => (
+      <>sqleditor.extension.form extension component</>
+    ));
+
+    setupExtensions();
+    const { findByText } = setup(mockedProps, store);
+    expect(
+      await findByText('sqleditor.extension.form extension component'),
+    ).toBeInTheDocument();
+  });
 });
diff --git a/superset-frontend/src/SqlLab/components/SqlEditor/index.jsx b/superset-frontend/src/SqlLab/components/SqlEditor/index.jsx
index 19464061f2046..0216ed2f5e29f 100644
--- a/superset-frontend/src/SqlLab/components/SqlEditor/index.jsx
+++ b/superset-frontend/src/SqlLab/components/SqlEditor/index.jsx
@@ -30,7 +30,14 @@ import { CSSTransition } from 'react-transition-group';
 import { shallowEqual, useDispatch, useSelector } from 'react-redux';
 import PropTypes from 'prop-types';
 import Split from 'react-split';
-import { css, FeatureFlag, styled, t, useTheme } from '@superset-ui/core';
+import {
+  css,
+  FeatureFlag,
+  styled,
+  t,
+  useTheme,
+  getExtensionsRegistry,
+} from '@superset-ui/core';
 import debounce from 'lodash/debounce';
 import throttle from 'lodash/throttle';
 import Modal from 'src/components/Modal';
@@ -205,6 +212,8 @@ const propTypes = {
   scheduleQueryWarning: PropTypes.string,
 };
 
+const extensionsRegistry = getExtensionsRegistry();
+
 const SqlEditor = ({
   tables,
   queryEditor,
@@ -253,6 +262,8 @@ const SqlEditor = ({
   const sqlEditorRef = useRef(null);
   const northPaneRef = useRef(null);
 
+  const SqlFormExtension = extensionsRegistry.get('sqleditor.extension.form');
+
   const startQuery = useCallback(
     (ctasArg = false, ctas_method = CtasEnum.TABLE) => {
       if (!database) {
@@ -673,6 +684,15 @@ const SqlEditor = ({
         onDragEnd={onResizeEnd}
       >
         <div ref={northPaneRef} className="north-pane">
+          {SqlFormExtension && (
+            <SqlFormExtension
+              queryEditorId={queryEditor.id}
+              setQueryEditorAndSaveSqlWithDebounce={
+                setQueryEditorAndSaveSqlWithDebounce
+              }
+              startQuery={startQuery}
+            />
+          )}
           <AceEditorWrapper
             autocomplete={autocompleteEnabled}
             onBlur={setQueryEditorAndSaveSql}