Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Add logging for ssh tunneling test_connection attempts #22625

Merged
merged 12 commits into from
Jan 12, 2023
38 changes: 30 additions & 8 deletions superset/databases/commands/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,17 @@
logger = logging.getLogger(__name__)


def get_log_connection_action(
action: str, ssh_tunnel: Optional[Any], exc: Optional[Exception] = None
) -> str:
action_modified = action
if exc:
action_modified += f".{exc.__class__.__name__}"
if ssh_tunnel:
action_modified += ".ssh_tunnel"
return action_modified


class TestConnectionDatabaseCommand(BaseCommand):
def __init__(self, data: Dict[str, Any]):
self._properties = data.copy()
Expand All @@ -59,6 +70,7 @@ def run(self) -> None: # pylint: disable=too-many-statements
uri = self._properties.get("sqlalchemy_uri", "")
if self._model and uri == self._model.safe_sqlalchemy_uri():
uri = self._model.sqlalchemy_uri_decrypted
ssh_tunnel = self._properties.get("ssh_tunnel")

# context for error messages
url = make_url_safe(uri)
Expand Down Expand Up @@ -94,7 +106,7 @@ def run(self) -> None: # pylint: disable=too-many-statements
database.db_engine_spec.mutate_db_for_connection_test(database)

# Generate tunnel if present in the properties
if ssh_tunnel := self._properties.get("ssh_tunnel"):
if ssh_tunnel:
# If there's an existing tunnel for that DB we need to use the stored
# password, private_key and private_key_password instead
if ssh_tunnel_id := ssh_tunnel.pop("id", None):
Expand All @@ -105,7 +117,7 @@ def run(self) -> None: # pylint: disable=too-many-statements
ssh_tunnel = SSHTunnel(**ssh_tunnel)

event_logger.log_with_context(
action="test_connection_attempt",
action=get_log_connection_action("test_connection_attempt", ssh_tunnel),
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are there tests for these?

engine=database.db_engine_spec.__name__,
)

Expand Down Expand Up @@ -147,13 +159,15 @@ def ping(engine: Engine) -> bool:

# Log succesful connection test with engine
event_logger.log_with_context(
action="test_connection_success",
action=get_log_connection_action("test_connection_success", ssh_tunnel),
engine=database.db_engine_spec.__name__,
)

except (NoSuchModuleError, ModuleNotFoundError) as ex:
event_logger.log_with_context(
action=f"test_connection_error.{ex.__class__.__name__}",
action=get_log_connection_action(
"test_connection_error", ssh_tunnel, ex
),
engine=database.db_engine_spec.__name__,
)
raise DatabaseTestConnectionDriverError(
Expand All @@ -163,29 +177,37 @@ def ping(engine: Engine) -> bool:
) from ex
except DBAPIError as ex:
event_logger.log_with_context(
action=f"test_connection_error.{ex.__class__.__name__}",
action=get_log_connection_action(
"test_connection_error", ssh_tunnel, ex
),
engine=database.db_engine_spec.__name__,
)
# check for custom errors (wrong username, wrong password, etc)
errors = database.db_engine_spec.extract_errors(ex, context)
raise SupersetErrorsException(errors) from ex
except SupersetSecurityException as ex:
event_logger.log_with_context(
action=f"test_connection_error.{ex.__class__.__name__}",
action=get_log_connection_action(
"test_connection_error", ssh_tunnel, ex
),
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__}",
action=get_log_connection_action(
"test_connection_error", ssh_tunnel, ex
),
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__}",
action=get_log_connection_action(
"test_connection_error", ssh_tunnel, ex
),
engine=database.db_engine_spec.__name__,
)
errors = database.db_engine_spec.extract_errors(ex, context)
Expand Down
2 changes: 1 addition & 1 deletion superset/db_engine_specs/impala.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ def execute(
cls,
cursor: Any,
query: str,
**kwargs: Any, # pylint: disable=unused-argument
**kwargs: Any,
) -> None:
try:
cursor.execute_async(query)
Expand Down
2 changes: 1 addition & 1 deletion superset/models/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -1326,7 +1326,7 @@ def get_sqla_col(self, col: Dict[str, Any]) -> Column:
col = sa.column(label, type_=col_type)
return self.make_sqla_column_compatible(col, label)

def get_sqla_query( # pylint: disable=too-many-arguments,too-many-locals,too-many-branches,too-many-statements,unused-argument
def get_sqla_query( # pylint: disable=too-many-arguments,too-many-locals,too-many-branches,too-many-statements
self,
apply_fetch_values_predicate: bool = False,
columns: Optional[List[Column]] = None,
Expand Down
32 changes: 32 additions & 0 deletions tests/unit_tests/databases/commands/test_connection_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.

from parameterized import parameterized

from superset.databases.commands.test_connection import get_log_connection_action
from superset.databases.ssh_tunnel.models import SSHTunnel


@parameterized.expand(
[
("foo", None, None, "foo"),
("foo", SSHTunnel, None, "foo.ssh_tunnel"),
("foo", SSHTunnel, Exception("oops"), "foo.Exception.ssh_tunnel"),
],
)
def test_get_log_connection_action(action, tunnel, exc, expected_result):
assert expected_result == get_log_connection_action(action, tunnel, exc)