diff --git a/superset/databases/commands/create.py b/superset/databases/commands/create.py index 373bcdfce365..2eabfd2950cb 100644 --- a/superset/databases/commands/create.py +++ b/superset/databases/commands/create.py @@ -21,7 +21,6 @@ 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 ( @@ -33,11 +32,9 @@ ) from superset.databases.commands.test_connection import TestConnectionDatabaseCommand from superset.databases.dao import DatabaseDAO -from superset.extensions import db, security_manager +from superset.extensions import db, event_logger, security_manager logger = logging.getLogger(__name__) -config = app.config -stats_logger = config["STATS_LOGGER"] class CreateDatabaseCommand(BaseCommand): @@ -54,10 +51,12 @@ def run(self) -> Model: try: TestConnectionDatabaseCommand(self._actor, self._properties).run() except Exception: - db.session.rollback() - stats_logger.incr( - f"db_connection_failed.{database.db_engine_spec.__name__}" - ) + with event_logger.log_context( + action="db_connection_failed", + engine=database.db_engine_spec.__name__, + ): + db.session.rollback() + raise DatabaseConnectionFailedError() # adding a new database we always want to force refresh schema list @@ -69,8 +68,9 @@ def run(self) -> Model: security_manager.add_permission_view_menu("database_access", database.perm) db.session.commit() except DAOCreateFailedError as ex: - logger.exception(ex.exception) - raise DatabaseCreateFailedError() + with event_logger.log_context(action=f"db_creation_failed.{ex.exception}"): + logger.exception(ex.exception) + raise DatabaseCreateFailedError() return database def validate(self) -> None: @@ -90,4 +90,6 @@ def validate(self) -> None: if exceptions: exception = DatabaseInvalidError() exception.add_list(exceptions) - raise exception + + with event_logger.log_context(action="db_connection_failed"): + raise exception diff --git a/superset/databases/commands/test_connection.py b/superset/databases/commands/test_connection.py index 4899350e45ab..7e7df456a23a 100644 --- a/superset/databases/commands/test_connection.py +++ b/superset/databases/commands/test_connection.py @@ -23,7 +23,6 @@ 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, @@ -33,11 +32,10 @@ ) from superset.databases.dao import DatabaseDAO from superset.exceptions import SupersetSecurityException +from superset.extensions import event_logger from superset.models.core import Database logger = logging.getLogger(__name__) -config = app.config -stats_logger = config["STATS_LOGGER"] class TestConnectionDatabaseCommand(BaseCommand): @@ -58,34 +56,44 @@ def run(self) -> None: impersonate_user=self._properties.get("impersonate_user", False), encrypted_extra=self._properties.get("encrypted_extra", "{}"), ) - if database is not None: - database.set_sqlalchemy_uri(uri) - database.db_engine_spec.mutate_db_for_connection_test(database) - username = self._actor.username if self._actor is not None else None - engine = database.get_sqla_engine(user_name=username) + if database is None: + raise DBAPIError(None, None, None) + + database.set_sqlalchemy_uri(uri) + database.db_engine_spec.mutate_db_for_connection_test(database) + username = self._actor.username if self._actor is not None else None + engine = database.get_sqla_engine(user_name=username) with closing(engine.raw_connection()) as conn: if not engine.dialect.do_ping(conn): raise DBAPIError(None, None, None) + + with event_logger.log_context( + action="test_connection_success", engine=make_url(uri).drivername, + ): + return except (NoSuchModuleError, ModuleNotFoundError): driver_name = make_url(uri).drivername raise DatabaseTestConnectionDriverError( message=_("Could not load database driver: {}").format(driver_name), ) except DBAPIError as ex: - stats_logger.incr( - f"test_connection_error.{make_url(uri).drivername}.{ex.__class__.__name__}" - ) - raise DatabaseTestConnectionFailedError() + with event_logger.log_context( + action=f"test_connection_error.{ex.__class__.__name__}", + engine=make_url(uri).drivername, + ): + 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 as ex: - stats_logger.incr( - f"test_connection_error.{make_url(uri).drivername}.{ex.__class__.__name__}" - ) - raise DatabaseTestConnectionUnexpectedError() + with event_logger.log_context( + action=f"test_connection_error.{ex.__class__.__name__}", + engine=make_url(uri).drivername, + ): + raise DatabaseSecurityUnsafeError(message=str(ex)) + except Exception as ex: # pylint: disable=broad-except + with event_logger.log_context( + action=f"test_connection_error.{ex.__class__.__name__}", + engine=make_url(uri).drivername, + ): + raise DatabaseTestConnectionUnexpectedError() def validate(self) -> None: database_name = self._properties.get("database_name") diff --git a/superset/utils/log.py b/superset/utils/log.py index 824487e1643c..5f4edf9362ed 100644 --- a/superset/utils/log.py +++ b/superset/utils/log.py @@ -86,13 +86,14 @@ def log_context( # pylint: disable=too-many-locals from superset.views.core import get_form_data start_time = time.time() - referrer = request.referrer[:1000] if request.referrer else None - user_id = g.user.get_id() if hasattr(g, "user") and g.user else None payload_override = {} # yield a helper to add additional payload yield lambda **kwargs: payload_override.update(kwargs) + referrer = request.referrer[:1000] if request.referrer else None + user_id = g.user.get_id() if hasattr(g, "user") and g.user else None + payload = collect_request_payload() if object_ref: payload["object_ref"] = object_ref diff --git a/tests/databases/commands_tests.py b/tests/databases/commands_tests.py index 8eaece501c6a..3ed5bb35bc5b 100644 --- a/tests/databases/commands_tests.py +++ b/tests/databases/commands_tests.py @@ -20,16 +20,20 @@ import pytest import yaml +from sqlalchemy.engine.url import make_url 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.dao.exceptions import DAOCreateFailedError +from superset.databases.commands.create import CreateDatabaseCommand from superset.databases.commands.exceptions import ( + DatabaseCreateFailedError, + DatabaseInvalidError, DatabaseNotFoundError, DatabaseSecurityUnsafeError, - DatabaseTestConnectionDriverError, DatabaseTestConnectionFailedError, DatabaseTestConnectionUnexpectedError, ) @@ -37,7 +41,7 @@ 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.errors import ErrorLevel, SupersetError, SupersetErrorType from superset.exceptions import SupersetSecurityException from superset.models.core import Database from superset.utils.core import backend, get_example_database @@ -523,65 +527,118 @@ def test_import_v1_rollback(self, mock_import_dataset): class TestTestConnectionDatabaseCommand(SupersetTestCase): - @mock.patch("superset.databases.dao.DatabaseDAO.build_db_for_connection_test") - @mock.patch("superset.databases.commands.test_connection.stats_logger.incr") + @patch("superset.extensions.event_logger.log_context") + @patch("superset.databases.dao.DatabaseDAO.build_db_for_connection_test") def test_connection_db_exception( - self, mock_stats_logger, mock_build_db_for_connection_test + self, mock_build_db_connection_test, mock_event_logger, ): - """Test that users can't export databases they don't have access to""" + """Test that exceptions are being properly logged""" + mock_build_db_connection_test.side_effect = [ + DBAPIError("An error occurred!", None, None), + SupersetSecurityException( + SupersetError( + "dummy", + SupersetErrorType.DATASOURCE_SECURITY_ACCESS_ERROR, + ErrorLevel.ERROR, + ) + ), + Exception("An error occurred!"), + ] 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 + test_item = DatabaseTestConnectionSchema().load(json_payload) + command = TestConnectionDatabaseCommand( + security_manager.find_user("admin"), test_item + ) + with self.assertRaises(DatabaseTestConnectionFailedError): + command.run() + + mock_event_logger.assert_called_with( + action="test_connection_error.DBAPIError", + engine=make_url(db_uri).drivername, + ) + + with self.assertRaises(DatabaseSecurityUnsafeError): + command.run() + + mock_event_logger.assert_called_with( + action="test_connection_error.SupersetSecurityException", + engine=make_url(db_uri).drivername, ) with self.assertRaises(DatabaseTestConnectionUnexpectedError): - command_without_db_name.run() + command.run() - mock_stats_logger.assert_called() + mock_event_logger.assert_called_with( + action="test_connection_error.Exception", engine=make_url(db_uri).drivername + ) - @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 + @patch("superset.extensions.event_logger.log_context") + def test_connection_db_success( + self, mock_event_logger, ): - """Test that users can't export databases they don't have access to""" + """Test that test_connection is logging on success""" 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 + test_item = DatabaseTestConnectionSchema().load(json_payload) + command = TestConnectionDatabaseCommand( + security_manager.find_user("admin"), test_item ) + command.run() - with self.assertRaises(DatabaseSecurityUnsafeError): - command_without_db_name.run() + mock_event_logger.assert_called_with( + action="test_connection_success", engine=make_url(db_uri).drivername + ) - 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 +class TestCreateDatabaseCommand(SupersetTestCase): + @patch("superset.databases.dao.DatabaseDAO.create") + @patch("superset.extensions.event_logger.log_context") + @patch( + "superset.databases.commands.test_connection.TestConnectionDatabaseCommand.run" + ) + def test_create_database_error( + self, mock_test_connection_db, mock_event_logger, mock_db_create ): - """Test that users can't export databases they don't have access to""" + """Test that exceptions are being properly logged""" + 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( + command_without_db_name = CreateDatabaseCommand( security_manager.find_user("admin"), json_payload ) - - with self.assertRaises(DatabaseTestConnectionFailedError): + # test with no db name + with self.assertRaises(DatabaseInvalidError): command_without_db_name.run() - mock_stats_logger.assert_called() + mock_event_logger.assert_called_with(action="db_connection_failed") + + # test when connection fails + mock_test_connection_db.side_effect = Exception("An error has occurred!") + mock_db_create.return_value = database + json_payload = {"sqlalchemy_uri": db_uri, "database_name": "foo"} + command = CreateDatabaseCommand( + security_manager.find_user("admin"), json_payload + ) + + with self.assertRaises(Exception): + command.run() + + mock_event_logger.assert_called_with( + action="db_connection_failed", engine=database.db_engine_spec.__name__ + ) + + # test when creation fails + mock_db_create.side_effect = DAOCreateFailedError( + exception=Exception("An error occurred") + ) + + with self.assertRaises(DatabaseCreateFailedError): + command.run() + + mock_event_logger.assert_called_with( + action="db_creation_failed.An error occurred" + )