diff --git a/superset/commands/database/utils.py b/superset/commands/database/utils.py index e850a9e89ee..0c25f8173c2 100644 --- a/superset/commands/database/utils.py +++ b/superset/commands/database/utils.py @@ -47,7 +47,8 @@ def ping(engine: Engine) -> bool: except (sqlite3.ProgrammingError, RuntimeError): # SQLite can't run on a separate thread, so ``utils.timeout`` fails # RuntimeError catches the equivalent error from duckdb. - return engine.dialect.do_ping(engine) + with closing(engine.raw_connection()) as conn: + return engine.dialect.do_ping(conn) def add_permissions(database: Database) -> None: diff --git a/tests/unit_tests/commands/databases/utils_test.py b/tests/unit_tests/commands/databases/utils_test.py index 04512e47f4f..a8fb3a36e54 100644 --- a/tests/unit_tests/commands/databases/utils_test.py +++ b/tests/unit_tests/commands/databases/utils_test.py @@ -77,7 +77,7 @@ def test_ping_sqlite_exception(mocker: MockerFixture, mock_engine: MockerFixture assert result is True mock_dialect.do_ping.assert_has_calls( - [mocker.call(mock_connection), mocker.call(mock_engine)] + [mocker.call(mock_connection), mocker.call(mock_connection)] ) @@ -85,7 +85,7 @@ def test_ping_runtime_exception(mocker: MockerFixture, mock_engine: MockerFixtur """ Test the ``ping`` method when a RuntimeError is raised. """ - mock_engine, _, mock_dialect = mock_engine + mock_engine, mock_connection, mock_dialect = mock_engine mock_timeout = mocker.patch("superset.commands.database.utils.timeout") mock_timeout.side_effect = RuntimeError("timeout") mock_dialect.do_ping.return_value = True @@ -93,7 +93,7 @@ def test_ping_runtime_exception(mocker: MockerFixture, mock_engine: MockerFixtur result = ping(mock_engine) assert result is True - mock_dialect.do_ping.assert_called_once_with(mock_engine) + mock_dialect.do_ping.assert_called_once_with(mock_connection) @pytest.fixture