Skip to content

Commit

Permalink
fix: dbmodal test connection error timeout (#20068)
Browse files Browse the repository at this point in the history
* fix: check for connect on ping

* clean up merge

* fix merge

* precommit
  • Loading branch information
pkdotson committed May 17, 2022
1 parent 9cdaa28 commit 5111011
Showing 1 changed file with 37 additions and 30 deletions.
67 changes: 37 additions & 30 deletions superset/databases/commands/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from flask_appbuilder.security.sqla.models import User
from flask_babel import gettext as _
from func_timeout import func_timeout, FunctionTimedOut
from sqlalchemy.engine import Engine
from sqlalchemy.exc import DBAPIError, NoSuchModuleError

from superset.commands.base import BaseCommand
Expand Down Expand Up @@ -82,36 +83,41 @@ def run(self) -> None:
action="test_connection_attempt",
engine=database.db_engine_spec.__name__,
)
with closing(engine.raw_connection()) as conn:
try:
alive = func_timeout(
int(
app.config[
"TEST_DATABASE_CONNECTION_TIMEOUT"
].total_seconds()
),
engine.dialect.do_ping,
args=(conn,),
)
except (sqlite3.ProgrammingError, RuntimeError):
# SQLite can't run on a separate thread, so ``func_timeout`` fails
# RuntimeError catches the equivalent error from duckdb.
alive = engine.dialect.do_ping(conn)
except FunctionTimedOut as ex:
raise SupersetTimeoutException(
error_type=SupersetErrorType.CONNECTION_DATABASE_TIMEOUT,
message=(
"Please check your connection details and database "
"settings, and ensure that your database is accepting "
"connections, then try connecting again."
),
level=ErrorLevel.ERROR,
extra={"sqlalchemy_uri": database.sqlalchemy_uri},
) from ex
except Exception: # pylint: disable=broad-except
alive = False
if not alive:
raise DBAPIError(None, None, None)

def ping(engine: Engine) -> bool:
with closing(engine.raw_connection()) as conn:
return engine.dialect.do_ping(conn)

try:
alive = func_timeout(
int(
app.config[
"TEST_DATABASE_CONNECTION_TIMEOUT"
].total_seconds()
),
ping,
args=(engine,),
)

except (sqlite3.ProgrammingError, RuntimeError):
# SQLite can't run on a separate thread, so ``func_timeout`` fails
# RuntimeError catches the equivalent error from duckdb.
alive = engine.dialect.do_ping(engine)
except FunctionTimedOut as ex:
raise SupersetTimeoutException(
error_type=SupersetErrorType.CONNECTION_DATABASE_TIMEOUT,
message=(
"Please check your connection details and database settings, "
"and ensure that your database is accepting connections, "
"then try connecting again."
),
level=ErrorLevel.ERROR,
extra={"sqlalchemy_uri": database.sqlalchemy_uri},
) from ex
except Exception: # pylint: disable=broad-except
alive = False
if not alive:
raise DBAPIError(None, None, None)

# Log succesful connection test with engine
event_logger.log_with_context(
Expand Down Expand Up @@ -144,6 +150,7 @@ def run(self) -> None:
)
raise DatabaseSecurityUnsafeError(message=str(ex)) from ex
except SupersetTimeoutException as ex:

event_logger.log_with_context(
action=f"test_connection_error.{ex.__class__.__name__}",
engine=database.db_engine_spec.__name__,
Expand Down

0 comments on commit 5111011

Please sign in to comment.