fix: Use superset generic db to catch external_metadata queries (#13974)

This commit is contained in:
Hugh A. Miles II 2021-04-10 10:15:03 -04:00 committed by GitHub
parent fe08b698f5
commit a4fd6b8f33
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 38 additions and 10 deletions

View File

@ -59,7 +59,11 @@ from superset import app, db, is_feature_enabled, security_manager
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 @@ class SqlaTable( # pylint: disable=too-many-public-methods,too-many-instance-at
)
# 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(

View File

@ -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 @@ class TestDatasource(SupersetTestCase):
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: