QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-04-07 18:35:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: default ports for SSH tunnel (#32403)

Browse Source
This commit is contained in:
Beto Dealmeida
2025-02-27 10:59:48 -05:00
committed by GitHub
parent 74733ae310
commit f4105e9ed2
7 changed files with 228 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
superset/commands/database/ssh_tunnel/create.py
View File

@@ -33,6 +33,7 @@ from superset.databases.utils import make_url_safe
from superset.extensions import event_logger
from superset.models.core import Database
from superset.utils.decorators import on_error, transaction
from superset.utils.ssh_tunnel import get_default_port
logger = logging.getLogger(__name__)
@@ -72,7 +73,9 @@ class CreateSSHTunnelCommand(BaseCommand):
"private_key_password"
)
url = make_url_safe(self._database.sqlalchemy_uri)
if not url.port:
backend = url.get_backend_name()
port = url.port or get_default_port(backend)
if not port:
raise SSHTunnelDatabasePortError()
if not server_address:
exceptions.append(SSHTunnelRequiredFieldValidationError("server_address"))

5
superset/commands/database/ssh_tunnel/update.py
View File

@@ -32,6 +32,7 @@ from superset.daos.database import SSHTunnelDAO
from superset.databases.ssh_tunnel.models import SSHTunnel
from superset.databases.utils import make_url_safe
from superset.utils.decorators import on_error, transaction
from superset.utils.ssh_tunnel import get_default_port
logger = logging.getLogger(__name__)
@@ -75,5 +76,7 @@ class UpdateSSHTunnelCommand(BaseCommand):
raise SSHTunnelInvalidError(
exceptions=[SSHTunnelRequiredFieldValidationError("private_key")]
)
if not url.port:
backend = url.get_backend_name()
port = url.port or get_default_port(backend)
if not port:
raise SSHTunnelDatabasePortError()

5
superset/extensions/ssh.py
View File

@@ -52,11 +52,14 @@ class SSHManager:
ssh_tunnel: "SSHTunnel",
sqlalchemy_database_uri: str,
) -> sshtunnel.SSHTunnelForwarder:
from superset.utils.ssh_tunnel import get_default_port
url = make_url_safe(sqlalchemy_database_uri)
backend = url.get_backend_name()
params = {
"ssh_address_or_host": (ssh_tunnel.server_address, ssh_tunnel.server_port),
"ssh_username": ssh_tunnel.username,
"remote_bind_address": (url.host, url.port),
"remote_bind_address": (url.host, url.port or get_default_port(backend)),
"local_bind_address": (self.local_bind_address,),
"debug_level": logging.getLogger("flask_appbuilder").level,
}

14
superset/utils/ssh_tunnel.py
View File

@@ -20,6 +20,13 @@ from typing import Any
from superset.constants import PASSWORD_MASK
from superset.databases.ssh_tunnel.models import SSHTunnel
DEFAULT_PORTS: dict[str, int] = {
"postgresql": 5432,
"mysql": 3306,
"oracle": 1521,
"mssql": 1433,
}
def mask_password_info(ssh_tunnel: dict[str, Any]) -> dict[str, Any]:
if ssh_tunnel.pop("password", None) is not None:
@@ -41,3 +48,10 @@ def unmask_password_info(
if ssh_tunnel.get("private_key_password") == PASSWORD_MASK:
ssh_tunnel["private_key_password"] = model.private_key_password
return ssh_tunnel
def get_default_port(backend: str) -> int | None:
"""
Get the default port for the given backend.
"""
return DEFAULT_PORTS.get(backend)

131
tests/integration_tests/databases/api_tests.py
View File

@@ -345,7 +345,59 @@ class TestDatabaseApi(SupersetTestCase):
@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_missing_port_raises_error(
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":
return
modified_sqlalchemy_uri = "postgresql://foo:bar@localhost/test-db"
ssh_tunnel_properties = {
"server_address": "123.132.123.1",
"server_port": 8080,
"username": "foo",
"password": "bar",
}
database_data_with_ssh_tunnel = {
"database_name": "test-db-with-ssh-tunnel",
"sqlalchemy_uri": modified_sqlalchemy_uri,
"ssh_tunnel": ssh_tunnel_properties,
}
uri = "api/v1/database/"
rv = self.client.post(uri, json=database_data_with_ssh_tunnel)
response = json.loads(rv.data.decode("utf-8"))
assert rv.status_code == 201
model_ssh_tunnel = (
db.session.query(SSHTunnel)
.filter(SSHTunnel.database_id == response.get("id"))
.one()
)
assert response.get("result")["ssh_tunnel"]["password"] == "XXXXXXXXXX" # noqa: S105
assert model_ssh_tunnel.database_id == response.get("id")
# Cleanup
model = db.session.query(Database).get(response.get("id"))
db.session.delete(model)
db.session.commit()
@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,
@@ -361,7 +413,7 @@ class TestDatabaseApi(SupersetTestCase):
if example_db.backend == "sqlite":
return
modified_sqlalchemy_uri = "postgresql://foo:bar@localhost/test-db"
modified_sqlalchemy_uri = "weird+db://foo:bar@localhost/test-db"
ssh_tunnel_properties = {
"server_address": "123.132.123.1",
@@ -369,13 +421,6 @@ class TestDatabaseApi(SupersetTestCase):
"username": "foo",
"password": "bar",
}
database_data_with_ssh_tunnel = {
"database_name": "test-db-with-ssh-tunnel",
"sqlalchemy_uri": modified_sqlalchemy_uri,
"ssh_tunnel": ssh_tunnel_properties,
}
database_data_with_ssh_tunnel = {
"database_name": "test-db-with-ssh-tunnel",
"sqlalchemy_uri": modified_sqlalchemy_uri,
@@ -459,7 +504,71 @@ class TestDatabaseApi(SupersetTestCase):
@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_missing_port_raises_error(
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,
):
"""
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":
return
modified_sqlalchemy_uri = "postgresql://foo:bar@localhost/test-db"
ssh_tunnel_properties = {
"server_address": "123.132.123.1",
"server_port": 8080,
"username": "foo",
"password": "bar",
}
database_data = {
"database_name": "test-db-with-ssh-tunnel",
"sqlalchemy_uri": example_db.sqlalchemy_uri_decrypted,
}
database_data_with_ssh_tunnel = {
"database_name": "test-db-with-ssh-tunnel",
"sqlalchemy_uri": modified_sqlalchemy_uri,
"ssh_tunnel": ssh_tunnel_properties,
}
uri = "api/v1/database/"
rv = self.client.post(uri, json=database_data)
response = json.loads(rv.data.decode("utf-8"))
assert rv.status_code == 201
uri = "api/v1/database/{}".format(response.get("id"))
rv = self.client.put(uri, json=database_data_with_ssh_tunnel)
response_update = json.loads(rv.data.decode("utf-8"))
assert rv.status_code == 200
model_ssh_tunnel = (
db.session.query(SSHTunnel)
.filter(SSHTunnel.database_id == response_update.get("id"))
.one()
)
assert model_ssh_tunnel.database_id == response_update.get("id")
# Cleanup
model = db.session.query(Database).get(response.get("id"))
db.session.delete(model)
db.session.commit()
@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,
@@ -477,7 +586,7 @@ class TestDatabaseApi(SupersetTestCase):
if example_db.backend == "sqlite":
return
modified_sqlalchemy_uri = "postgresql://foo:bar@localhost/test-db"
modified_sqlalchemy_uri = "weird+db://foo:bar@localhost/test-db"
ssh_tunnel_properties = {
"server_address": "123.132.123.1",

57
tests/unit_tests/databases/ssh_tunnel/commands/create_test.py
View File

@@ -17,6 +17,7 @@
import pytest
from sqlalchemy.orm.session import Session
from superset.commands.database.ssh_tunnel.exceptions import (
SSHTunnelDatabasePortError,
@@ -24,13 +25,16 @@ from superset.commands.database.ssh_tunnel.exceptions import (
)
def test_create_ssh_tunnel_command() -> None:
def test_create_ssh_tunnel_command(session: Session) -> None:
from superset import db
from superset.commands.database.ssh_tunnel.create import CreateSSHTunnelCommand
from superset.databases.ssh_tunnel.models import SSHTunnel
from superset.models.core import Database
engine = db.session.get_bind()
Database.metadata.create_all(engine) # pylint: disable=no-member
database = Database(
id=1,
database_name="my_database",
sqlalchemy_uri="postgresql://u:p@localhost:5432/db",
)
@@ -49,12 +53,15 @@ def test_create_ssh_tunnel_command() -> None:
assert isinstance(result, SSHTunnel)
def test_create_ssh_tunnel_command_invalid_params() -> None:
def test_create_ssh_tunnel_command_invalid_params(session: Session) -> None:
from superset import db
from superset.commands.database.ssh_tunnel.create import CreateSSHTunnelCommand
from superset.models.core import Database
engine = db.session.get_bind()
Database.metadata.create_all(engine) # pylint: disable=no-member
database = Database(
id=1,
database_name="my_database",
sqlalchemy_uri="postgresql://u:p@localhost:5432/db",
)
@@ -76,14 +83,52 @@ def test_create_ssh_tunnel_command_invalid_params() -> None:
assert str(excinfo.value) == ("SSH Tunnel parameters are invalid.")
def test_create_ssh_tunnel_command_no_port() -> None:
def test_create_ssh_tunnel_command_no_port(session: Session) -> None:
"""
Test that SSH Tunnel can be created without explicit port but with a default one.
"""
from superset import db
from superset.commands.database.ssh_tunnel.create import CreateSSHTunnelCommand
from superset.databases.ssh_tunnel.models import SSHTunnel
from superset.models.core import Database
engine = db.session.get_bind()
Database.metadata.create_all(engine) # pylint: disable=no-member
database = Database(
database_name="my_database",
sqlalchemy_uri="postgresql://u:p@localhost/db",
)
properties = {
"database": database,
"server_address": "123.132.123.1",
"server_port": "3005",
"username": "foo",
"password": "bar",
}
result = CreateSSHTunnelCommand(database, properties).run()
assert result is not None
assert isinstance(result, SSHTunnel)
def test_create_ssh_tunnel_command_no_port_no_default(session: Session) -> None:
"""
Test that error is raised when creating SSH Tunnel without explicit/default ports.
"""
from superset import db
from superset.commands.database.ssh_tunnel.create import CreateSSHTunnelCommand
from superset.models.core import Database
engine = db.session.get_bind()
Database.metadata.create_all(engine) # pylint: disable=no-member
database = Database(
id=1,
database_name="my_database",
sqlalchemy_uri="postgresql://u:p@localhost/db",
sqlalchemy_uri="weird+db://u:p@localhost/db",
)
properties = {

31
tests/unit_tests/databases/ssh_tunnel/commands/update_test.py
View File

@@ -103,6 +103,37 @@ def test_update_shh_tunnel_invalid_params(session_with_data: Session) -> None:
"session_with_data", ["postgresql://u:p@localhost/testdb"], indirect=True
)
def test_update_shh_tunnel_no_port(session_with_data: Session) -> None:
"""
Test that SSH Tunnel can be updated without explicit port but with a default one.
"""
from superset.commands.database.ssh_tunnel.update import UpdateSSHTunnelCommand
from superset.daos.database import DatabaseDAO
from superset.databases.ssh_tunnel.models import SSHTunnel
result = DatabaseDAO.get_ssh_tunnel(1)
assert result
assert isinstance(result, SSHTunnel)
assert 1 == result.database_id
assert "Test" == result.server_address
update_payload = {"server_address": "Test2"}
UpdateSSHTunnelCommand(1, update_payload).run()
result = DatabaseDAO.get_ssh_tunnel(1)
assert result
assert isinstance(result, SSHTunnel)
assert "Test2" == result.server_address
@pytest.mark.parametrize(
"session_with_data", ["weird+db://u:p@localhost/testdb"], indirect=True
)
def test_update_shh_tunnel_no_port_no_default(session_with_data: Session) -> None:
"""
Test that error is raised when updating SSH Tunnel without explicit/default ports.
"""
from superset.commands.database.ssh_tunnel.update import UpdateSSHTunnelCommand
from superset.daos.database import DatabaseDAO
from superset.databases.ssh_tunnel.models import SSHTunnel