feat: default ports for SSH tunnel (#32403)

tests/integration_tests/databases/api_tests.py

@@ -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",

tests/unit_tests/databases/ssh_tunnel/commands/create_test.py

@@ -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 = {

tests/unit_tests/databases/ssh_tunnel/commands/update_test.py

@@ -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