Skip to content

Commit

Permalink
feat: shorter timeout on test_connection (#18001)
Browse files Browse the repository at this point in the history
* feat: shorter timeout on test_connection

* pip-compile-multi --no-upgrade

* Fix for SQLite

* Return 408

* Add test
  • Loading branch information
betodealmeida authored Jan 12, 2022
1 parent 6bd4dd2 commit 51090c3
Show file tree
Hide file tree
Showing 8 changed files with 89 additions and 3 deletions.
2 changes: 2 additions & 0 deletions requirements/base.txt
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,8 @@ flask-wtf==0.14.3
# via
# apache-superset
# flask-appbuilder
func-timeout==4.3.5
# via apache-superset
geographiclib==1.52
# via geopy
geopy==2.2.0
Expand Down
1 change: 1 addition & 0 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@ def get_git_sha() -> str:
"flask-talisman",
"flask-migrate",
"flask-wtf",
"func_timeout",
"geopy",
"graphlib-backport",
"gunicorn>=20.1.0",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,12 +29,14 @@ const SqlAlchemyTab = ({
testConnection,
conf,
isEditMode = false,
testInProgress = false,
}: {
db: DatabaseObject | null;
onInputChange: EventHandler<ChangeEvent<HTMLInputElement>>;
testConnection: EventHandler<MouseEvent<HTMLElement>>;
conf: { SQLALCHEMY_DOCS_URL: string; SQLALCHEMY_DISPLAY_TEXT: string };
isEditMode?: boolean;
testInProgress?: boolean;
}) => (
<>
<StyledInputContainer>
Expand Down Expand Up @@ -88,6 +90,7 @@ const SqlAlchemyTab = ({
</StyledInputContainer>
<Button
onClick={testConnection}
disabled={testInProgress}
cta
buttonStyle="link"
css={(theme: SupersetTheme) => wideButton(theme)}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -440,6 +440,7 @@ const DatabaseModal: FunctionComponent<DatabaseModalProps> = ({
const [dbName, setDbName] = useState('');
const [editNewDb, setEditNewDb] = useState<boolean>(false);
const [isLoading, setLoading] = useState<boolean>(false);
const [testInProgress, setTestInProgress] = useState<boolean>(false);
const conf = useCommonConf();
const dbImages = getDatabaseImages();
const connectionAlert = getConnectionAlert();
Expand Down Expand Up @@ -494,7 +495,18 @@ const DatabaseModal: FunctionComponent<DatabaseModalProps> = ({
encrypted_extra: db?.encrypted_extra || '',
server_cert: db?.server_cert || undefined,
};
testDatabaseConnection(connection, addDangerToast, addSuccessToast);
setTestInProgress(true);
testDatabaseConnection(
connection,
(errorMsg: string) => {
setTestInProgress(false);
addDangerToast(errorMsg);
},
(errorMsg: string) => {
setTestInProgress(false);
addSuccessToast(errorMsg);
},
);
};

const onClose = () => {
Expand Down Expand Up @@ -1047,6 +1059,7 @@ const DatabaseModal: FunctionComponent<DatabaseModalProps> = ({
conf={conf}
testConnection={testConnection}
isEditMode={isEditMode}
testInProgress={testInProgress}
/>
{isDynamic(db?.backend || db?.engine) && !isEditMode && (
<div css={(theme: SupersetTheme) => infoTooltip(theme)}>
Expand Down
5 changes: 5 additions & 0 deletions superset/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -1171,6 +1171,11 @@ def SQL_QUERY_MUTATOR( # pylint: disable=invalid-name,unused-argument
"SQLite",
# etc.
]
# When adding a new database we try to connect to it. Depending on which parameters are
# incorrect this could take a couple minutes, until the SQLAlchemy driver pinging the
# database times out. Instead of relying on the driver timeout we can specify a shorter
# one here.
TEST_DATABASE_CONNECTION_TIMEOUT = timedelta(seconds=30)

# Do you want Talisman enabled?
TALISMAN_ENABLED = False
Expand Down
35 changes: 34 additions & 1 deletion superset/databases/commands/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,14 @@
# specific language governing permissions and limitations
# under the License.
import logging
import sqlite3
from contextlib import closing
from typing import Any, Dict, Optional

from flask import current_app as app
from flask_appbuilder.security.sqla.models import User
from flask_babel import gettext as _
from func_timeout import func_timeout, FunctionTimedOut
from sqlalchemy.engine.url import make_url
from sqlalchemy.exc import DBAPIError, NoSuchModuleError

Expand All @@ -31,7 +34,8 @@
DatabaseTestConnectionUnexpectedError,
)
from superset.databases.dao import DatabaseDAO
from superset.exceptions import SupersetSecurityException
from superset.errors import ErrorLevel, SupersetErrorType
from superset.exceptions import SupersetSecurityException, SupersetTimeoutException
from superset.extensions import event_logger
from superset.models.core import Database

Expand Down Expand Up @@ -78,7 +82,29 @@ def run(self) -> None:
)
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:
# SQLite can't run on a separate thread, so ``func_timeout`` fails
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:
Expand Down Expand Up @@ -114,6 +140,13 @@ def run(self) -> None:
engine=database.db_engine_spec.__name__,
)
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__,
)
# bubble up the exception to return a 408
raise ex
except Exception as ex:
event_logger.log_with_context(
action=f"test_connection_error.{ex.__class__.__name__}",
Expand Down
2 changes: 2 additions & 0 deletions superset/errors.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ class SupersetErrorType(str, Enum):
CONNECTION_MISSING_PARAMETERS_ERROR = "CONNECTION_MISSING_PARAMETERS_ERROR"
OBJECT_DOES_NOT_EXIST_ERROR = "OBJECT_DOES_NOT_EXIST_ERROR"
SYNTAX_ERROR = "SYNTAX_ERROR"
CONNECTION_DATABASE_TIMEOUT = "CONNECTION_DATABASE_TIMEOUT"

# Viz errors
VIZ_GET_DF_ERROR = "VIZ_GET_DF_ERROR"
Expand Down Expand Up @@ -172,6 +173,7 @@ class SupersetErrorType(str, Enum):
SupersetErrorType.CONNECTION_INVALID_PORT_ERROR: [1034],
SupersetErrorType.ASYNC_WORKERS_ERROR: [1035],
SupersetErrorType.DATABASE_NOT_FOUND_ERROR: [1011, 1036],
SupersetErrorType.CONNECTION_DATABASE_TIMEOUT: [1001, 1009],
}


Expand Down
29 changes: 28 additions & 1 deletion tests/integration_tests/databases/commands_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

import pytest
import yaml
from func_timeout import FunctionTimedOut
from sqlalchemy.exc import DBAPIError

from superset import db, event_logger, security_manager
Expand All @@ -38,7 +39,11 @@
from superset.databases.commands.validate import ValidateDatabaseParametersCommand
from superset.databases.schemas import DatabaseTestConnectionSchema
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import SupersetErrorsException, SupersetSecurityException
from superset.exceptions import (
SupersetErrorsException,
SupersetSecurityException,
SupersetTimeoutException,
)
from superset.models.core import Database
from superset.utils.core import backend, get_example_database
from tests.integration_tests.base_tests import SupersetTestCase
Expand Down Expand Up @@ -612,6 +617,28 @@ def test_connection_do_ping_exception(
== SupersetErrorType.GENERIC_DB_ENGINE_ERROR
)

@mock.patch("superset.databases.commands.test_connection.func_timeout")
@mock.patch(
"superset.databases.commands.test_connection.event_logger.log_with_context"
)
def test_connection_do_ping_timeout(self, mock_event_logger, mock_func_timeout):
"""Test to make sure do_ping exceptions gets captured"""
database = get_example_database()
mock_func_timeout.side_effect = FunctionTimedOut("Time out")
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(SupersetTimeoutException) as excinfo:
command_without_db_name.run()
assert excinfo.value.status == 408
assert (
excinfo.value.error.error_type
== SupersetErrorType.CONNECTION_DATABASE_TIMEOUT
)

@mock.patch("superset.databases.dao.Database.get_sqla_engine")
@mock.patch(
"superset.databases.commands.test_connection.event_logger.log_with_context"
Expand Down

0 comments on commit 51090c3

Please sign in to comment.