diff --git a/superset/cli/importexport.py b/superset/cli/importexport.py index d6bbeb67fe287..1ddaf7cfc8d00 100755 --- a/superset/cli/importexport.py +++ b/superset/cli/importexport.py @@ -74,9 +74,8 @@ def export_dashboards(dashboard_file: Optional[str] = None) -> None: from superset.dashboards.commands.export import ExportDashboardsCommand from superset.models.dashboard import Dashboard - g.user = security_manager.find_user( # pylint: disable=assigning-non-slot - username="admin" - ) + # pylint: disable=assigning-non-slot + g.user = security_manager.find_user(username="admin") dashboard_ids = [id_ for (id_,) in db.session.query(Dashboard.id).all()] timestamp = datetime.now().strftime("%Y%m%dT%H%M%S") @@ -110,9 +109,8 @@ def export_datasources(datasource_file: Optional[str] = None) -> None: from superset.connectors.sqla.models import SqlaTable from superset.datasets.commands.export import ExportDatasetsCommand - g.user = security_manager.find_user( # pylint: disable=assigning-non-slot - username="admin" - ) + # pylint: disable=assigning-non-slot + g.user = security_manager.find_user(username="admin") dataset_ids = [id_ for (id_,) in db.session.query(SqlaTable.id).all()] timestamp = datetime.now().strftime("%Y%m%dT%H%M%S") @@ -153,9 +151,8 @@ def import_dashboards(path: str, username: Optional[str]) -> None: ) if username is not None: - g.user = security_manager.find_user( # pylint: disable=assigning-non-slot - username=username - ) + # pylint: disable=assigning-non-slot + g.user = security_manager.find_user(username=username) if is_zipfile(path): with ZipFile(path) as bundle: contents = get_contents_from_bundle(bundle) @@ -320,9 +317,8 @@ def import_dashboards(path: str, recursive: bool, username: str) -> None: elif path_object.exists() and recursive: files.extend(path_object.rglob("*.json")) if username is not None: - g.user = security_manager.find_user( # pylint: disable=assigning-non-slot - username=username - ) + # pylint: disable=assigning-non-slot + g.user = security_manager.find_user(username=username) contents = {} for path_ in files: with open(path_) as file: diff --git a/superset/databases/commands/test_connection.py b/superset/databases/commands/test_connection.py index 1155f4774db42..6fec767a1f52d 100644 --- a/superset/databases/commands/test_connection.py +++ b/superset/databases/commands/test_connection.py @@ -23,7 +23,6 @@ from flask_appbuilder.security.sqla.models import User from flask_babel import gettext as _ from func_timeout import func_timeout, FunctionTimedOut -from sqlalchemy.engine.url import make_url from sqlalchemy.exc import DBAPIError, NoSuchModuleError from superset.commands.base import BaseCommand @@ -34,6 +33,7 @@ DatabaseTestConnectionUnexpectedError, ) from superset.databases.dao import DatabaseDAO +from superset.databases.utils import make_url_safe from superset.errors import ErrorLevel, SupersetErrorType from superset.exceptions import SupersetSecurityException, SupersetTimeoutException from superset.extensions import event_logger @@ -55,7 +55,7 @@ def run(self) -> None: uri = self._model.sqlalchemy_uri_decrypted # context for error messages - url = make_url(uri) + url = make_url_safe(uri) context = { "hostname": url.host, "password": url.password, diff --git a/superset/databases/commands/validate.py b/superset/databases/commands/validate.py index 91e76d8d55efb..fa05dc7c3030c 100644 --- a/superset/databases/commands/validate.py +++ b/superset/databases/commands/validate.py @@ -20,7 +20,6 @@ from flask_appbuilder.security.sqla.models import User from flask_babel import gettext as __ -from sqlalchemy.engine.url import make_url from superset.commands.base import BaseCommand from superset.databases.commands.exceptions import ( @@ -30,6 +29,7 @@ InvalidParametersError, ) from superset.databases.dao import DatabaseDAO +from superset.databases.utils import make_url_safe from superset.db_engine_specs import get_engine_specs from superset.db_engine_specs.base import BasicParametersMixin from superset.errors import ErrorLevel, SupersetError, SupersetErrorType @@ -121,7 +121,7 @@ def run(self) -> None: with closing(engine.raw_connection()) as conn: alive = engine.dialect.do_ping(conn) except Exception as ex: - url = make_url(sqlalchemy_uri) + url = make_url_safe(sqlalchemy_uri) context = { "hostname": url.host, "password": url.password, diff --git a/superset/databases/schemas.py b/superset/databases/schemas.py index 4fa38415ef2fc..dd60e87d167f0 100644 --- a/superset/databases/schemas.py +++ b/superset/databases/schemas.py @@ -24,10 +24,10 @@ from marshmallow.validate import Length, ValidationError from marshmallow_enum import EnumField from sqlalchemy import MetaData -from sqlalchemy.engine.url import make_url -from sqlalchemy.exc import ArgumentError from superset import db +from superset.databases.commands.exceptions import DatabaseInvalidError +from superset.databases.utils import make_url_safe from superset.db_engine_specs import BaseEngineSpec, get_engine_specs from superset.exceptions import CertificateException, SupersetSecurityException from superset.models.core import ConfigurationMethod, Database, PASSWORD_MASK @@ -144,8 +144,8 @@ def sqlalchemy_uri_validator(value: str) -> str: Validate if it's a valid SQLAlchemy URI and refuse SQLLite by default """ try: - uri = make_url(value.strip()) - except (ArgumentError, AttributeError, ValueError) as ex: + uri = make_url_safe(value.strip()) + except DatabaseInvalidError as ex: raise ValidationError( [ _( @@ -649,7 +649,7 @@ def validate_password(self, data: Dict[str, Any], **kwargs: Any) -> None: return uri = data["sqlalchemy_uri"] - password = make_url(uri).password + password = make_url_safe(uri).password if password == PASSWORD_MASK and data.get("password") is None: raise ValidationError("Must provide a password for the database") diff --git a/superset/databases/utils.py b/superset/databases/utils.py index 4d475177f7211..cf54f6da6aecf 100644 --- a/superset/databases/utils.py +++ b/superset/databases/utils.py @@ -16,14 +16,15 @@ # under the License. from typing import Any, Dict, List, Optional -from superset import app -from superset.models.core import Database +from sqlalchemy.engine.url import make_url, URL -custom_password_store = app.config["SQLALCHEMY_CUSTOM_PASSWORD_STORE"] +from superset.databases.commands.exceptions import DatabaseInvalidError def get_foreign_keys_metadata( - database: Database, table_name: str, schema_name: Optional[str] + database: Any, + table_name: str, + schema_name: Optional[str], ) -> List[Dict[str, Any]]: foreign_keys = database.get_foreign_keys(table_name, schema_name) for fk in foreign_keys: @@ -33,7 +34,7 @@ def get_foreign_keys_metadata( def get_indexes_metadata( - database: Database, table_name: str, schema_name: Optional[str] + database: Any, table_name: str, schema_name: Optional[str] ) -> List[Dict[str, Any]]: indexes = database.get_indexes(table_name, schema_name) for idx in indexes: @@ -51,7 +52,7 @@ def get_col_type(col: Dict[Any, Any]) -> str: def get_table_metadata( - database: Database, table_name: str, schema_name: Optional[str] + database: Any, table_name: str, schema_name: Optional[str] ) -> Dict[str, Any]: """ Get table metadata information, including type, pk, fks. @@ -101,3 +102,17 @@ def get_table_metadata( "indexes": keys, "comment": table_comment, } + + +def make_url_safe(raw_url: str) -> URL: + """ + Wrapper for SQLAlchemy's make_url(), which tends to raise too detailed of + errors, which inevitably find their way into server logs. ArgumentErrors + tend to contain usernames and passwords, which makes them non-log-friendly + :param raw_url: + :return: + """ + try: + return make_url(raw_url.strip()) + except Exception: + raise DatabaseInvalidError() # pylint: disable=raise-missing-from diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index e60ce78d72b10..1393fcdac5915 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -48,7 +48,7 @@ from sqlalchemy.engine.base import Engine from sqlalchemy.engine.interfaces import Compiled, Dialect from sqlalchemy.engine.reflection import Inspector -from sqlalchemy.engine.url import make_url, URL +from sqlalchemy.engine.url import URL from sqlalchemy.ext.compiler import compiles from sqlalchemy.orm import Session from sqlalchemy.sql import quoted_name, text @@ -58,6 +58,7 @@ from typing_extensions import TypedDict from superset import security_manager, sql_parse +from superset.databases.utils import make_url_safe from superset.errors import ErrorLevel, SupersetError, SupersetErrorType from superset.models.sql_lab import Query from superset.models.sql_types.base import literal_dttm_type_factory @@ -1633,7 +1634,7 @@ def build_sqlalchemy_uri( # pylint: disable=unused-argument def get_parameters_from_uri( # pylint: disable=unused-argument cls, uri: str, encrypted_extra: Optional[Dict[str, Any]] = None ) -> BasicParametersType: - url = make_url(uri) + url = make_url_safe(uri) query = { key: value for (key, value) in url.query.items() diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py index 2c9f81b1bdde0..d844b342cee97 100644 --- a/superset/db_engine_specs/bigquery.py +++ b/superset/db_engine_specs/bigquery.py @@ -28,11 +28,11 @@ from marshmallow.exceptions import ValidationError from sqlalchemy import column from sqlalchemy.engine.base import Engine -from sqlalchemy.engine.url import make_url from sqlalchemy.sql import sqltypes from typing_extensions import TypedDict from superset.databases.schemas import encrypted_field_properties, EncryptedString +from superset.databases.utils import make_url_safe from superset.db_engine_specs.base import BaseEngineSpec from superset.db_engine_specs.exceptions import SupersetDBAPIDisconnectionError from superset.errors import SupersetError, SupersetErrorType @@ -377,7 +377,7 @@ def build_sqlalchemy_uri( def get_parameters_from_uri( cls, uri: str, encrypted_extra: Optional[Dict[str, str]] = None ) -> Any: - value = make_url(uri) + value = make_url_safe(uri) # Building parameters from encrypted_extra and uri if encrypted_extra: diff --git a/superset/db_engine_specs/hive.py b/superset/db_engine_specs/hive.py index 5f7e01d50271f..484e867ab0a88 100644 --- a/superset/db_engine_specs/hive.py +++ b/superset/db_engine_specs/hive.py @@ -31,11 +31,12 @@ from sqlalchemy import Column, text from sqlalchemy.engine.base import Engine from sqlalchemy.engine.reflection import Inspector -from sqlalchemy.engine.url import make_url, URL +from sqlalchemy.engine.url import URL from sqlalchemy.orm import Session from sqlalchemy.sql.expression import ColumnClause, Select from superset.common.db_query_status import QueryStatus +from superset.databases.utils import make_url_safe from superset.db_engine_specs.base import BaseEngineSpec from superset.db_engine_specs.presto import PrestoEngineSpec from superset.exceptions import SupersetException @@ -510,7 +511,7 @@ def update_impersonation_config( :param username: Effective username :return: None """ - url = make_url(uri) + url = make_url_safe(uri) backend_name = url.get_backend_name() # Must be Hive connection, enable impersonation, and set optional param diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py index 77c3d4c2e3383..8675607848328 100644 --- a/superset/db_engine_specs/presto.py +++ b/superset/db_engine_specs/presto.py @@ -34,12 +34,13 @@ from sqlalchemy.engine.base import Engine from sqlalchemy.engine.reflection import Inspector from sqlalchemy.engine.result import RowProxy -from sqlalchemy.engine.url import make_url, URL +from sqlalchemy.engine.url import URL from sqlalchemy.orm import Session from sqlalchemy.sql.expression import ColumnClause, Select from superset import cache_manager, is_feature_enabled from superset.common.db_query_status import QueryStatus +from superset.databases.utils import make_url_safe from superset.db_engine_specs.base import BaseEngineSpec, ColumnTypeMapping from superset.errors import SupersetErrorType from superset.exceptions import SupersetTemplateException @@ -235,7 +236,7 @@ def update_impersonation_config( :param username: Effective username :return: None """ - url = make_url(uri) + url = make_url_safe(uri) backend_name = url.get_backend_name() # Must be Presto connection, enable impersonation, and set optional param diff --git a/superset/db_engine_specs/snowflake.py b/superset/db_engine_specs/snowflake.py index 058ca89c6af29..cf645f8b74c24 100644 --- a/superset/db_engine_specs/snowflake.py +++ b/superset/db_engine_specs/snowflake.py @@ -24,9 +24,10 @@ from apispec.ext.marshmallow import MarshmallowPlugin from flask_babel import gettext as __ from marshmallow import fields, Schema -from sqlalchemy.engine.url import make_url, URL +from sqlalchemy.engine.url import URL from typing_extensions import TypedDict +from superset.databases.utils import make_url_safe from superset.db_engine_specs.postgres import PostgresBaseEngineSpec from superset.errors import ErrorLevel, SupersetError, SupersetErrorType from superset.models.sql_lab import Query @@ -220,7 +221,7 @@ def get_parameters_from_uri( Dict[str, str] ] = None, ) -> Any: - url = make_url(uri) + url = make_url_safe(uri) query = dict(url.query.items()) return { "username": url.username, diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py index 31e9a0aa7b3a3..8ff4cfde59676 100644 --- a/superset/db_engine_specs/trino.py +++ b/superset/db_engine_specs/trino.py @@ -21,8 +21,9 @@ import simplejson as json from flask import current_app -from sqlalchemy.engine.url import make_url, URL +from sqlalchemy.engine.url import URL +from superset.databases.utils import make_url_safe from superset.db_engine_specs.base import BaseEngineSpec from superset.utils import core as utils @@ -107,7 +108,7 @@ def update_impersonation_config( :param username: Effective username :return: None """ - url = make_url(uri) + url = make_url_safe(uri) backend_name = url.get_backend_name() # Must be Trino connection, enable impersonation, and set optional param diff --git a/superset/models/core.py b/superset/models/core.py index fcc7cf16d8ef2..daa0fb9a7ddfc 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -43,9 +43,9 @@ Table, Text, ) -from sqlalchemy.engine import Connection, Dialect, Engine, url +from sqlalchemy.engine import Connection, Dialect, Engine from sqlalchemy.engine.reflection import Inspector -from sqlalchemy.engine.url import make_url, URL +from sqlalchemy.engine.url import URL from sqlalchemy.exc import ArgumentError from sqlalchemy.ext.hybrid import hybrid_property from sqlalchemy.orm import relationship @@ -54,6 +54,7 @@ from sqlalchemy.sql import expression, Select from superset import app, db_engine_specs, is_feature_enabled +from superset.databases.utils import make_url_safe from superset.db_engine_specs.base import TimeGrain from superset.extensions import cache_manager, encrypted_field_factory, security_manager from superset.models.helpers import AuditMixinNullable, ImportExportMixin @@ -242,16 +243,16 @@ def unique_name(self) -> str: @property def url_object(self) -> URL: - return make_url(self.sqlalchemy_uri_decrypted) + return make_url_safe(self.sqlalchemy_uri_decrypted) @property def backend(self) -> str: - sqlalchemy_url = make_url(self.sqlalchemy_uri_decrypted) - return sqlalchemy_url.get_backend_name() # pylint: disable=no-member + sqlalchemy_url = make_url_safe(self.sqlalchemy_uri_decrypted) + return sqlalchemy_url.get_backend_name() @property def parameters(self) -> Dict[str, Any]: - uri = make_url(self.sqlalchemy_uri_decrypted) + uri = make_url_safe(self.sqlalchemy_uri_decrypted) encrypted_extra = self.get_encrypted_extra() try: # pylint: disable=useless-suppression @@ -303,7 +304,7 @@ def connect_args(self) -> Dict[str, Any]: def get_password_masked_url_from_uri( # pylint: disable=invalid-name cls, uri: str ) -> URL: - sqlalchemy_url = make_url(uri) + sqlalchemy_url = make_url_safe(uri) return cls.get_password_masked_url(sqlalchemy_url) @classmethod @@ -314,7 +315,7 @@ def get_password_masked_url(cls, masked_url: URL) -> URL: return url_copy def set_sqlalchemy_uri(self, uri: str) -> None: - conn = sqla.engine.url.make_url(uri.strip()) + conn = make_url_safe(uri.strip()) if conn.password != PASSWORD_MASK and not custom_password_store: # do not over-write the password with the password mask self.password = conn.password @@ -361,7 +362,7 @@ def get_sqla_engine( source: Optional[utils.QuerySource] = None, ) -> Engine: extra = self.get_extra() - sqlalchemy_url = make_url(self.sqlalchemy_uri_decrypted) + sqlalchemy_url = make_url_safe(self.sqlalchemy_uri_decrypted) self.db_engine_spec.adjust_database_uri(sqlalchemy_url, schema) effective_username = self.get_effective_user(sqlalchemy_url, user_name) # If using MySQL or Presto for example, will set url.username @@ -723,7 +724,7 @@ def get_schema_access_for_file_upload( # pylint: disable=invalid-name @property def sqlalchemy_uri_decrypted(self) -> str: try: - conn = sqla.engine.url.make_url(self.sqlalchemy_uri) + conn = make_url_safe(self.sqlalchemy_uri) except (ArgumentError, ValueError): # if the URI is invalid, ignore and return a placeholder url # (so users see 500 less often) @@ -783,7 +784,7 @@ def has_view_by_name(self, view_name: str, schema: Optional[str] = None) -> bool @memoized def get_dialect(self) -> Dialect: - sqla_url = url.make_url(self.sqlalchemy_uri_decrypted) + sqla_url = make_url_safe(self.sqlalchemy_uri_decrypted) return sqla_url.get_dialect()() diff --git a/superset/sql_parse.py b/superset/sql_parse.py index 6bfb63c425c48..e3b2e7c196834 100644 --- a/superset/sql_parse.py +++ b/superset/sql_parse.py @@ -574,7 +574,6 @@ def get_rls_for_table( return None template_processor = dataset.get_template_processor() - # pylint: disable=protected-access predicate = " AND ".join( str(filter_) for filter_ in dataset.get_sqla_row_level_filters(template_processor) diff --git a/superset/tasks/async_queries.py b/superset/tasks/async_queries.py index 6a42d961e9d9d..74adcd080c0c3 100644 --- a/superset/tasks/async_queries.py +++ b/superset/tasks/async_queries.py @@ -47,17 +47,16 @@ def ensure_user_is_set(user_id: Optional[int]) -> None: user_is_not_set = not (hasattr(g, "user") and g.user is not None) if user_is_not_set and user_id is not None: - g.user = security_manager.get_user_by_id( # pylint: disable=assigning-non-slot - user_id - ) + # pylint: disable=assigning-non-slot + g.user = security_manager.get_user_by_id(user_id) elif user_is_not_set: - g.user = ( # pylint: disable=assigning-non-slot - security_manager.get_anonymous_user() - ) + # pylint: disable=assigning-non-slot + g.user = security_manager.get_anonymous_user() def set_form_data(form_data: Dict[str, Any]) -> None: - g.form_data = form_data # pylint: disable=assigning-non-slot + # pylint: disable=assigning-non-slot + g.form_data = form_data def _create_query_context_from_form(form_data: Dict[str, Any]) -> QueryContext: diff --git a/superset/views/core.py b/superset/views/core.py index 9c52c3c0f51af..f806ea7b0f466 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -39,8 +39,7 @@ from flask_appbuilder.security.sqla import models as ab_models from flask_babel import gettext as __, lazy_gettext as _ from sqlalchemy import and_, or_ -from sqlalchemy.engine.url import make_url -from sqlalchemy.exc import ArgumentError, DBAPIError, NoSuchModuleError, SQLAlchemyError +from sqlalchemy.exc import DBAPIError, NoSuchModuleError, SQLAlchemyError from sqlalchemy.orm.session import Session from sqlalchemy.sql import functions as func @@ -73,8 +72,10 @@ from superset.dashboards.dao import DashboardDAO from superset.dashboards.permalink.commands.get import GetDashboardPermalinkCommand from superset.dashboards.permalink.exceptions import DashboardPermalinkGetFailedError +from superset.databases.commands.exceptions import DatabaseInvalidError from superset.databases.dao import DatabaseDAO from superset.databases.filters import DatabaseFilter +from superset.databases.utils import make_url_safe from superset.datasets.commands.exceptions import DatasetNotFoundError from superset.errors import ErrorLevel, SupersetError, SupersetErrorType from superset.exceptions import ( @@ -1330,7 +1331,7 @@ def testconn(self) -> FlaskResponse: # pylint: disable=no-self-use uri = request.json.get("uri") try: if app.config["PREVENT_UNSAFE_DB_CONNECTIONS"]: - check_sqlalchemy_uri(make_url(uri)) + check_sqlalchemy_uri(make_url_safe(uri)) # if the database already exists in the database, only its safe # (password-masked) URI would be shown in the UI and would be passed in the # form data so if the database already exists and the form was submitted @@ -1371,7 +1372,7 @@ def testconn(self) -> FlaskResponse: # pylint: disable=no-self-use return json_error_response(ex.message) except (NoSuchModuleError, ModuleNotFoundError): logger.info("Invalid driver") - driver_name = make_url(uri).drivername + driver_name = make_url_safe(uri).drivername return json_error_response( _( "Could not load database driver: %(driver_name)s", @@ -1379,7 +1380,7 @@ def testconn(self) -> FlaskResponse: # pylint: disable=no-self-use ), 400, ) - except ArgumentError: + except DatabaseInvalidError: logger.info("Invalid URI") return json_error_response( _( @@ -1815,7 +1816,8 @@ def warm_up_cache( # pylint: disable=too-many-locals,no-self-use force=True, ) - g.form_data = form_data # pylint: disable=assigning-non-slot + # pylint: disable=assigning-non-slot + g.form_data = form_data payload = obj.get_payload() delattr(g, "form_data") error = payload["errors"] or None diff --git a/superset/views/database/mixins.py b/superset/views/database/mixins.py index d5a5157ef4f7b..f6f7f1115e201 100644 --- a/superset/views/database/mixins.py +++ b/superset/views/database/mixins.py @@ -19,10 +19,10 @@ from flask import Markup from flask_babel import lazy_gettext as _ from sqlalchemy import MetaData -from sqlalchemy.engine.url import make_url from superset import app, security_manager from superset.databases.filters import DatabaseFilter +from superset.databases.utils import make_url_safe from superset.exceptions import SupersetException from superset.models.core import Database from superset.security.analytics_db_safety import check_sqlalchemy_uri @@ -209,7 +209,7 @@ class DatabaseMixin: def _pre_add_update(self, database: Database) -> None: if app.config["PREVENT_UNSAFE_DB_CONNECTIONS"]: - check_sqlalchemy_uri(make_url(database.sqlalchemy_uri)) + check_sqlalchemy_uri(make_url_safe(database.sqlalchemy_uri)) self.check_extra(database) self.check_encrypted_extra(database) if database.server_cert: diff --git a/superset/views/database/validators.py b/superset/views/database/validators.py index 2b2aa264407ef..93723ac38b8f2 100644 --- a/superset/views/database/validators.py +++ b/superset/views/database/validators.py @@ -19,10 +19,10 @@ from flask_babel import lazy_gettext as _ from marshmallow import ValidationError -from sqlalchemy.engine.url import make_url -from sqlalchemy.exc import ArgumentError from superset import security_manager +from superset.databases.commands.exceptions import DatabaseInvalidError +from superset.databases.utils import make_url_safe from superset.models.core import Database @@ -33,8 +33,8 @@ def sqlalchemy_uri_validator( Check if a user has submitted a valid SQLAlchemy URI """ try: - make_url(uri.strip()) - except (ArgumentError, AttributeError) as ex: + make_url_safe(uri.strip()) + except DatabaseInvalidError as ex: raise exception( [ _( diff --git a/tests/integration_tests/utils_tests.py b/tests/integration_tests/utils_tests.py index 765f586ced6c5..5add2c5f6e014 100644 --- a/tests/integration_tests/utils_tests.py +++ b/tests/integration_tests/utils_tests.py @@ -24,6 +24,8 @@ import re from typing import Any, Tuple, List, Optional from unittest.mock import Mock, patch + +from superset.databases.commands.exceptions import DatabaseInvalidError from tests.integration_tests.fixtures.birth_names_dashboard import ( load_birth_names_dashboard_with_slices, load_birth_names_data, @@ -736,7 +738,7 @@ def test_get_or_create_db(self): db.session.commit() def test_get_or_create_db_invalid_uri(self): - with self.assertRaises(ArgumentError): + with self.assertRaises(DatabaseInvalidError): get_or_create_db("test_db", "yoursql:superset.db/()") def test_get_iterable(self):