fix: better handle datasource exceptions (#13578)

* fix: handle datasource injected security exception

* add tests

* fix error text on create update dbs

* fix lint

* revert create update message

* fix test

* add sqlalchemy exceptions

(cherry picked from commit 1e88408418)

superset-frontend/spec/javascripts/sqllab/TemplateParamsEditor_spec.tsx

@@ -0,0 +1,57 @@
+/**
+ * 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, { ReactNode } from 'react';
+import {
+  render,
+  fireEvent,
+  getByText,
+  waitFor,
+} from 'spec/helpers/testing-library';
+import brace from 'brace';
+import { ThemeProvider, supersetTheme } from '@superset-ui/core';
+
+import TemplateParamsEditor from 'src/SqlLab/components/TemplateParamsEditor';
+
+const ThemeWrapper = ({ children }: { children: ReactNode }) => (
+  <ThemeProvider theme={supersetTheme}>{children}</ThemeProvider>
+);
+
+describe('TemplateParamsEditor', () => {
+  it('should render with a title', () => {
+    const { container } = render(
+      <TemplateParamsEditor code="FOO" language="json" onChange={() => {}} />,
+      { wrapper: ThemeWrapper },
+    );
+    expect(container.querySelector('div[role="button"]')).toBeInTheDocument();
+  });
+
+  it('should open a modal with the ace editor', async () => {
+    const { container, baseElement } = render(
+      <TemplateParamsEditor code="FOO" language="json" onChange={() => {}} />,
+      { wrapper: ThemeWrapper },
+    );
+    fireEvent.click(getByText(container, 'Parameters'));
+    const spy = jest.spyOn(brace, 'acequire');
+    spy.mockReturnValue({ setCompleters: () => 'foo' });
+    await waitFor(() => {
+      expect(baseElement.querySelector('#brace-editor')).toBeInTheDocument();
+    });
+  });
+});

superset-frontend/src/SqlLab/components/TemplateParamsEditor.tsx

@@ -17,7 +17,6 @@
  * under the License.
  */
 import React, { useState, useEffect } from 'react';
-import PropTypes from 'prop-types';
 import Badge from 'src/common/components/Badge';
 import { t, styled } from '@superset-ui/core';
 import { InfoTooltipWithTrigger } from '@superset-ui/chart-controls';
@@ -26,17 +25,7 @@ import { debounce } from 'lodash';
 import ModalTrigger from 'src/components/ModalTrigger';
 import { ConfigEditor } from 'src/components/AsyncAceEditor';
 import { FAST_DEBOUNCE } from 'src/constants';
-
-const propTypes = {
-  onChange: PropTypes.func,
-  code: PropTypes.string,
-  language: PropTypes.oneOf(['yaml', 'json']),
-};
-
-const defaultProps = {
-  onChange: () => {},
-  code: '{}',
-};
+import { Tooltip } from 'src/common/components/Tooltip';
 
 const StyledConfigEditor = styled(ConfigEditor)`
   &.ace_editor {
@@ -44,8 +33,16 @@ const StyledConfigEditor = styled(ConfigEditor)`
   }
 `;
 
-function TemplateParamsEditor({ code, language, onChange }) {
-  const [parsedJSON, setParsedJSON] = useState();
+function TemplateParamsEditor({
+  code = '{}',
+  language,
+  onChange = () => {},
+}: {
+  code: string;
+  language: 'yaml' | 'json';
+  onChange: () => void;
+}) {
+  const [parsedJSON, setParsedJSON] = useState({});
   const [isValid, setIsValid] = useState(true);
 
   useEffect(() => {
@@ -53,7 +50,7 @@ function TemplateParamsEditor({ code, language, onChange }) {
       setParsedJSON(JSON.parse(code));
       setIsValid(true);
     } catch {
-      setParsedJSON({});
+      setParsedJSON({} as any);
       setIsValid(false);
     }
   }, [code]);
@@ -96,25 +93,29 @@ function TemplateParamsEditor({ code, language, onChange }) {
     <ModalTrigger
       modalTitle={t('Template parameters')}
       triggerNode={
-        <div tooltip={t('Edit template parameters')} buttonSize="small">
-          {`${t('Parameters')} `}
-          <Badge count={paramCount} />
-          {!isValid && (
-            <InfoTooltipWithTrigger
-              icon="exclamation-triangle"
-              bsStyle="danger"
-              tooltip={t('Invalid JSON')}
-              label="invalid-json"
-            />
-          )}
-        </div>
+        <Tooltip
+          id="parameters-tooltip"
+          placement="top"
+          title={t('Edit template parameters')}
+          trigger={['hover']}
+        >
+          <div role="button">
+            {`${t('Parameters')} `}
+            <Badge count={paramCount} />
+            {!isValid && (
+              <InfoTooltipWithTrigger
+                icon="exclamation-triangle"
+                bsStyle="danger"
+                tooltip={t('Invalid JSON')}
+                label="invalid-json"
+              />
+            )}
+          </div>
+        </Tooltip>
       }
       modalBody={modalBody}
     />
   );
 }
 
-TemplateParamsEditor.propTypes = propTypes;
-TemplateParamsEditor.defaultProps = defaultProps;
-
 export default TemplateParamsEditor;

superset-frontend/src/SqlLab/reducers/getInitialState.js

@@ -74,13 +74,13 @@ export default function getInitialState({
         id: id.toString(),
         loaded: true,
         title: activeTab.label,
-        sql: activeTab.sql,
-        selectedText: null,
+        sql: activeTab.sql || undefined,
+        selectedText: undefined,
         latestQueryId: activeTab.latest_query
           ? activeTab.latest_query.id
           : null,
         autorun: activeTab.autorun,
-        templateParams: activeTab.template_params,
+        templateParams: activeTab.template_params || undefined,
         dbId: activeTab.database_id,
         functionNames: [],
         schema: activeTab.schema,

superset-frontend/src/SqlLab/reducers/getInitialState.test.ts

@@ -36,8 +36,19 @@ const apiData = {
     username: 'some name',
   },
 };
+const apiDataWithTabState = {
+  ...apiData,
+  tab_state_ids: [{ id: 1 }],
+  active_tab: { id: 1, table_schemas: [] },
+};
 describe('getInitialState', () => {
   it('should output the user that is passed in', () => {
     expect(getInitialState(apiData).sqlLab.user.userId).toEqual(1);
   });
+  it('should return undefined instead of null for templateParams', () => {
+    expect(
+      getInitialState(apiDataWithTabState).sqlLab.queryEditors[0]
+        .templateParams,
+    ).toBeUndefined();
+  });
 });

superset-frontend/src/components/AsyncAceEditor/index.tsx

@@ -119,6 +119,8 @@ export default function AsyncAceEditor(
           mode = inferredMode,
           theme = inferredTheme,
           tabSize = defaultTabSize,
+          defaultValue = '',
+          value = '',
           ...props
         },
         ref,
@@ -150,6 +152,8 @@ export default function AsyncAceEditor(
             mode={mode}
             theme={theme}
             tabSize={tabSize}
+            defaultValue={defaultValue}
+            value={value || ''}
             {...props}
           />
         );

superset/databases/commands/update.py

@@ -54,7 +54,7 @@ class UpdateDatabaseCommand(BaseCommand):
             # TODO Improve this simplistic implementation for catching DB conn fails
             try:
                 schemas = database.get_all_schema_names()
-            except Exception:
+            except Exception as ex:
                 db.session.rollback()
                 raise DatabaseConnectionFailedError()
             for schema in schemas:

superset/migrations/versions/67da9ef1ef9c_add_hide_left_bar_to_tabstate.py

@@ -29,12 +29,18 @@ down_revision = "1412ec1e5a7b"
 import sqlalchemy as sa
 from alembic import op
 from sqlalchemy.dialects import mysql
+from sqlalchemy.sql import expression
 
 
 def upgrade():
     with op.batch_alter_table("tab_state") as batch_op:
         batch_op.add_column(
-            sa.Column("hide_left_bar", sa.Boolean(), nullable=False, default=False)
+            sa.Column(
+                "hide_left_bar",
+                sa.Boolean(),
+                nullable=False,
+                server_default=expression.false(),
+            )
         )
 
 

superset/models/dashboard.py

@@ -38,6 +38,7 @@ from sqlalchemy import (
     UniqueConstraint,
 )
 from sqlalchemy.engine.base import Connection
+from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.orm import relationship, sessionmaker, subqueryload
 from sqlalchemy.orm.mapper import Mapper
 from sqlalchemy.orm.session import object_session
@@ -48,6 +49,7 @@ from superset import app, ConnectorRegistry, db, is_feature_enabled, security_ma
 from superset.connectors.base.models import BaseDatasource
 from superset.connectors.druid.models import DruidColumn, DruidMetric
 from superset.connectors.sqla.models import SqlMetric, TableColumn
+from superset.exceptions import SupersetException
 from superset.extensions import cache_manager
 from superset.models.helpers import AuditMixinNullable, ImportExportMixin
 from superset.models.slice import Slice
@@ -241,18 +243,22 @@ class Dashboard(  # pylint: disable=too-many-instance-attributes
         """Bootstrap data for rendering the dashboard page."""
         slices = self.slices
         datasource_slices = utils.indexed(slices, "datasource")
+        try:
+            datasources = {
+                # Filter out unneeded fields from the datasource payload
+                datasource.uid: datasource.data_for_slices(slices)
+                for datasource, slices in datasource_slices.items()
+                if datasource
+            }
+        except (SupersetException, SQLAlchemyError):
+            datasources = {}
         return {
             # dashboard metadata
             "dashboard": self.data,
             # slices metadata
             "slices": [slc.data for slc in slices],
             # datasource metadata
-            "datasources": {
-                # Filter out unneeded fields from the datasource payload
-                datasource.uid: datasource.data_for_slices(slices)
-                for datasource, slices in datasource_slices.items()
-                if datasource
-            },
+            "datasources": datasources,
         }
 
     @property  # type: ignore

superset/views/core.py

@@ -666,7 +666,7 @@ class Superset(BaseSupersetView):  # pylint: disable=too-many-public-methods
     @event_logger.log_this
     @expose("/explore/<datasource_type>/<int:datasource_id>/", methods=["GET", "POST"])
     @expose("/explore/", methods=["GET", "POST"])
-    def explore(  # pylint: disable=too-many-locals,too-many-return-statements
+    def explore(  # pylint: disable=too-many-locals,too-many-return-statements,too-many-statements
         self, datasource_type: Optional[str] = None, datasource_id: Optional[int] = None
     ) -> FlaskResponse:
         user_id = g.user.get_id() if g.user else None
@@ -789,12 +789,18 @@ class Superset(BaseSupersetView):  # pylint: disable=too-many-public-methods
             "name": datasource_name,
             "columns": [],
             "metrics": [],
+            "database": {"id": 0, "backend": "",},
         }
+        try:
+            datasource_data = datasource.data if datasource else dummy_datasource_data
+        except (SupersetException, SQLAlchemyError):
+            datasource_data = dummy_datasource_data
+
         bootstrap_data = {
             "can_add": slice_add_perm,
             "can_download": slice_download_perm,
             "can_overwrite": slice_overwrite_perm,
-            "datasource": datasource.data if datasource else dummy_datasource_data,
+            "datasource": datasource_data,
             "form_data": form_data,
             "datasource_id": datasource_id,
             "datasource_type": datasource_type,

tests/core_tests.py

@@ -36,7 +36,7 @@ from unittest import mock, skipUnless
 
 import pandas as pd
 import sqlalchemy as sqla
-
+from sqlalchemy.exc import SQLAlchemyError
 from superset.models.cache import CacheKey
 from superset.utils.core import get_example_database
 from tests.fixtures.energy_dashboard import load_energy_table_with_slice
@@ -53,6 +53,7 @@ from superset import (
 from superset.connectors.sqla.models import SqlaTable
 from superset.db_engine_specs.base import BaseEngineSpec
 from superset.db_engine_specs.mssql import MssqlEngineSpec
+from superset.exceptions import SupersetException
 from superset.extensions import async_query_manager
 from superset.models import core as models
 from superset.models.annotations import Annotation, AnnotationLayer
@@ -1474,6 +1475,59 @@ class TestCore(SupersetTestCase):
             "my_col"
         ]
 
+    @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices")
+    @mock.patch("superset.models.core.DB_CONNECTION_MUTATOR")
+    def test_explore_injected_exceptions(self, mock_db_connection_mutator):
+        """
+        Handle injected exceptions from the db mutator
+        """
+        # Assert we can handle a custom exception at the mutator level
+        exception = SupersetException("Error message")
+        mock_db_connection_mutator.side_effect = exception
+        slice = db.session.query(Slice).first()
+        url = f"/superset/explore/?form_data=%7B%22slice_id%22%3A%20{slice.id}%7D"
+
+        self.login()
+        data = self.get_resp(url)
+        self.assertIn("Error message", data)
+
+        # Assert we can handle a driver exception at the mutator level
+        exception = SQLAlchemyError("Error message")
+        mock_db_connection_mutator.side_effect = exception
+        slice = db.session.query(Slice).first()
+        url = f"/superset/explore/?form_data=%7B%22slice_id%22%3A%20{slice.id}%7D"
+
+        self.login()
+        data = self.get_resp(url)
+        self.assertIn("Error message", data)
+
+    @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices")
+    @mock.patch("superset.models.core.DB_CONNECTION_MUTATOR")
+    def test_dashboard_injected_exceptions(self, mock_db_connection_mutator):
+        """
+        Handle injected exceptions from the db mutator
+        """
+
+        # Assert we can handle a custom excetion at the mutator level
+        exception = SupersetException("Error message")
+        mock_db_connection_mutator.side_effect = exception
+        dash = db.session.query(Dashboard).first()
+        url = f"/superset/dashboard/{dash.id}/"
+
+        self.login()
+        data = self.get_resp(url)
+        self.assertIn("Error message", data)
+
+        # Assert we can handle a driver exception at the mutator level
+        exception = SQLAlchemyError("Error message")
+        mock_db_connection_mutator.side_effect = exception
+        dash = db.session.query(Dashboard).first()
+        url = f"/superset/dashboard/{dash.id}/"
+
+        self.login()
+        data = self.get_resp(url)
+        self.assertIn("Error message", data)
+
 
 if __name__ == "__main__":
     unittest.main()