From 446627c89c6e25aac295ea5c7b34e3c690dff200 Mon Sep 17 00:00:00 2001
From: "Hugh A. Miles II" <hughmil3s@gmail.com>
Date: Sat, 10 Apr 2021 10:15:03 -0400
Subject: [PATCH] fix: Use superset generic db to catch external_metadata
 queries (#13974)

---
 superset/connectors/sqla/models.py | 27 +++++++++++++++++----------
 tests/datasource_tests.py          | 21 +++++++++++++++++++++
 2 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
index 74b5a7de5a32b..dc12740ad76ff 100644
--- a/superset/connectors/sqla/models.py
+++ b/superset/connectors/sqla/models.py
@@ -59,7 +59,11 @@
 from superset.connectors.base.models import BaseColumn, BaseDatasource, BaseMetric
 from superset.db_engine_specs.base import TimestampExpression
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
-from superset.exceptions import QueryObjectValidationError, SupersetSecurityException
+from superset.exceptions import (
+    QueryObjectValidationError,
+    SupersetGenericDBErrorException,
+    SupersetSecurityException,
+)
 from superset.jinja_context import (
     BaseTemplateProcessor,
     ExtraCache,
@@ -645,15 +649,18 @@ def external_metadata(self) -> List[Dict[str, str]]:
                 )
             # TODO(villebro): refactor to use same code that's used by
             #  sql_lab.py:execute_sql_statements
-            with closing(engine.raw_connection()) as conn:
-                cursor = conn.cursor()
-                query = self.database.apply_limit_to_sql(statements[0])
-                db_engine_spec.execute(cursor, query)
-                result = db_engine_spec.fetch_data(cursor, limit=1)
-                result_set = SupersetResultSet(
-                    result, cursor.description, db_engine_spec
-                )
-                cols = result_set.columns
+            try:
+                with closing(engine.raw_connection()) as conn:
+                    cursor = conn.cursor()
+                    query = self.database.apply_limit_to_sql(statements[0])
+                    db_engine_spec.execute(cursor, query)
+                    result = db_engine_spec.fetch_data(cursor, limit=1)
+                    result_set = SupersetResultSet(
+                        result, cursor.description, db_engine_spec
+                    )
+                    cols = result_set.columns
+            except Exception as exc:
+                raise SupersetGenericDBErrorException(message=str(exc))
         else:
             db_dialect = self.database.get_dialect()
             cols = self.database.get_columns(
diff --git a/tests/datasource_tests.py b/tests/datasource_tests.py
index 31a05981fece1..438079fc5be85 100644
--- a/tests/datasource_tests.py
+++ b/tests/datasource_tests.py
@@ -17,12 +17,14 @@
 """Unit tests for Superset"""
 import json
 from copy import deepcopy
+from unittest import mock
 
 import pytest
 
 from superset import app, ConnectorRegistry, db
 from superset.connectors.sqla.models import SqlaTable
 from superset.datasets.commands.exceptions import DatasetNotFoundError
+from superset.exceptions import SupersetException, SupersetGenericDBErrorException
 from superset.utils.core import get_example_database
 from tests.fixtures.birth_names_dashboard import load_birth_names_dashboard_with_slices
 
@@ -96,6 +98,25 @@ def test_external_metadata_for_mutistatement_virtual_table(self):
             resp = self.get_json_resp(url)
             self.assertEqual(resp["error"], "Only single queries supported")
 
+    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
+    @mock.patch("superset.connectors.sqla.models.SqlaTable.external_metadata")
+    def test_external_metadata_error_return_400(self, mock_get_datasource):
+        self.login(username="admin")
+        tbl = self.get_table_by_name("birth_names")
+        url = f"/datasource/external_metadata/table/{tbl.id}/"
+
+        mock_get_datasource.side_effect = SupersetGenericDBErrorException("oops")
+
+        pytest.raises(
+            SupersetGenericDBErrorException,
+            lambda: ConnectorRegistry.get_datasource(
+                "table", tbl.id, db.session
+            ).external_metadata(),
+        )
+
+        resp = self.client.get(url)
+        assert resp.status_code == 400
+
     def compare_lists(self, l1, l2, key):
         l2_lookup = {o.get(key): o for o in l2}
         for obj1 in l1: