fix: permission checks on import (#23200)

2026-05-08 17:35:33 +00:00 · 2023-03-15 08:31:09 -07:00
parent 831978f0f7
commit cfc2ca672e
17 changed files with 621 additions and 99 deletions
--- a/tests/integration_tests/databases/commands_tests.py
+++ b/tests/integration_tests/databases/commands_tests.py
@@ -40,7 +40,7 @@ from superset.databases.commands.importers.v1 import ImportDatabasesCommand
 from superset.databases.commands.tables import TablesDatabaseCommand
 from superset.databases.commands.test_connection import TestConnectionDatabaseCommand
 from superset.databases.commands.validate import ValidateDatabaseParametersCommand
-from superset.databases.schemas import DatabaseTestConnectionSchema
+from superset.databases.ssh_tunnel.models import SSHTunnel
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.exceptions import (
    SupersetErrorsException,
@@ -63,16 +63,19 @@ from tests.integration_tests.fixtures.energy_dashboard import (
 from tests.integration_tests.fixtures.importexport import (
    database_config,
    database_metadata_config,
+    database_with_ssh_tunnel_config_mix_credentials,
+    database_with_ssh_tunnel_config_no_credentials,
+    database_with_ssh_tunnel_config_password,
+    database_with_ssh_tunnel_config_private_key,
+    database_with_ssh_tunnel_config_private_pass_only,
    dataset_config,
    dataset_metadata_config,
 )


 class TestCreateDatabaseCommand(SupersetTestCase):
-    @mock.patch(
-        "superset.databases.commands.test_connection.event_logger.log_with_context"
-    )
-    @mock.patch("superset.utils.core.g")
+    @patch("superset.databases.commands.test_connection.event_logger.log_with_context")
+    @patch("superset.utils.core.g")
    def test_create_duplicate_error(self, mock_g, mock_logger):
        example_db = get_example_database()
        mock_g.user = security_manager.find_user("admin")
@@ -90,10 +93,8 @@ class TestCreateDatabaseCommand(SupersetTestCase):
            "DatabaseRequiredFieldValidationError"
        )

-    @mock.patch(
-        "superset.databases.commands.test_connection.event_logger.log_with_context"
-    )
-    @mock.patch("superset.utils.core.g")
+    @patch("superset.databases.commands.test_connection.event_logger.log_with_context")
+    @patch("superset.utils.core.g")
    def test_multiple_error_logging(self, mock_g, mock_logger):
        mock_g.user = security_manager.find_user("admin")
        command = CreateDatabaseCommand({})
@@ -395,8 +396,11 @@ class TestExportDatabasesCommand(SupersetTestCase):


 class TestImportDatabasesCommand(SupersetTestCase):
-    def test_import_v1_database(self):
+    @patch("superset.security.manager.g")
+    def test_import_v1_database(self, mock_g):
        """Test that a database can be imported"""
+        mock_g.user = security_manager.find_user("admin")
+
        contents = {
            "metadata.yaml": yaml.safe_dump(database_metadata_config),
            "databases/imported_database.yaml": yaml.safe_dump(database_config),
@@ -421,7 +425,8 @@ class TestImportDatabasesCommand(SupersetTestCase):
        db.session.delete(database)
        db.session.commit()

-    def test_import_v1_database_broken_csv_fields(self):
+    @patch("superset.security.manager.g")
+    def test_import_v1_database_broken_csv_fields(self, mock_g):
        """
        Test that a database can be imported with broken schema.

@@ -429,6 +434,8 @@ class TestImportDatabasesCommand(SupersetTestCase):
        the V1 schema. This test ensures that we can import databases that were
        exported with the broken schema.
        """
+        mock_g.user = security_manager.find_user("admin")
+
        broken_config = database_config.copy()
        broken_config["allow_file_upload"] = broken_config.pop("allow_csv_upload")
        broken_config["extra"] = {"schemas_allowed_for_file_upload": ["upload"]}
@@ -457,8 +464,11 @@ class TestImportDatabasesCommand(SupersetTestCase):
        db.session.delete(database)
        db.session.commit()

-    def test_import_v1_database_multiple(self):
+    @patch("superset.security.manager.g")
+    def test_import_v1_database_multiple(self, mock_g):
        """Test that a database can be imported multiple times"""
+        mock_g.user = security_manager.find_user("admin")
+
        num_databases = db.session.query(Database).count()

        contents = {
@@ -498,8 +508,11 @@ class TestImportDatabasesCommand(SupersetTestCase):
        db.session.delete(database)
        db.session.commit()

-    def test_import_v1_database_with_dataset(self):
+    @patch("superset.security.manager.g")
+    def test_import_v1_database_with_dataset(self, mock_g):
        """Test that a database can be imported with datasets"""
+        mock_g.user = security_manager.find_user("admin")
+
        contents = {
            "databases/imported_database.yaml": yaml.safe_dump(database_config),
            "datasets/imported_dataset.yaml": yaml.safe_dump(dataset_config),
@@ -518,8 +531,11 @@ class TestImportDatabasesCommand(SupersetTestCase):
        db.session.delete(database)
        db.session.commit()

-    def test_import_v1_database_with_dataset_multiple(self):
+    @patch("superset.security.manager.g")
+    def test_import_v1_database_with_dataset_multiple(self, mock_g):
        """Test that a database can be imported multiple times w/o changing datasets"""
+        mock_g.user = security_manager.find_user("admin")
+
        contents = {
            "databases/imported_database.yaml": yaml.safe_dump(database_config),
            "datasets/imported_dataset.yaml": yaml.safe_dump(dataset_config),
@@ -623,6 +639,200 @@ class TestImportDatabasesCommand(SupersetTestCase):
            }
        }

+    @patch("superset.databases.schemas.is_feature_enabled")
+    def test_import_v1_database_masked_ssh_tunnel_password(
+        self, mock_schema_is_feature_enabled
+    ):
+        """Test that database imports with masked ssh_tunnel passwords are rejected"""
+        mock_schema_is_feature_enabled.return_value = True
+        masked_database_config = database_with_ssh_tunnel_config_password.copy()
+        contents = {
+            "metadata.yaml": yaml.safe_dump(database_metadata_config),
+            "databases/imported_database.yaml": yaml.safe_dump(masked_database_config),
+        }
+        command = ImportDatabasesCommand(contents)
+        with pytest.raises(CommandInvalidError) as excinfo:
+            command.run()
+        assert str(excinfo.value) == "Error importing database"
+        assert excinfo.value.normalized_messages() == {
+            "databases/imported_database.yaml": {
+                "_schema": ["Must provide a password for the ssh tunnel"]
+            }
+        }
+
+    @patch("superset.databases.schemas.is_feature_enabled")
+    def test_import_v1_database_masked_ssh_tunnel_private_key_and_password(
+        self, mock_schema_is_feature_enabled
+    ):
+        """Test that database imports with masked ssh_tunnel private_key and private_key_password are rejected"""
+        mock_schema_is_feature_enabled.return_value = True
+        masked_database_config = database_with_ssh_tunnel_config_private_key.copy()
+        contents = {
+            "metadata.yaml": yaml.safe_dump(database_metadata_config),
+            "databases/imported_database.yaml": yaml.safe_dump(masked_database_config),
+        }
+        command = ImportDatabasesCommand(contents)
+        with pytest.raises(CommandInvalidError) as excinfo:
+            command.run()
+        assert str(excinfo.value) == "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",
+                ]
+            }
+        }
+
+    @patch("superset.databases.schemas.is_feature_enabled")
+    @patch("superset.security.manager.g")
+    def test_import_v1_database_with_ssh_tunnel_password(
+        self,
+        mock_g,
+        mock_schema_is_feature_enabled,
+    ):
+        """Test that a database with ssh_tunnel password can be imported"""
+        mock_g.user = security_manager.find_user("admin")
+        mock_schema_is_feature_enabled.return_value = True
+        masked_database_config = database_with_ssh_tunnel_config_password.copy()
+        masked_database_config["ssh_tunnel"]["password"] = "TEST"
+        contents = {
+            "metadata.yaml": yaml.safe_dump(database_metadata_config),
+            "databases/imported_database.yaml": yaml.safe_dump(masked_database_config),
+        }
+        command = ImportDatabasesCommand(contents)
+        command.run()
+
+        database = (
+            db.session.query(Database).filter_by(uuid=database_config["uuid"]).one()
+        )
+        assert database.allow_file_upload
+        assert database.allow_ctas
+        assert database.allow_cvas
+        assert database.allow_dml
+        assert not database.allow_run_async
+        assert database.cache_timeout is None
+        assert database.database_name == "imported_database"
+        assert database.expose_in_sqllab
+        assert database.extra == "{}"
+        assert database.sqlalchemy_uri == "sqlite:///test.db"
+
+        model_ssh_tunnel = (
+            db.session.query(SSHTunnel)
+            .filter(SSHTunnel.database_id == database.id)
+            .one()
+        )
+        self.assertEqual(model_ssh_tunnel.password, "TEST")
+
+        db.session.delete(database)
+        db.session.commit()
+
+    @patch("superset.databases.schemas.is_feature_enabled")
+    @patch("superset.security.manager.g")
+    def test_import_v1_database_with_ssh_tunnel_private_key_and_password(
+        self,
+        mock_g,
+        mock_schema_is_feature_enabled,
+    ):
+        """Test that a database with ssh_tunnel private_key and private_key_password can be imported"""
+        mock_g.user = security_manager.find_user("admin")
+
+        mock_schema_is_feature_enabled.return_value = True
+        masked_database_config = database_with_ssh_tunnel_config_private_key.copy()
+        masked_database_config["ssh_tunnel"]["private_key"] = "TestPrivateKey"
+        masked_database_config["ssh_tunnel"]["private_key_password"] = "TEST"
+        contents = {
+            "metadata.yaml": yaml.safe_dump(database_metadata_config),
+            "databases/imported_database.yaml": yaml.safe_dump(masked_database_config),
+        }
+        command = ImportDatabasesCommand(contents)
+        command.run()
+
+        database = (
+            db.session.query(Database).filter_by(uuid=database_config["uuid"]).one()
+        )
+        assert database.allow_file_upload
+        assert database.allow_ctas
+        assert database.allow_cvas
+        assert database.allow_dml
+        assert not database.allow_run_async
+        assert database.cache_timeout is None
+        assert database.database_name == "imported_database"
+        assert database.expose_in_sqllab
+        assert database.extra == "{}"
+        assert database.sqlalchemy_uri == "sqlite:///test.db"
+
+        model_ssh_tunnel = (
+            db.session.query(SSHTunnel)
+            .filter(SSHTunnel.database_id == database.id)
+            .one()
+        )
+        self.assertEqual(model_ssh_tunnel.private_key, "TestPrivateKey")
+        self.assertEqual(model_ssh_tunnel.private_key_password, "TEST")
+
+        db.session.delete(database)
+        db.session.commit()
+
+    @patch("superset.databases.schemas.is_feature_enabled")
+    def test_import_v1_database_masked_ssh_tunnel_no_credentials(
+        self, mock_schema_is_feature_enabled
+    ):
+        """Test that databases with ssh_tunnels that have no credentials are rejected"""
+        mock_schema_is_feature_enabled.return_value = True
+        masked_database_config = database_with_ssh_tunnel_config_no_credentials.copy()
+        contents = {
+            "metadata.yaml": yaml.safe_dump(database_metadata_config),
+            "databases/imported_database.yaml": yaml.safe_dump(masked_database_config),
+        }
+        command = ImportDatabasesCommand(contents)
+        with pytest.raises(CommandInvalidError) as excinfo:
+            command.run()
+        assert str(excinfo.value) == "Must provide credentials for the SSH Tunnel"
+
+    @patch("superset.databases.schemas.is_feature_enabled")
+    def test_import_v1_database_masked_ssh_tunnel_multiple_credentials(
+        self, mock_schema_is_feature_enabled
+    ):
+        """Test that databases with ssh_tunnels that have multiple credentials are rejected"""
+        mock_schema_is_feature_enabled.return_value = True
+        masked_database_config = database_with_ssh_tunnel_config_mix_credentials.copy()
+        contents = {
+            "metadata.yaml": yaml.safe_dump(database_metadata_config),
+            "databases/imported_database.yaml": yaml.safe_dump(masked_database_config),
+        }
+        command = ImportDatabasesCommand(contents)
+        with pytest.raises(CommandInvalidError) as excinfo:
+            command.run()
+        assert (
+            str(excinfo.value) == "Cannot have multiple credentials for the SSH Tunnel"
+        )
+
+    @patch("superset.databases.schemas.is_feature_enabled")
+    def test_import_v1_database_masked_ssh_tunnel_only_priv_key_psswd(
+        self, mock_schema_is_feature_enabled
+    ):
+        """Test that databases with ssh_tunnels that have multiple credentials are rejected"""
+        mock_schema_is_feature_enabled.return_value = True
+        masked_database_config = (
+            database_with_ssh_tunnel_config_private_pass_only.copy()
+        )
+        contents = {
+            "metadata.yaml": yaml.safe_dump(database_metadata_config),
+            "databases/imported_database.yaml": yaml.safe_dump(masked_database_config),
+        }
+        command = ImportDatabasesCommand(contents)
+        with pytest.raises(CommandInvalidError) as excinfo:
+            command.run()
+        assert str(excinfo.value) == "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",
+                ]
+            }
+        }
+
    @patch("superset.databases.commands.importers.v1.import_dataset")
    def test_import_v1_rollback(self, mock_import_dataset):
        """Test than on an exception everything is rolled back"""
@@ -648,11 +858,9 @@ class TestImportDatabasesCommand(SupersetTestCase):


 class TestTestConnectionDatabaseCommand(SupersetTestCase):
-    @mock.patch("superset.databases.dao.Database._get_sqla_engine")
-    @mock.patch(
-        "superset.databases.commands.test_connection.event_logger.log_with_context"
-    )
-    @mock.patch("superset.utils.core.g")
+    @patch("superset.databases.dao.Database._get_sqla_engine")
+    @patch("superset.databases.commands.test_connection.event_logger.log_with_context")
+    @patch("superset.utils.core.g")
    def test_connection_db_exception(
        self, mock_g, mock_event_logger, mock_get_sqla_engine
    ):
@@ -671,11 +879,9 @@ class TestTestConnectionDatabaseCommand(SupersetTestCase):
            )
        mock_event_logger.assert_called()

-    @mock.patch("superset.databases.dao.Database._get_sqla_engine")
-    @mock.patch(
-        "superset.databases.commands.test_connection.event_logger.log_with_context"
-    )
-    @mock.patch("superset.utils.core.g")
+    @patch("superset.databases.dao.Database._get_sqla_engine")
+    @patch("superset.databases.commands.test_connection.event_logger.log_with_context")
+    @patch("superset.utils.core.g")
    def test_connection_do_ping_exception(
        self, mock_g, mock_event_logger, mock_get_sqla_engine
    ):
@@ -696,11 +902,9 @@ class TestTestConnectionDatabaseCommand(SupersetTestCase):
            == 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"
-    )
-    @mock.patch("superset.utils.core.g")
+    @patch("superset.databases.commands.test_connection.func_timeout")
+    @patch("superset.databases.commands.test_connection.event_logger.log_with_context")
+    @patch("superset.utils.core.g")
    def test_connection_do_ping_timeout(
        self, mock_g, mock_event_logger, mock_func_timeout
    ):
@@ -720,11 +924,9 @@ class TestTestConnectionDatabaseCommand(SupersetTestCase):
            == 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"
-    )
-    @mock.patch("superset.utils.core.g")
+    @patch("superset.databases.dao.Database._get_sqla_engine")
+    @patch("superset.databases.commands.test_connection.event_logger.log_with_context")
+    @patch("superset.utils.core.g")
    def test_connection_superset_security_connection(
        self, mock_g, mock_event_logger, mock_get_sqla_engine
    ):
@@ -745,11 +947,9 @@ class TestTestConnectionDatabaseCommand(SupersetTestCase):

        mock_event_logger.assert_called()

-    @mock.patch("superset.databases.dao.Database._get_sqla_engine")
-    @mock.patch(
-        "superset.databases.commands.test_connection.event_logger.log_with_context"
-    )
-    @mock.patch("superset.utils.core.g")
+    @patch("superset.databases.dao.Database._get_sqla_engine")
+    @patch("superset.databases.commands.test_connection.event_logger.log_with_context")
+    @patch("superset.utils.core.g")
    def test_connection_db_api_exc(
        self, mock_g, mock_event_logger, mock_get_sqla_engine
    ):
@@ -772,9 +972,9 @@ class TestTestConnectionDatabaseCommand(SupersetTestCase):
        mock_event_logger.assert_called()


-@mock.patch("superset.db_engine_specs.base.is_hostname_valid")
-@mock.patch("superset.db_engine_specs.base.is_port_open")
-@mock.patch("superset.databases.commands.validate.DatabaseDAO")
+@patch("superset.db_engine_specs.base.is_hostname_valid")
+@patch("superset.db_engine_specs.base.is_port_open")
+@patch("superset.databases.commands.validate.DatabaseDAO")
 def test_validate(DatabaseDAO, is_port_open, is_hostname_valid, app_context):
    """
    Test parameter validation.
@@ -797,8 +997,8 @@ def test_validate(DatabaseDAO, is_port_open, is_hostname_valid, app_context):
    command.run()


-@mock.patch("superset.db_engine_specs.base.is_hostname_valid")
-@mock.patch("superset.db_engine_specs.base.is_port_open")
+@patch("superset.db_engine_specs.base.is_hostname_valid")
+@patch("superset.db_engine_specs.base.is_port_open")
 def test_validate_partial(is_port_open, is_hostname_valid, app_context):
    """
    Test parameter validation when only some parameters are present.
@@ -838,7 +1038,7 @@ def test_validate_partial(is_port_open, is_hostname_valid, app_context):
    ]


-@mock.patch("superset.db_engine_specs.base.is_hostname_valid")
+@patch("superset.db_engine_specs.base.is_hostname_valid")
 def test_validate_partial_invalid_hostname(is_hostname_valid, app_context):
    """
    Test parameter validation when only some parameters are present.
@@ -892,7 +1092,7 @@ def test_validate_partial_invalid_hostname(is_hostname_valid, app_context):


 class TestTablesDatabaseCommand(SupersetTestCase):
-    @mock.patch("superset.databases.dao.DatabaseDAO.find_by_id")
+    @patch("superset.databases.dao.DatabaseDAO.find_by_id")
    def test_database_tables_list_with_unknown_database(self, mock_find_by_id):
        mock_find_by_id.return_value = None
        command = TablesDatabaseCommand(1, "test", False)
@@ -901,9 +1101,9 @@ class TestTablesDatabaseCommand(SupersetTestCase):
            command.run()
            assert str(excinfo.value) == ("Database not found.")

-    @mock.patch("superset.databases.dao.DatabaseDAO.find_by_id")
-    @mock.patch("superset.security.manager.SupersetSecurityManager.can_access_database")
-    @mock.patch("superset.utils.core.g")
+    @patch("superset.databases.dao.DatabaseDAO.find_by_id")
+    @patch("superset.security.manager.SupersetSecurityManager.can_access_database")
+    @patch("superset.utils.core.g")
    def test_database_tables_superset_exception(
        self, mock_g, mock_can_access_database, mock_find_by_id
    ):
@@ -920,9 +1120,9 @@ class TestTablesDatabaseCommand(SupersetTestCase):
            command.run()
            assert str(excinfo.value) == "Test Error"

-    @mock.patch("superset.databases.dao.DatabaseDAO.find_by_id")
-    @mock.patch("superset.security.manager.SupersetSecurityManager.can_access_database")
-    @mock.patch("superset.utils.core.g")
+    @patch("superset.databases.dao.DatabaseDAO.find_by_id")
+    @patch("superset.security.manager.SupersetSecurityManager.can_access_database")
+    @patch("superset.utils.core.g")
    def test_database_tables_exception(
        self, mock_g, mock_can_access_database, mock_find_by_id
    ):
@@ -939,9 +1139,9 @@ class TestTablesDatabaseCommand(SupersetTestCase):
                == "Unexpected error occurred, please check your logs for details"
            )

-    @mock.patch("superset.databases.dao.DatabaseDAO.find_by_id")
-    @mock.patch("superset.security.manager.SupersetSecurityManager.can_access_database")
-    @mock.patch("superset.utils.core.g")
+    @patch("superset.databases.dao.DatabaseDAO.find_by_id")
+    @patch("superset.security.manager.SupersetSecurityManager.can_access_database")
+    @patch("superset.utils.core.g")
    def test_database_tables_list_tables(
        self, mock_g, mock_can_access_database, mock_find_by_id
    ):