diff --git a/superset/databases/commands/create.py b/superset/databases/commands/create.py index c438444d04..cb4c3baead 100644 --- a/superset/databases/commands/create.py +++ b/superset/databases/commands/create.py @@ -32,7 +32,7 @@ from superset.databases.commands.exceptions import ( ) 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__) @@ -50,8 +50,12 @@ class CreateDatabaseCommand(BaseCommand): try: TestConnectionDatabaseCommand(self._actor, self._properties).run() - except Exception: + except Exception as ex: # pylint: disable=broad-except db.session.rollback() + event_logger.log_with_context( + action=f"db_creation_failed.{ex.__class__.__name__}", + engine=database.db_engine_spec.__name__, + ) raise DatabaseConnectionFailedError() # adding a new database we always want to force refresh schema list @@ -63,7 +67,10 @@ class CreateDatabaseCommand(BaseCommand): security_manager.add_permission_view_menu("database_access", database.perm) db.session.commit() except DAOCreateFailedError as ex: - logger.exception(ex.exception) + event_logger.log_with_context( + action=f"db_creation_failed.{ex.__class__.__name__}", + engine=database.db_engine_spec.__name__, + ) raise DatabaseCreateFailedError() return database @@ -84,4 +91,7 @@ class CreateDatabaseCommand(BaseCommand): if exceptions: exception = DatabaseInvalidError() exception.add_list(exceptions) + event_logger.log_with_context( + action=f"db_connection_failed.{exception.__class__.__name__}" + ) raise exception diff --git a/superset/databases/commands/test_connection.py b/superset/databases/commands/test_connection.py index d35d07aaa7..101b07220f 100644 --- a/superset/databases/commands/test_connection.py +++ b/superset/databases/commands/test_connection.py @@ -20,7 +20,6 @@ from typing import Any, Dict, Optional from flask_appbuilder.security.sqla.models import User from flask_babel import gettext as _ -from sqlalchemy.engine.url import make_url from sqlalchemy.exc import DBAPIError, NoSuchModuleError from superset.commands.base import BaseCommand @@ -32,6 +31,7 @@ from superset.databases.commands.exceptions import ( ) 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__) @@ -55,24 +55,42 @@ class TestConnectionDatabaseCommand(BaseCommand): 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) + + 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) - except (NoSuchModuleError, ModuleNotFoundError): - driver_name = make_url(uri).drivername - raise DatabaseTestConnectionDriverError( - message=_("Could not load database driver: {}").format(driver_name), + + except (NoSuchModuleError, ModuleNotFoundError) as ex: + event_logger.log_with_context( + action=f"test_connection_error.{ex.__class__.__name__}", + engine=database.db_engine_spec.__name__, + ) + raise DatabaseTestConnectionDriverError( + message=_("Could not load database driver: {}").format( + database.db_engine_spec.__name__ + ), + ) + except DBAPIError as ex: + event_logger.log_with_context( + action=f"test_connection_error.{ex.__class__.__name__}", + engine=database.db_engine_spec.__name__, ) - except DBAPIError: raise DatabaseTestConnectionFailedError() except SupersetSecurityException as ex: + event_logger.log_with_context( + action=f"test_connection_error.{ex.__class__.__name__}", + engine=database.db_engine_spec.__name__, + ) raise DatabaseSecurityUnsafeError(message=str(ex)) - except Exception: + except Exception as ex: # pylint: disable=broad-except + event_logger.log_with_context( + action=f"test_connection_error.{ex.__class__.__name__}", + engine=database.db_engine_spec.__name__, + ) raise DatabaseTestConnectionUnexpectedError() def validate(self) -> None: diff --git a/superset/databases/dao.py b/superset/databases/dao.py index 2e89ad0735..e24cdcdee1 100644 --- a/superset/databases/dao.py +++ b/superset/databases/dao.py @@ -56,7 +56,7 @@ class DatabaseDAO(BaseDAO): @staticmethod def build_db_for_connection_test( server_cert: str, extra: str, impersonate_user: bool, encrypted_extra: str - ) -> Optional[Database]: + ) -> Database: return Database( server_cert=server_cert, extra=extra, diff --git a/superset/utils/log.py b/superset/utils/log.py index 0994535afe..8482871c8c 100644 --- a/superset/utils/log.py +++ b/superset/utils/log.py @@ -104,7 +104,7 @@ class AbstractEventLogger(ABC): def log_with_context( # pylint: disable=too-many-locals self, action: str, - duration: timedelta, + duration: Optional[timedelta] = None, object_ref: Optional[str] = None, log_to_statsd: bool = True, **payload_override: Optional[Dict[str, Any]], @@ -112,6 +112,9 @@ class AbstractEventLogger(ABC): from superset.views.core import get_form_data referrer = request.referrer[:1000] if request.referrer else None + + duration_ms = int(duration.total_seconds() * 1000) if duration else None + try: user_id = g.user.get_id() except Exception as ex: # pylint: disable=broad-except @@ -158,7 +161,7 @@ class AbstractEventLogger(ABC): records=records, dashboard_id=dashboard_id, slice_id=slice_id, - duration_ms=int(duration.total_seconds() * 1000), + duration_ms=duration_ms, referrer=referrer, ) diff --git a/tests/databases/api_tests.py b/tests/databases/api_tests.py index 825ea85437..36d2179aef 100644 --- a/tests/databases/api_tests.py +++ b/tests/databases/api_tests.py @@ -819,7 +819,7 @@ class TestDatabaseApi(SupersetTestCase): self.assertEqual(rv.headers["Content-Type"], "application/json; charset=utf-8") response = json.loads(rv.data.decode("utf-8")) expected_response = { - "message": "Could not load database driver: broken", + "message": "Could not load database driver: BaseEngineSpec", } self.assertEqual(response, expected_response) @@ -834,7 +834,7 @@ class TestDatabaseApi(SupersetTestCase): self.assertEqual(rv.headers["Content-Type"], "application/json; charset=utf-8") response = json.loads(rv.data.decode("utf-8")) expected_response = { - "message": "Could not load database driver: mssql+pymssql", + "message": "Could not load database driver: MssqlEngineSpec", } self.assertEqual(response, expected_response) diff --git a/tests/databases/commands_tests.py b/tests/databases/commands_tests.py index 3b1767fdc2..2e2caaedcc 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 import db, event_logger, 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,75 @@ class TestImportDatabasesCommand(SupersetTestCase): # 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.Database.get_sqla_engine") + @mock.patch( + "superset.databases.commands.test_connection.event_logger.log_with_context" + ) + def test_connection_db_exception(self, mock_event_logger, mock_get_sqla_engine): + """Test to make sure event_logger is called when an exception is raised""" + database = get_example_database() + mock_get_sqla_engine.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 pytest.raises(DatabaseTestConnectionUnexpectedError) as excinfo: + command_without_db_name.run() + assert str(excinfo.value) == ( + "Unexpected error occurred, please check your logs for details" + ) + mock_event_logger.assert_called() + + @mock.patch("superset.databases.dao.Database.get_sqla_engine") + @mock.patch( + "superset.databases.commands.test_connection.event_logger.log_with_context" + ) + def test_connection_superset_security_connection( + self, mock_event_logger, mock_get_sqla_engine + ): + """Test to make sure event_logger is called when security + connection exc is raised""" + database = get_example_database() + mock_get_sqla_engine.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 pytest.raises(DatabaseSecurityUnsafeError) as excinfo: + command_without_db_name.run() + assert str(excinfo.value) == ("Stopped an unsafe database connection") + + mock_event_logger.assert_called() + + @mock.patch("superset.databases.dao.Database.get_sqla_engine") + @mock.patch( + "superset.databases.commands.test_connection.event_logger.log_with_context" + ) + def test_connection_db_api_exc(self, mock_event_logger, mock_get_sqla_engine): + """Test to make sure event_logger is called when DBAPIError is raised""" + database = get_example_database() + mock_get_sqla_engine.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 pytest.raises(DatabaseTestConnectionFailedError) as excinfo: + command_without_db_name.run() + assert str(excinfo.value) == ( + "Connection failed, please check your connection settings" + ) + + mock_event_logger.assert_called()