feat(ssh_tunnel): SSH Tunnel Switch extension (#22967)

superset-frontend/packages/superset-ui-core/src/ui-overrides/ExtensionsRegistry.ts

@@ -49,6 +49,16 @@ interface MenuObjectChildProps {
   disable?: boolean;
 }
 
+export interface SwitchProps {
+  isEditMode: boolean;
+  dbFetched: any;
+  disableSSHTunnelingForEngine?: boolean;
+  useSSHTunneling: boolean;
+  setUseSSHTunneling: React.Dispatch<React.SetStateAction<boolean>>;
+  setDB: React.Dispatch<any>;
+  isSSHTunneling: boolean;
+}
+
 type ConfigDetailsProps = {
   embeddedId: string;
 };
@@ -69,6 +79,7 @@ export type Extensions = Partial<{
   'welcome.message': React.ComponentType;
   'welcome.banner': React.ComponentType;
   'welcome.main.replacement': React.ComponentType;
+  'ssh_tunnel.form.switch': React.ComponentType<SwitchProps>;
 }>;
 
 /**

superset-frontend/src/views/CRUD/data/database/DatabaseModal/SSHTunnelForm.tsx

@@ -17,16 +17,13 @@
  * under the License.
  */
 import React, { EventHandler, ChangeEvent, useState } from 'react';
-import { t, SupersetTheme, styled } from '@superset-ui/core';
-import { AntdForm, AntdSwitch, Col, Row } from 'src/components';
-import InfoTooltip from 'src/components/InfoTooltip';
+import { t, styled } from '@superset-ui/core';
+import { AntdForm, Col, Row } from 'src/components';
 import { Form, FormLabel } from 'src/components/Form';
 import { Radio } from 'src/components/Radio';
 import { Input, TextArea } from 'src/components/Input';
 import { Input as AntdInput, Tooltip } from 'antd';
 import { EyeInvisibleOutlined, EyeOutlined } from '@ant-design/icons';
-import { isEmpty } from 'lodash';
-import { infoTooltip, toggleStyle } from './styles';
 
 import { DatabaseObject } from '../types';
 import { AuthType } from '.';
@@ -54,79 +51,143 @@ const StyledInputPassword = styled(AntdInput.Password)`
 
 const SSHTunnelForm = ({
   db,
-  dbFetched,
-  isEditMode,
-  isSSHTunneling,
   onSSHTunnelParametersChange,
   setSSHTunnelLoginMethod,
-  removeSSHTunnelConfig,
 }: {
   db: DatabaseObject | null;
-  dbFetched: DatabaseObject | null;
-  isEditMode: boolean;
-  isSSHTunneling: boolean;
   onSSHTunnelParametersChange: EventHandler<
     ChangeEvent<HTMLInputElement | HTMLTextAreaElement>
   >;
   setSSHTunnelLoginMethod: (method: AuthType) => void;
-  removeSSHTunnelConfig: () => void;
 }) => {
-  const [useSSHTunneling, setUseSSHTunneling] = useState<boolean>(
-    !isEmpty(db?.ssh_tunnel),
-  );
   const [usePassword, setUsePassword] = useState<AuthType>(AuthType.password);
 
   return (
     <Form>
-      <div css={(theme: SupersetTheme) => infoTooltip(theme)}>
-        <AntdSwitch
-          disabled={
-            !isSSHTunneling || (isEditMode && !isEmpty(dbFetched?.ssh_tunnel))
-          }
-          checked={useSSHTunneling}
-          onChange={changed => {
-            setUseSSHTunneling(changed);
-            if (!changed) removeSSHTunnelConfig();
-          }}
-          data-test="ssh-tunnel-switch"
-        />
-        <span css={toggleStyle}>SSH Tunnel</span>
-        <InfoTooltip
-          tooltip={t('SSH Tunnel configuration parameters')}
-          placement="right"
-          viewBox="0 -5 24 24"
-        />
-      </div>
-      {useSSHTunneling && (
+      <StyledRow gutter={16}>
+        <Col xs={24} md={12}>
+          <StyledDiv>
+            <FormLabel htmlFor="server_address" required>
+              {t('SSH Host')}
+            </FormLabel>
+            <Input
+              name="server_address"
+              type="text"
+              placeholder={t('e.g. 127.0.0.1')}
+              value={db?.ssh_tunnel?.server_address || ''}
+              onChange={onSSHTunnelParametersChange}
+              data-test="ssh-tunnel-server_address-input"
+            />
+          </StyledDiv>
+        </Col>
+        <Col xs={24} md={12}>
+          <StyledDiv>
+            <FormLabel htmlFor="server_port" required>
+              {t('SSH Port')}
+            </FormLabel>
+            <Input
+              name="server_port"
+              type="text"
+              placeholder={t('22')}
+              value={db?.ssh_tunnel?.server_port || ''}
+              onChange={onSSHTunnelParametersChange}
+              data-test="ssh-tunnel-server_port-input"
+            />
+          </StyledDiv>
+        </Col>
+      </StyledRow>
+      <StyledRow gutter={16}>
+        <Col xs={24}>
+          <StyledDiv>
+            <FormLabel htmlFor="username" required>
+              {t('Username')}
+            </FormLabel>
+            <Input
+              name="username"
+              type="text"
+              placeholder={t('e.g. Analytics')}
+              value={db?.ssh_tunnel?.username || ''}
+              onChange={onSSHTunnelParametersChange}
+              data-test="ssh-tunnel-username-input"
+            />
+          </StyledDiv>
+        </Col>
+      </StyledRow>
+      <StyledRow gutter={16}>
+        <Col xs={24}>
+          <StyledDiv>
+            <FormLabel htmlFor="use_password" required>
+              {t('Login with')}
+            </FormLabel>
+            <StyledFormItem name="use_password" initialValue={usePassword}>
+              <Radio.Group
+                onChange={({ target: { value } }) => {
+                  setUsePassword(value);
+                  setSSHTunnelLoginMethod(value);
+                }}
+              >
+                <Radio
+                  value={AuthType.password}
+                  data-test="ssh-tunnel-use_password-radio"
+                >
+                  {t('Password')}
+                </Radio>
+                <Radio
+                  value={AuthType.privateKey}
+                  data-test="ssh-tunnel-use_private_key-radio"
+                >
+                  {t('Private Key & Password')}
+                </Radio>
+              </Radio.Group>
+            </StyledFormItem>
+          </StyledDiv>
+        </Col>
+      </StyledRow>
+      {usePassword === AuthType.password && (
+        <StyledRow gutter={16}>
+          <Col xs={24}>
+            <StyledDiv>
+              <FormLabel htmlFor="password" required>
+                {t('SSH Password')}
+              </FormLabel>
+              <StyledInputPassword
+                name="password"
+                placeholder={t('e.g. ********')}
+                value={db?.ssh_tunnel?.password || ''}
+                onChange={onSSHTunnelParametersChange}
+                data-test="ssh-tunnel-password-input"
+                iconRender={visible =>
+                  visible ? (
+                    <Tooltip title="Hide password.">
+                      <EyeInvisibleOutlined />
+                    </Tooltip>
+                  ) : (
+                    <Tooltip title="Show password.">
+                      <EyeOutlined />
+                    </Tooltip>
+                  )
+                }
+                role="textbox"
+              />
+            </StyledDiv>
+          </Col>
+        </StyledRow>
+      )}
+      {usePassword === AuthType.privateKey && (
         <>
           <StyledRow gutter={16}>
-            <Col xs={24} md={12}>
+            <Col xs={24}>
               <StyledDiv>
-                <FormLabel htmlFor="server_address" required>
-                  {t('SSH Host')}
+                <FormLabel htmlFor="private_key" required>
+                  {t('Private Key')}
                 </FormLabel>
-                <Input
-                  name="server_address"
-                  type="text"
-                  placeholder={t('e.g. 127.0.0.1')}
-                  value={db?.ssh_tunnel?.server_address || ''}
+                <TextArea
+                  name="private_key"
+                  placeholder={t('Paste Private Key here')}
+                  value={db?.ssh_tunnel?.private_key || ''}
                   onChange={onSSHTunnelParametersChange}
-                  data-test="ssh-tunnel-server_address-input"
-                />
-              </StyledDiv>
-            </Col>
-            <Col xs={24} md={12}>
-              <StyledDiv>
-                <FormLabel htmlFor="server_port" required>
-                  {t('SSH Port')}
-                </FormLabel>
-                <Input
-                  name="server_port"
-                  type="text"
-                  placeholder={t('22')}
-                  value={db?.ssh_tunnel?.server_port || ''}
-                  onChange={onSSHTunnelParametersChange}
-                  data-test="ssh-tunnel-server_port-input"
+                  data-test="ssh-tunnel-private_key-input"
+                  rows={4}
                 />
               </StyledDiv>
             </Col>
@@ -134,129 +195,31 @@ const SSHTunnelForm = ({
           <StyledRow gutter={16}>
             <Col xs={24}>
               <StyledDiv>
-                <FormLabel htmlFor="username" required>
-                  {t('Username')}
+                <FormLabel htmlFor="private_key_password" required>
+                  {t('Private Key Password')}
                 </FormLabel>
-                <Input
-                  name="username"
-                  type="text"
-                  placeholder={t('e.g. Analytics')}
-                  value={db?.ssh_tunnel?.username || ''}
+                <StyledInputPassword
+                  name="private_key_password"
+                  placeholder={t('e.g. ********')}
+                  value={db?.ssh_tunnel?.private_key_password || ''}
                   onChange={onSSHTunnelParametersChange}
-                  data-test="ssh-tunnel-username-input"
+                  data-test="ssh-tunnel-private_key_password-input"
+                  iconRender={visible =>
+                    visible ? (
+                      <Tooltip title="Hide password.">
+                        <EyeInvisibleOutlined />
+                      </Tooltip>
+                    ) : (
+                      <Tooltip title="Show password.">
+                        <EyeOutlined />
+                      </Tooltip>
+                    )
+                  }
+                  role="textbox"
                 />
               </StyledDiv>
             </Col>
           </StyledRow>
-          <StyledRow gutter={16}>
-            <Col xs={24}>
-              <StyledDiv>
-                <FormLabel htmlFor="use_password" required>
-                  {t('Login with')}
-                </FormLabel>
-                <StyledFormItem name="use_password" initialValue={usePassword}>
-                  <Radio.Group
-                    onChange={({ target: { value } }) => {
-                      setUsePassword(value);
-                      setSSHTunnelLoginMethod(value);
-                    }}
-                  >
-                    <Radio
-                      value={AuthType.password}
-                      data-test="ssh-tunnel-use_password-radio"
-                    >
-                      {t('Password')}
-                    </Radio>
-                    <Radio
-                      value={AuthType.privateKey}
-                      data-test="ssh-tunnel-use_private_key-radio"
-                    >
-                      {t('Private Key & Password')}
-                    </Radio>
-                  </Radio.Group>
-                </StyledFormItem>
-              </StyledDiv>
-            </Col>
-          </StyledRow>
-          {usePassword === AuthType.password && (
-            <StyledRow gutter={16}>
-              <Col xs={24}>
-                <StyledDiv>
-                  <FormLabel htmlFor="password" required>
-                    {t('SSH Password')}
-                  </FormLabel>
-                  <StyledInputPassword
-                    name="password"
-                    placeholder={t('e.g. ********')}
-                    value={db?.ssh_tunnel?.password || ''}
-                    onChange={onSSHTunnelParametersChange}
-                    data-test="ssh-tunnel-password-input"
-                    iconRender={visible =>
-                      visible ? (
-                        <Tooltip title={t('Hide password.')}>
-                          <EyeInvisibleOutlined />
-                        </Tooltip>
-                      ) : (
-                        <Tooltip title={t('Show password.')}>
-                          <EyeOutlined />
-                        </Tooltip>
-                      )
-                    }
-                    role="textbox"
-                  />
-                </StyledDiv>
-              </Col>
-            </StyledRow>
-          )}
-          {usePassword === AuthType.privateKey && (
-            <>
-              <StyledRow gutter={16}>
-                <Col xs={24}>
-                  <StyledDiv>
-                    <FormLabel htmlFor="private_key" required>
-                      {t('Private Key')}
-                    </FormLabel>
-                    <TextArea
-                      name="private_key"
-                      placeholder={t('Paste Private Key here')}
-                      value={db?.ssh_tunnel?.private_key || ''}
-                      onChange={onSSHTunnelParametersChange}
-                      data-test="ssh-tunnel-private_key-input"
-                      rows={4}
-                    />
-                  </StyledDiv>
-                </Col>
-              </StyledRow>
-              <StyledRow gutter={16}>
-                <Col xs={24}>
-                  <StyledDiv>
-                    <FormLabel htmlFor="private_key_password" required>
-                      {t('Private Key Password')}
-                    </FormLabel>
-                    <StyledInputPassword
-                      name="private_key_password"
-                      placeholder={t('e.g. ********')}
-                      value={db?.ssh_tunnel?.private_key_password || ''}
-                      onChange={onSSHTunnelParametersChange}
-                      data-test="ssh-tunnel-private_key_password-input"
-                      iconRender={visible =>
-                        visible ? (
-                          <Tooltip title="Hide password.">
-                            <EyeInvisibleOutlined />
-                          </Tooltip>
-                        ) : (
-                          <Tooltip title="Show password.">
-                            <EyeOutlined />
-                          </Tooltip>
-                        )
-                      }
-                      role="textbox"
-                    />
-                  </StyledDiv>
-                </Col>
-              </StyledRow>
-            </>
-          )}
         </>
       )}
     </Form>

superset-frontend/src/views/CRUD/data/database/DatabaseModal/SSHTunnelSwitch.tsx

@@ -0,0 +1,58 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import React from 'react';
+import { t, SupersetTheme, SwitchProps } from '@superset-ui/core';
+import { AntdSwitch } from 'src/components';
+import InfoTooltip from 'src/components/InfoTooltip';
+import { isEmpty } from 'lodash';
+import { ActionType } from '.';
+import { infoTooltip, toggleStyle } from './styles';
+
+const SSHTunnelSwitch = ({
+  isEditMode,
+  dbFetched,
+  useSSHTunneling,
+  setUseSSHTunneling,
+  setDB,
+  isSSHTunneling,
+}: SwitchProps) =>
+  isSSHTunneling ? (
+    <div css={(theme: SupersetTheme) => infoTooltip(theme)}>
+      <AntdSwitch
+        disabled={isEditMode && !isEmpty(dbFetched?.ssh_tunnel)}
+        checked={useSSHTunneling}
+        onChange={changed => {
+          setUseSSHTunneling(changed);
+          if (!changed) {
+            setDB({
+              type: ActionType.removeSSHTunnelConfig,
+            });
+          }
+        }}
+        data-test="ssh-tunnel-switch"
+      />
+      <span css={toggleStyle}>{t('SSH Tunnel')}</span>
+      <InfoTooltip
+        tooltip={t('SSH Tunnel configuration parameters')}
+        placement="right"
+        viewBox="0 -5 24 24"
+      />
+    </div>
+  ) : null;
+export default SSHTunnelSwitch;

superset-frontend/src/views/CRUD/data/database/DatabaseModal/index.test.tsx

@@ -31,6 +31,8 @@ import {
   DatabaseObject,
   CONFIGURATION_METHOD,
 } from 'src/views/CRUD/data/database/types';
+import { getExtensionsRegistry } from '@superset-ui/core';
+import setupExtensions from 'src/setup/setupExtensions';
 import * as hooks from 'src/views/CRUD/hooks';
 import DatabaseModal, {
   dbReducer,
@@ -1588,6 +1590,36 @@ describe('DatabaseModal', () => {
       expect(errorTitleMessage).toBeVisible();
     });
   });
+
+  describe('DatabaseModal w Extensions', () => {
+    const renderAndWait = async () => {
+      const extensionsRegistry = getExtensionsRegistry();
+
+      extensionsRegistry.set('ssh_tunnel.form.switch', () => (
+        <>ssh_tunnel.form.switch extension component</>
+      ));
+
+      setupExtensions();
+
+      const mounted = act(async () => {
+        render(<DatabaseModal {...dbProps} dbEngine="SQLite" />, {
+          useRedux: true,
+        });
+      });
+
+      return mounted;
+    };
+
+    beforeEach(async () => {
+      await renderAndWait();
+    });
+
+    test('should render an extension component if one is supplied', () => {
+      expect(
+        screen.getByText('ssh_tunnel.form.switch extension component'),
+      ).toBeInTheDocument();
+    });
+  });
 });
 
 describe('dbReducer', () => {

superset-frontend/src/views/CRUD/data/database/DatabaseModal/index.tsx

@@ -22,6 +22,7 @@ import {
   SupersetTheme,
   FeatureFlag,
   isFeatureEnabled,
+  getExtensionsRegistry,
 } from '@superset-ui/core';
 import React, {
   FunctionComponent,
@@ -63,7 +64,7 @@ import {
   ExtraJson,
 } from 'src/views/CRUD/data/database/types';
 import Loading from 'src/components/Loading';
-import { pick } from 'lodash';
+import { isEmpty, pick } from 'lodash';
 import ExtraOptions from './ExtraOptions';
 import SqlAlchemyForm from './SqlAlchemyForm';
 import DatabaseConnectionForm from './DatabaseConnectionForm';
@@ -90,6 +91,9 @@ import {
 } from './styles';
 import ModalHeader, { DOCUMENTATION_LINK } from './ModalHeader';
 import SSHTunnelForm from './SSHTunnelForm';
+import SSHTunnelSwitch from './SSHTunnelSwitch';
+
+const extensionsRegistry = getExtensionsRegistry();
 
 const DEFAULT_EXTRA = JSON.stringify({ allows_virtual_table_explore: true });
 
@@ -557,6 +561,11 @@ const DatabaseModal: FunctionComponent<DatabaseModalProps> = ({
   const [importingErrorMessage, setImportingErrorMessage] = useState<string>();
   const [passwordFields, setPasswordFields] = useState<string[]>([]);
 
+  const SSHTunnelSwitchComponent =
+    extensionsRegistry.get('ssh_tunnel.form.switch') ?? SSHTunnelSwitch;
+
+  const [useSSHTunneling, setUseSSHTunneling] = useState<boolean>(false);
+
   const conf = useCommonConf();
   const dbImages = getDatabaseImages();
   const connectionAlert = getConnectionAlert();
@@ -572,7 +581,7 @@ const DatabaseModal: FunctionComponent<DatabaseModalProps> = ({
   )?.engine_information?.disable_ssh_tunneling;
   const isSSHTunneling =
     isFeatureEnabled(FeatureFlag.SSH_TUNNELING) &&
-    disableSSHTunnelingForEngine !== undefined;
+    !disableSSHTunnelingForEngine;
   const hasAlert =
     connectionAlert || !!(db?.engine && engineSpecificAlertMapping[db.engine]);
   const useSqlAlchemyForm =
@@ -583,8 +592,6 @@ const DatabaseModal: FunctionComponent<DatabaseModalProps> = ({
       (DB: DatabaseObject) => DB.backend === engine || DB.engine === engine,
     )?.parameters !== undefined;
   const showDBError = validationErrors || dbErrors;
-  const isEmpty = (data?: Object | null) =>
-    !data || (data && Object.keys(data).length === 0);
 
   const dbModel: DatabaseForm =
     availableDbs?.databases?.find(
@@ -652,6 +659,7 @@ const DatabaseModal: FunctionComponent<DatabaseModalProps> = ({
     setPasswordFields([]);
     setPasswords({});
     setConfirmedOverwrite(false);
+    setUseSSHTunneling(false);
     onHide();
   };
 
@@ -1145,6 +1153,12 @@ const DatabaseModal: FunctionComponent<DatabaseModalProps> = ({
     setPasswordFields([...passwordsNeeded]);
   }, [passwordsNeeded]);
 
+  useEffect(() => {
+    if (db) {
+      setUseSSHTunneling(!isEmpty(db?.ssh_tunnel));
+    }
+  }, [db]);
+
   const onDbImport = async (info: UploadChangeParam) => {
     setImportingErrorMessage('');
     setPasswordFields([]);
@@ -1325,10 +1339,7 @@ const DatabaseModal: FunctionComponent<DatabaseModalProps> = ({
 
   const renderSSHTunnelForm = () => (
     <SSHTunnelForm
-      isEditMode={isEditMode}
-      isSSHTunneling={isSSHTunneling}
       db={db as DatabaseObject}
-      dbFetched={dbFetched as DatabaseObject}
       onSSHTunnelParametersChange={({
         target,
       }: {
@@ -1346,11 +1357,6 @@ const DatabaseModal: FunctionComponent<DatabaseModalProps> = ({
           payload: { login_method: method },
         })
       }
-      removeSSHTunnelConfig={() =>
-        setDB({
-          type: ActionType.removeSSHTunnelConfig,
-        })
-      }
     />
   );
 
@@ -1558,7 +1564,16 @@ const DatabaseModal: FunctionComponent<DatabaseModalProps> = ({
                 testConnection={testConnection}
                 testInProgress={testInProgress}
               >
-                {isSSHTunneling && renderSSHTunnelForm()}
+                <SSHTunnelSwitchComponent
+                  isEditMode={isEditMode}
+                  dbFetched={dbFetched}
+                  disableSSHTunnelingForEngine={disableSSHTunnelingForEngine}
+                  useSSHTunneling={useSSHTunneling}
+                  setUseSSHTunneling={setUseSSHTunneling}
+                  setDB={setDB}
+                  isSSHTunneling={isSSHTunneling}
+                />
+                {useSSHTunneling && renderSSHTunnelForm()}
               </SqlAlchemyForm>
               {isDynamic(db?.backend || db?.engine) && !isEditMode && (
                 <div css={(theme: SupersetTheme) => infoTooltip(theme)}>
@@ -1832,7 +1847,18 @@ const DatabaseModal: FunctionComponent<DatabaseModalProps> = ({
                   validationErrors={validationErrors}
                   getPlaceholder={getPlaceholder}
                 />
-                {isSSHTunneling && (
+                <SSHTunnelContainer>
+                  <SSHTunnelSwitchComponent
+                    isEditMode={isEditMode}
+                    dbFetched={dbFetched}
+                    disableSSHTunnelingForEngine={disableSSHTunnelingForEngine}
+                    useSSHTunneling={useSSHTunneling}
+                    setUseSSHTunneling={setUseSSHTunneling}
+                    setDB={setDB}
+                    isSSHTunneling={isSSHTunneling}
+                  />
+                </SSHTunnelContainer>
+                {useSSHTunneling && (
                   <SSHTunnelContainer>
                     {renderSSHTunnelForm()}
                   </SSHTunnelContainer>