QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-08-12 11:11:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: cleanup ssh tunnel (#34388)

Browse Source
This commit is contained in:
Beto Dealmeida
2025-12-03 14:26:35 -05:00
committed by GitHub
parent 70aec7fa76
commit c458f99dd4
35 changed files with 304 additions and 1287 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/integration_tests/commands_test.py
+2 -2
View File
@@ -169,7 +169,7 @@ class TestImportAssetsCommand(SupersetTestCase):
assert dashboard.owners == [self.user]
mock_add_permissions.assert_called_with(database, None)
mock_add_permissions.assert_called_with(database)
db.session.delete(dashboard)
db.session.delete(chart)
@@ -214,7 +214,7 @@ class TestImportAssetsCommand(SupersetTestCase):
dataset = chart.table
database = dataset.database
mock_add_permissions.assert_called_with(database, None)
mock_add_permissions.assert_called_with(database)
db.session.delete(dashboard)
db.session.delete(chart)
tests/integration_tests/databases/api_tests.py
+39 -66
View File
@@ -289,23 +289,21 @@ class TestDatabaseApi(SupersetTestCase):
db.session.delete(model)
db.session.commit()
@with_feature_flags(SSH_TUNNELING=True)
@mock.patch(
"superset.commands.database.test_connection.TestConnectionDatabaseCommand.run",
)
@mock.patch("superset.commands.database.create.is_feature_enabled")
@mock.patch("superset.models.core.Database.get_all_catalog_names")
@mock.patch("superset.models.core.Database.get_all_schema_names")
def test_create_database_with_ssh_tunnel(
self,
mock_get_all_schema_names,
mock_get_all_catalog_names,
mock_create_is_feature_enabled,
mock_test_connection_database_command_run,
):
"""
Database API: Test create with SSH Tunnel
"""
mock_create_is_feature_enabled.return_value = True
self.login(ADMIN_USERNAME)
example_db = get_example_database()
if example_db.backend == "sqlite":
@@ -337,23 +335,21 @@ class TestDatabaseApi(SupersetTestCase):
db.session.delete(model)
db.session.commit()
@with_feature_flags(SSH_TUNNELING=True)
@mock.patch(
"superset.commands.database.test_connection.TestConnectionDatabaseCommand.run",
)
@mock.patch("superset.commands.database.create.is_feature_enabled")
@mock.patch("superset.models.core.Database.get_all_catalog_names")
@mock.patch("superset.models.core.Database.get_all_schema_names")
def test_create_database_with_ssh_tunnel_no_port(
self,
mock_get_all_schema_names,
mock_get_all_catalog_names,
mock_create_is_feature_enabled,
mock_test_connection_database_command_run,
):
"""
Database API: Test create with SSH Tunnel
"""
mock_create_is_feature_enabled.return_value = True
self.login(ADMIN_USERNAME)
example_db = get_example_database()
if example_db.backend == "sqlite":
@@ -390,23 +386,21 @@ class TestDatabaseApi(SupersetTestCase):
db.session.commit()
@pytest.mark.skip("buggy")
@with_feature_flags(SSH_TUNNELING=True)
@mock.patch(
"superset.commands.database.test_connection.TestConnectionDatabaseCommand.run",
)
@mock.patch("superset.commands.database.create.is_feature_enabled")
@mock.patch("superset.models.core.Database.get_all_catalog_names")
@mock.patch("superset.models.core.Database.get_all_schema_names")
def test_create_database_with_ssh_tunnel_no_port_no_default(
self,
mock_get_all_schema_names,
mock_get_all_catalog_names,
mock_create_is_feature_enabled,
mock_test_connection_database_command_run,
):
"""
Database API: Test that missing port raises SSHTunnelDatabaseError
"""
mock_create_is_feature_enabled.return_value = True
self.login(username="admin")
example_db = get_example_database()
if example_db.backend == "sqlite":
@@ -435,30 +429,25 @@ class TestDatabaseApi(SupersetTestCase):
== "A database port is required when connecting via SSH Tunnel."
)
@with_feature_flags(SSH_TUNNELING=True)
@mock.patch(
"superset.commands.database.sync_permissions.SyncPermissionsCommand.run",
)
@mock.patch(
"superset.commands.database.test_connection.TestConnectionDatabaseCommand.run",
)
@mock.patch("superset.commands.database.create.is_feature_enabled")
@mock.patch("superset.commands.database.update.is_feature_enabled")
@mock.patch("superset.models.core.Database.get_all_catalog_names")
@mock.patch("superset.models.core.Database.get_all_schema_names")
def test_update_database_with_ssh_tunnel(
self,
mock_get_all_schema_names,
mock_get_all_catalog_names,
mock_update_is_feature_enabled,
mock_create_is_feature_enabled,
mock_test_connection_database_command_run,
mock_sync_perms_command,
):
"""
Database API: Test update Database with SSH Tunnel
"""
mock_create_is_feature_enabled.return_value = True
mock_update_is_feature_enabled.return_value = True
self.login(ADMIN_USERNAME)
example_db = get_example_database()
if example_db.backend == "sqlite":
@@ -500,30 +489,25 @@ class TestDatabaseApi(SupersetTestCase):
db.session.delete(model)
db.session.commit()
@with_feature_flags(SSH_TUNNELING=True)
@mock.patch(
"superset.commands.database.sync_permissions.SyncPermissionsCommand.run",
)
@mock.patch(
"superset.commands.database.test_connection.TestConnectionDatabaseCommand.run",
)
@mock.patch("superset.commands.database.create.is_feature_enabled")
@mock.patch("superset.commands.database.update.is_feature_enabled")
@mock.patch("superset.models.core.Database.get_all_catalog_names")
@mock.patch("superset.models.core.Database.get_all_schema_names")
def test_update_database_with_ssh_tunnel_no_port(
self,
mock_get_all_schema_names,
mock_get_all_catalog_names,
mock_update_is_feature_enabled,
mock_create_is_feature_enabled,
mock_test_connection_database_command_run,
mock_sync_perms_cmmd_run,
):
"""
Database API: Test update Database with SSH Tunnel
"""
mock_create_is_feature_enabled.return_value = True
mock_update_is_feature_enabled.return_value = True
self.login(ADMIN_USERNAME)
example_db = get_example_database()
if example_db.backend == "sqlite":
@@ -568,26 +552,21 @@ class TestDatabaseApi(SupersetTestCase):
db.session.delete(model)
db.session.commit()
@with_feature_flags(SSH_TUNNELING=True)
@mock.patch(
"superset.commands.database.test_connection.TestConnectionDatabaseCommand.run",
)
@mock.patch("superset.commands.database.create.is_feature_enabled")
@mock.patch("superset.commands.database.update.is_feature_enabled")
@mock.patch("superset.models.core.Database.get_all_catalog_names")
@mock.patch("superset.models.core.Database.get_all_schema_names")
def test_update_database_no_port_no_default(
self,
mock_get_all_schema_names,
mock_get_all_catalog_names,
mock_update_is_feature_enabled,
mock_create_is_feature_enabled,
mock_test_connection_database_command_run,
):
"""
Database API: Test that missing port raises SSHTunnelDatabaseError
"""
mock_create_is_feature_enabled.return_value = True
mock_update_is_feature_enabled.return_value = True
self.login(username="admin")
example_db = get_example_database()
if example_db.backend == "sqlite":
@@ -632,33 +611,25 @@ class TestDatabaseApi(SupersetTestCase):
db.session.delete(model)
db.session.commit()
@with_feature_flags(SSH_TUNNELING=True)
@mock.patch(
"superset.commands.database.sync_permissions.SyncPermissionsCommand.run",
)
@mock.patch(
"superset.commands.database.test_connection.TestConnectionDatabaseCommand.run",
)
@mock.patch("superset.commands.database.create.is_feature_enabled")
@mock.patch("superset.commands.database.update.is_feature_enabled")
@mock.patch("superset.commands.database.ssh_tunnel.delete.is_feature_enabled")
@mock.patch("superset.models.core.Database.get_all_catalog_names")
@mock.patch("superset.models.core.Database.get_all_schema_names")
def test_delete_ssh_tunnel(
self,
mock_get_all_schema_names,
mock_get_all_catalog_names,
mock_delete_is_feature_enabled,
mock_update_is_feature_enabled,
mock_create_is_feature_enabled,
mock_test_connection_database_command_run,
mock_sync_perms_command,
):
"""
Database API: Test deleting a SSH tunnel via Database update
"""
mock_create_is_feature_enabled.return_value = True
mock_update_is_feature_enabled.return_value = True
mock_delete_is_feature_enabled.return_value = True
self.login(username="admin")
example_db = get_example_database()
if example_db.backend == "sqlite":
@@ -720,30 +691,25 @@ class TestDatabaseApi(SupersetTestCase):
db.session.delete(model)
db.session.commit()
@with_feature_flags(SSH_TUNNELING=True)
@mock.patch(
"superset.commands.database.sync_permissions.SyncPermissionsCommand.run",
)
@mock.patch(
"superset.commands.database.test_connection.TestConnectionDatabaseCommand.run",
)
@mock.patch("superset.commands.database.create.is_feature_enabled")
@mock.patch("superset.commands.database.update.is_feature_enabled")
@mock.patch("superset.models.core.Database.get_all_catalog_names")
@mock.patch("superset.models.core.Database.get_all_schema_names")
def test_update_ssh_tunnel_via_database_api(
self,
mock_get_all_schema_names,
mock_get_all_catalog_names,
mock_update_is_feature_enabled,
mock_create_is_feature_enabled,
mock_test_connection_database_command_run,
mock_sync_perms_command,
):
"""
Database API: Test update SSH Tunnel via Database API
"""
mock_create_is_feature_enabled.return_value = True
mock_update_is_feature_enabled.return_value = True
self.login(ADMIN_USERNAME)
example_db = get_example_database()
@@ -802,15 +768,14 @@ class TestDatabaseApi(SupersetTestCase):
db.session.delete(model)
db.session.commit()
@with_feature_flags(SSH_TUNNELING=True)
@mock.patch(
"superset.commands.database.test_connection.TestConnectionDatabaseCommand.run",
)
@mock.patch("superset.models.core.Database.get_all_catalog_names")
@mock.patch("superset.models.core.Database.get_all_schema_names")
@mock.patch("superset.commands.database.create.is_feature_enabled")
def test_cascade_delete_ssh_tunnel(
self,
mock_create_is_feature_enabled,
mock_get_all_schema_names,
mock_get_all_catalog_names,
mock_test_connection_database_command_run,
@@ -818,7 +783,6 @@ class TestDatabaseApi(SupersetTestCase):
"""
Database API: SSH Tunnel gets deleted if Database gets deleted
"""
mock_create_is_feature_enabled.return_value = True
self.login(ADMIN_USERNAME)
example_db = get_example_database()
if example_db.backend == "sqlite":
@@ -837,6 +801,7 @@ class TestDatabaseApi(SupersetTestCase):
uri = "api/v1/database/"
rv = self.client.post(uri, json=database_data)
response = json.loads(rv.data.decode("utf-8"))
print(rv.text)
assert rv.status_code == 201
model_ssh_tunnel = (
db.session.query(SSHTunnel)
@@ -855,10 +820,10 @@ class TestDatabaseApi(SupersetTestCase):
)
assert model_ssh_tunnel is None
@with_feature_flags(SSH_TUNNELING=True)
@mock.patch(
"superset.commands.database.test_connection.TestConnectionDatabaseCommand.run",
)
@mock.patch("superset.commands.database.create.is_feature_enabled")
@mock.patch("superset.models.core.Database.get_all_catalog_names")
@mock.patch("superset.models.core.Database.get_all_schema_names")
@mock.patch("superset.extensions.db.session.rollback")
@@ -866,14 +831,12 @@ class TestDatabaseApi(SupersetTestCase):
self,
mock_get_all_schema_names,
mock_get_all_catalog_names,
mock_create_is_feature_enabled,
mock_test_connection_database_command_run,
mock_rollback,
):
"""
Database API: Test rollback is called if SSH Tunnel creation fails
"""
mock_create_is_feature_enabled.return_value = True
self.login(ADMIN_USERNAME)
example_db = get_example_database()
if example_db.backend == "sqlite":
@@ -897,7 +860,11 @@ class TestDatabaseApi(SupersetTestCase):
"sqlalchemy_uri": example_db.sqlalchemy_uri_decrypted,
"ssh_tunnel": ssh_tunnel_properties,
}
fail_message = {"message": "SSH Tunnel parameters are invalid."}
fail_message = {
"message": {
"ssh_tunnel": {"password": "Either password or private_key is required"}
}
}
uri = "api/v1/database/"
rv = self.client.post(uri, json=database_data)
@@ -912,9 +879,6 @@ class TestDatabaseApi(SupersetTestCase):
assert model_ssh_tunnel is None
assert response == fail_message
# Check that rollback was called
mock_rollback.assert_called()
# Clean up any database that might have been created
created_db = (
db.session.query(Database)
@@ -925,23 +889,21 @@ class TestDatabaseApi(SupersetTestCase):
db.session.delete(created_db)
db.session.commit()
@with_feature_flags(SSH_TUNNELING=True)
@mock.patch(
"superset.commands.database.test_connection.TestConnectionDatabaseCommand.run",
)
@mock.patch("superset.commands.database.create.is_feature_enabled")
@mock.patch("superset.models.core.Database.get_all_catalog_names")
@mock.patch("superset.models.core.Database.get_all_schema_names")
def test_get_database_returns_related_ssh_tunnel(
self,
mock_get_all_schema_names,
mock_get_all_catalog_names,
mock_create_is_feature_enabled,
mock_test_connection_database_command_run,
):
"""
Database API: Test GET Database returns its related SSH Tunnel
"""
mock_create_is_feature_enabled.return_value = True
self.login(ADMIN_USERNAME)
example_db = get_example_database()
if example_db.backend == "sqlite":
@@ -1091,7 +1053,7 @@ class TestDatabaseApi(SupersetTestCase):
]
}
}
assert rv.status_code == 400
assert rv.status_code == 422
def test_create_database_no_configuration_method(self):
"""
@@ -1146,7 +1108,7 @@ class TestDatabaseApi(SupersetTestCase):
rv = self.client.post(uri, json=database_data)
response = json.loads(rv.data.decode("utf-8"))
expected_response = {"message": {"server_cert": ["Invalid certificate"]}}
assert rv.status_code == 400
assert rv.status_code == 422
assert response == expected_response
def test_create_database_json_validate(self):
@@ -1181,7 +1143,7 @@ class TestDatabaseApi(SupersetTestCase):
],
}
}
assert rv.status_code == 400
assert rv.status_code == 422
assert response == expected_response
def test_create_database_extra_metadata_validate(self):
@@ -1217,7 +1179,7 @@ class TestDatabaseApi(SupersetTestCase):
]
}
}
assert rv.status_code == 400
assert rv.status_code == 422
assert response == expected_response
def test_create_database_unique_validate(self):
@@ -1260,7 +1222,7 @@ class TestDatabaseApi(SupersetTestCase):
uri = "api/v1/database/"
rv = self.client.post(uri, json=database_data)
response = json.loads(rv.data.decode("utf-8"))
assert rv.status_code == 400
assert rv.status_code == 422
assert "Invalid connection string" in response["message"]["sqlalchemy_uri"][0]
@with_config({"PREVENT_UNSAFE_DB_CONNECTIONS": True})
@@ -1287,7 +1249,7 @@ class TestDatabaseApi(SupersetTestCase):
}
}
assert response_data == expected_response
assert response.status_code == 400
assert response.status_code == 422
def test_create_database_conn_fail(self):
"""
@@ -2343,7 +2305,7 @@ class TestDatabaseApi(SupersetTestCase):
expected_response = {
"errors": [
{
"message": "Could not load database driver: BaseEngineSpec",
"message": "Could not load database driver for: broken",
"error_type": "GENERIC_COMMAND_ERROR",
"level": "warning",
"extra": {
@@ -2372,7 +2334,7 @@ class TestDatabaseApi(SupersetTestCase):
expected_response = {
"errors": [
{
"message": "Could not load database driver: MssqlEngineSpec",
"message": "Could not load database driver for: mssql",
"error_type": "GENERIC_COMMAND_ERROR",
"level": "warning",
"extra": {
@@ -2991,16 +2953,27 @@ class TestDatabaseApi(SupersetTestCase):
assert response == {
"errors": [
{
"message": "Must provide credentials for the SSH Tunnel",
"message": (
"Error importing database: databases/database_1.yaml: "
"{'ssh_tunnel': {'password': 'Either password or private_key "
"is required'}}"
),
"error_type": "GENERIC_COMMAND_ERROR",
"level": "warning",
"extra": {
"databases/database_1.yaml": {
"ssh_tunnel": {
"password": (
"Either password or private_key is required"
),
}
},
"issue_codes": [
{
"code": 1010,
"message": (
"Issue 1010 - Superset encountered an "
"error while running a command."
"Issue 1010 - Superset encountered an error while "
"running a command."
),
}
],
tests/integration_tests/databases/commands_tests.py
+10 -5
View File
@@ -813,7 +813,10 @@ class TestImportDatabasesCommand(SupersetTestCase):
command = ImportDatabasesCommand(contents)
with pytest.raises(CommandInvalidError) as excinfo:
command.run()
assert str(excinfo.value) == "Must provide credentials for the SSH Tunnel"
assert str(excinfo.value) == (
"Error importing database: databases/imported_database.yaml: "
"{'ssh_tunnel': {'password': 'Either password or private_key is required'}}"
)
@patch("superset.databases.schemas.is_feature_enabled")
@patch("superset.commands.database.importers.v1.utils.add_permissions")
@@ -858,10 +861,12 @@ class TestImportDatabasesCommand(SupersetTestCase):
assert str(excinfo.value).startswith("Error importing database")
assert excinfo.value.normalized_messages() == {
"databases/imported_database.yaml": {
"_schema": [
"Must provide a private key for the ssh tunnel",
"Must provide a private key password for the ssh tunnel",
]
"ssh_tunnel": {
"password": "Either password or private_key is required",
"private_key": (
"private_key is required when private_key_password is provided"
),
}
}
}
tests/integration_tests/databases/ssh_tunnel/commands/__init__.py
-16
View File
@@ -1,16 +0,0 @@
# 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.
tests/integration_tests/databases/ssh_tunnel/commands/commands_tests.py
-58
View File
@@ -1,58 +0,0 @@
# 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 unittest import mock
import pytest
from superset import security_manager
from superset.commands.database.ssh_tunnel.delete import DeleteSSHTunnelCommand
from superset.commands.database.ssh_tunnel.exceptions import (
SSHTunnelNotFoundError,
)
from superset.commands.database.ssh_tunnel.update import UpdateSSHTunnelCommand
from tests.integration_tests.base_tests import SupersetTestCase
class TestUpdateSSHTunnelCommand(SupersetTestCase):
@mock.patch("superset.utils.core.g")
def test_update_ssh_tunnel_not_found(self, mock_g):
mock_g.user = security_manager.find_user("admin")
# We have not created a SSH Tunnel yet so id = 1 is invalid
command = UpdateSSHTunnelCommand(
1,
{
"server_address": "127.0.0.1",
"server_port": 5432,
"username": "test_user",
},
)
with pytest.raises(SSHTunnelNotFoundError) as excinfo:
command.run()
assert str(excinfo.value) == ("SSH Tunnel not found.")
class TestDeleteSSHTunnelCommand(SupersetTestCase):
@mock.patch("superset.utils.core.g")
@mock.patch("superset.commands.database.ssh_tunnel.delete.is_feature_enabled")
def test_delete_ssh_tunnel_not_found(self, mock_g, mock_delete_is_feature_enabled):
mock_g.user = security_manager.find_user("admin")
mock_delete_is_feature_enabled.return_value = True
# We have not created a SSH Tunnel yet so id = 1 is invalid
command = DeleteSSHTunnelCommand(1)
with pytest.raises(SSHTunnelNotFoundError) as excinfo:
command.run()
assert str(excinfo.value) == ("SSH Tunnel not found.")
tests/integration_tests/tags/api_tests.py
+4 -3
View File
@@ -850,9 +850,10 @@ class TestTagApi(InsertChartMixin, SupersetTestCase):
rv = self.client.post(uri, json=data)
# Should succeed without SQL errors (201 for created or 200 for success)
assert rv.status_code in [200, 201], (
f"Tag creation should succeed, got {rv.status_code}"
)
assert rv.status_code in [
200,
201,
], f"Tag creation should succeed, got {rv.status_code}"
# Query the database to verify the tag was created correctly
created_tag = db.session.query(Tag).filter_by(name=tag_name).first()