diff --git a/superset/databases/commands/create.py b/superset/databases/commands/create.py index c438444d043e9..373bcdfce365a 100644 --- a/superset/databases/commands/create.py +++ b/superset/databases/commands/create.py @@ -21,6 +21,7 @@ from flask_appbuilder.security.sqla.models import User from marshmallow import ValidationError +from superset import app from superset.commands.base import BaseCommand from superset.dao.exceptions import DAOCreateFailedError from superset.databases.commands.exceptions import ( @@ -35,6 +36,8 @@ from superset.extensions import db, security_manager logger = logging.getLogger(__name__) +config = app.config +stats_logger = config["STATS_LOGGER"] class CreateDatabaseCommand(BaseCommand): @@ -52,6 +55,9 @@ def run(self) -> Model: TestConnectionDatabaseCommand(self._actor, self._properties).run() except Exception: db.session.rollback() + stats_logger.incr( + f"db_connection_failed.{database.db_engine_spec.__name__}" + ) raise DatabaseConnectionFailedError() # adding a new database we always want to force refresh schema list diff --git a/superset/databases/commands/test_connection.py b/superset/databases/commands/test_connection.py index d35d07aaa781e..4899350e45abc 100644 --- a/superset/databases/commands/test_connection.py +++ b/superset/databases/commands/test_connection.py @@ -23,6 +23,7 @@ from sqlalchemy.engine.url import make_url from sqlalchemy.exc import DBAPIError, NoSuchModuleError +from superset import app from superset.commands.base import BaseCommand from superset.databases.commands.exceptions import ( DatabaseSecurityUnsafeError, @@ -35,6 +36,8 @@ from superset.models.core import Database logger = logging.getLogger(__name__) +config = app.config +stats_logger = config["STATS_LOGGER"] class TestConnectionDatabaseCommand(BaseCommand): @@ -68,11 +71,20 @@ def run(self) -> None: raise DatabaseTestConnectionDriverError( message=_("Could not load database driver: {}").format(driver_name), ) - except DBAPIError: + except DBAPIError as ex: + stats_logger.incr( + f"test_connection_error.{make_url(uri).drivername}.{ex.__class__.__name__}" + ) raise DatabaseTestConnectionFailedError() except SupersetSecurityException as ex: + stats_logger.incr( + f"test_connection_error.{make_url(uri).drivername}.{ex.__class__.__name__}" + ) raise DatabaseSecurityUnsafeError(message=str(ex)) - except Exception: + except Exception as ex: + stats_logger.incr( + f"test_connection_error.{make_url(uri).drivername}.{ex.__class__.__name__}" + ) raise DatabaseTestConnectionUnexpectedError() def validate(self) -> None: diff --git a/tests/databases/commands_tests.py b/tests/databases/commands_tests.py index 3b1767fdc2397..8eaece501c6a0 100644 --- a/tests/databases/commands_tests.py +++ b/tests/databases/commands_tests.py @@ -15,18 +15,30 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=no-self-use, invalid-name +from unittest import mock from unittest.mock import patch import pytest import yaml +from sqlalchemy.exc import DBAPIError from superset import db, security_manager from superset.commands.exceptions import CommandInvalidError from superset.commands.importers.exceptions import IncorrectVersionError from superset.connectors.sqla.models import SqlaTable -from superset.databases.commands.exceptions import DatabaseNotFoundError +from superset.databases.commands.exceptions import ( + DatabaseNotFoundError, + DatabaseSecurityUnsafeError, + DatabaseTestConnectionDriverError, + DatabaseTestConnectionFailedError, + DatabaseTestConnectionUnexpectedError, +) from superset.databases.commands.export import ExportDatabasesCommand from superset.databases.commands.importers.v1 import ImportDatabasesCommand +from superset.databases.commands.test_connection import TestConnectionDatabaseCommand +from superset.databases.schemas import DatabaseTestConnectionSchema +from superset.errors import SupersetError +from superset.exceptions import SupersetSecurityException from superset.models.core import Database from superset.utils.core import backend, get_example_database from tests.base_tests import SupersetTestCase @@ -508,3 +520,68 @@ def test_import_v1_rollback(self, mock_import_dataset): # verify that the database was not added new_num_databases = db.session.query(Database).count() assert new_num_databases == num_databases + + +class TestTestConnectionDatabaseCommand(SupersetTestCase): + @mock.patch("superset.databases.dao.DatabaseDAO.build_db_for_connection_test") + @mock.patch("superset.databases.commands.test_connection.stats_logger.incr") + def test_connection_db_exception( + self, mock_stats_logger, mock_build_db_for_connection_test + ): + """Test that users can't export databases they don't have access to""" + database = get_example_database() + mock_build_db_for_connection_test.side_effect = Exception( + "An error has occurred!" + ) + db_uri = database.sqlalchemy_uri_decrypted + json_payload = {"sqlalchemy_uri": db_uri} + command_without_db_name = TestConnectionDatabaseCommand( + security_manager.find_user("admin"), json_payload + ) + + with self.assertRaises(DatabaseTestConnectionUnexpectedError): + command_without_db_name.run() + + mock_stats_logger.assert_called() + + @mock.patch("superset.databases.dao.DatabaseDAO.build_db_for_connection_test") + @mock.patch("superset.databases.commands.test_connection.stats_logger.incr") + def test_connection_superset_security_connection( + self, mock_stats_logger, mock_build_db_for_connection_test + ): + """Test that users can't export databases they don't have access to""" + database = get_example_database() + mock_build_db_for_connection_test.side_effect = SupersetSecurityException( + SupersetError(error_type=500, message="test", level="info", extra={}) + ) + db_uri = database.sqlalchemy_uri_decrypted + json_payload = {"sqlalchemy_uri": db_uri} + command_without_db_name = TestConnectionDatabaseCommand( + security_manager.find_user("admin"), json_payload + ) + + with self.assertRaises(DatabaseSecurityUnsafeError): + command_without_db_name.run() + + mock_stats_logger.assert_called() + + @mock.patch("superset.databases.dao.DatabaseDAO.build_db_for_connection_test") + @mock.patch("superset.databases.commands.test_connection.stats_logger.incr") + def test_connection_db_api_exc( + self, mock_stats_logger, mock_build_db_for_connection_test + ): + """Test that users can't export databases they don't have access to""" + database = get_example_database() + mock_build_db_for_connection_test.side_effect = DBAPIError( + statement="error", params={}, orig={} + ) + db_uri = database.sqlalchemy_uri_decrypted + json_payload = {"sqlalchemy_uri": db_uri} + command_without_db_name = TestConnectionDatabaseCommand( + security_manager.find_user("admin"), json_payload + ) + + with self.assertRaises(DatabaseTestConnectionFailedError): + command_without_db_name.run() + + mock_stats_logger.assert_called()