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

feat: support for import/export masked_encrypted_extra (backend) (#38077)

Browse Source
This commit is contained in:
Vitor Avila
2026-03-04 16:26:28 -03:00
committed by GitHub
parent 63e7ee70bf
commit 8c9efe5659
16 changed files with 799 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

70
tests/unit_tests/databases/api_test.py
View File

@@ -451,6 +451,76 @@ def test_import(
ssh_tunnel_passwords=None,
ssh_tunnel_private_keys=None,
ssh_tunnel_priv_key_passwords=None,
encrypted_extra_secrets=None,
)
def test_import_with_encrypted_extra_secrets(
mocker: MockerFixture,
client: Any,
full_api_access: None,
) -> None:
"""
Test that encrypted_extra_secrets are passed to ImportDatabasesCommand.
"""
contents = {
"metadata.yaml": yaml.safe_dump(
{
"version": "1.0.0",
"type": "Database",
"timestamp": "2021-01-01T00:00:00Z",
}
),
"databases/test.yaml": yaml.safe_dump(
{
"database_name": "test",
"sqlalchemy_uri": "bigquery://gcp-project-id/",
"cache_timeout": 0,
"expose_in_sqllab": True,
"allow_run_async": False,
"allow_ctas": False,
"allow_cvas": False,
"allow_dml": False,
"allow_file_upload": False,
"masked_encrypted_extra": json.dumps(
{"credentials_info": {"private_key": "XXXXXXXXXX"}}
),
"extra": json.dumps({"allows_virtual_table_explore": True}),
"uuid": "00000000-0000-0000-0000-123456789001",
}
),
}
mocker.patch("superset.databases.api.is_zipfile", return_value=True)
mocker.patch("superset.databases.api.ZipFile")
mocker.patch(
"superset.databases.api.get_contents_from_bundle",
return_value=contents,
)
command = mocker.patch("superset.databases.api.ImportDatabasesCommand")
secrets = {
"databases/test.yaml": {
"$.credentials_info.private_key": "-----BEGIN PRIVATE KEY-----"
}
}
form_data = {
"formData": (BytesIO(b"test"), "test.zip"),
"encrypted_extra_secrets": json.dumps(secrets),
}
client.post(
"/api/v1/database/import/",
data=form_data,
content_type="multipart/form-data",
)
command.assert_called_with(
contents,
passwords=None,
overwrite=False,
ssh_tunnel_passwords=None,
ssh_tunnel_private_keys=None,
ssh_tunnel_priv_key_passwords=None,
encrypted_extra_secrets=secrets,
)

123
tests/unit_tests/databases/commands/importers/v1/import_test.py
View File

@@ -250,3 +250,126 @@ def test_import_database_with_user_impersonation(
database = import_database(config)
assert database.impersonate_user is True
def test_import_database_with_masked_encrypted_extra_new_db(
mocker: MockerFixture,
session: Session,
) -> None:
"""
Test importing a new database with masked_encrypted_extra.
When no existing DB matches the UUID, the masked_encrypted_extra value
should be stored as-is in encrypted_extra.
"""
from superset import security_manager
from superset.commands.database.importers.v1.utils import import_database
from superset.models.core import Database
from tests.integration_tests.fixtures.importexport import (
database_config_with_masked_encrypted_extra,
)
mocker.patch.object(security_manager, "can_access", return_value=True)
mocker.patch("superset.commands.database.importers.v1.utils.add_permissions")
engine = db.session.get_bind()
Database.metadata.create_all(engine) # pylint: disable=no-member
config = copy.deepcopy(database_config_with_masked_encrypted_extra)
database = import_database(config)
assert database.database_name == "imported_database_encrypted"
# masked_encrypted_extra should be stored as encrypted_extra
assert database.encrypted_extra is not None
encrypted = json.loads(database.encrypted_extra)
assert encrypted["credentials_info"]["type"] == "service_account"
assert encrypted["credentials_info"]["project_id"] == "test-project"
assert encrypted["credentials_info"]["private_key"] == (
"-----BEGIN PRIVATE KEY-----\nMyPriVaTeKeY\n-----END PRIVATE KEY-----\n"
)
def test_import_database_with_masked_encrypted_extra_existing_db(
mocker: MockerFixture,
session: Session,
) -> None:
"""
Test importing over an existing database with masked_encrypted_extra.
When the import carries PASSWORD_MASK values for sensitive fields and
an existing DB has the real values, reveal_sensitive should restore
the original values from the existing DB's encrypted_extra.
"""
from superset import security_manager
from superset.commands.database.importers.v1.utils import import_database
from superset.constants import PASSWORD_MASK
from superset.models.core import Database
from tests.integration_tests.fixtures.importexport import (
database_config_with_masked_encrypted_extra,
)
mocker.patch.object(security_manager, "can_access", return_value=True)
mocker.patch("superset.commands.database.importers.v1.utils.add_permissions")
engine = db.session.get_bind()
Database.metadata.create_all(engine) # pylint: disable=no-member
# First, create the existing database with real encrypted_extra
config = copy.deepcopy(database_config_with_masked_encrypted_extra)
import_database(config)
db.session.flush()
# Now import again with masked values (simulating re-import)
config2 = copy.deepcopy(database_config_with_masked_encrypted_extra)
config2["masked_encrypted_extra"] = json.dumps(
{
"credentials_info": {
"type": "service_account",
"project_id": "updated-project",
"private_key": PASSWORD_MASK,
}
}
)
database2 = import_database(config2, overwrite=True)
# The masked private_key should be revealed from the existing DB
encrypted = json.loads(database2.encrypted_extra)
assert encrypted["credentials_info"]["project_id"] == "updated-project"
assert encrypted["credentials_info"]["private_key"] == (
"-----BEGIN PRIVATE KEY-----\nMyPriVaTeKeY\n-----END PRIVATE KEY-----\n"
)
assert encrypted["credentials_info"]["private_key"] != PASSWORD_MASK
def test_import_database_oauth2_redirect_is_nonfatal(
mocker: MockerFixture,
session: Session,
) -> None:
"""
Test that an OAuth2RedirectError from add_permissions is logged
and does not prevent the import from succeeding.
"""
from superset import security_manager
from superset.commands.database.importers.v1.utils import import_database
from superset.exceptions import OAuth2RedirectError
from superset.models.core import Database
from tests.integration_tests.fixtures.importexport import database_config
mocker.patch.object(security_manager, "can_access", return_value=True)
mock_add_perms = mocker.patch(
"superset.commands.database.importers.v1.utils.add_permissions",
side_effect=OAuth2RedirectError(
url="https://oauth.example.com/authorize",
tab_id="abc-123",
redirect_uri="https://superset.example.com/callback",
),
)
engine = db.session.get_bind()
Database.metadata.create_all(engine) # pylint: disable=no-member
config = copy.deepcopy(database_config)
database = import_database(config)
assert database.database_name == "imported_database"
mock_add_perms.assert_called_once_with(database)

106
tests/unit_tests/databases/schema_tests.py
View File

@@ -23,6 +23,8 @@ import pytest
from marshmallow import fields, Schema, ValidationError
from pytest_mock import MockerFixture
from superset.utils import json
if TYPE_CHECKING:
from superset.databases.schemas import DatabaseParametersSchemaMixin
@@ -61,6 +63,24 @@ def dummy_engine(mocker: MockerFixture) -> None:
mocker.patch("superset.databases.schemas.get_engine_spec", return_value=DummyEngine)
@pytest.fixture
def mock_bq_engine(mocker: MockerFixture) -> None:
"""
Fixture providing a mocked BQ engine spec.
"""
from superset.db_engine_specs.bigquery import BigQueryEngineSpec
mock_url = mocker.MagicMock()
mock_url.get_backend_name.return_value = "bigquery"
mock_url.get_driver_name.return_value = "bigquery"
mocker.patch("superset.databases.schemas.make_url_safe", return_value=mock_url)
mocker.patch(
"superset.databases.schemas.get_engine_spec",
return_value=BigQueryEngineSpec,
)
def test_database_parameters_schema_mixin(
dummy_engine: None,
dummy_schema: "Schema",
@@ -272,3 +292,89 @@ def test_oauth2_schema_extra() -> None:
}
)
assert payload == {"code": "SECRET", "state": "12345"}
def test_import_schema_rejects_both_encrypted_and_masked() -> None:
"""
Test that ImportV1DatabaseSchema rejects configs with both
encrypted_extra and masked_encrypted_extra.
"""
from superset.databases.schemas import ImportV1DatabaseSchema
schema = ImportV1DatabaseSchema()
config = {
"database_name": "test_db",
"sqlalchemy_uri": "bigquery://test/",
"uuid": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee",
"encrypted_extra": json.dumps({"secret": "value"}),
"masked_encrypted_extra": json.dumps({"secret": "XXXXXXXXXX"}),
"extra": {},
"version": "1.0.0",
}
with pytest.raises(ValidationError) as exc_info:
schema.load(config)
assert "File contains both" in str(exc_info.value)
def test_import_schema_rejects_masked_fields_for_new_db(
mock_bq_engine: None,
mocker: MockerFixture,
) -> None:
"""
Test that ImportV1DatabaseSchema rejects configs with PASSWORD_MASK
values for a new DB (no existing UUID match).
"""
from superset.databases.schemas import ImportV1DatabaseSchema
mock_session = mocker.patch("superset.databases.schemas.db.session")
mock_session.query.return_value.filter_by.return_value.first.return_value = None
schema = ImportV1DatabaseSchema()
config = {
"database_name": "test_db",
"sqlalchemy_uri": "bigquery://test/",
"uuid": "bbbbbbbb-aaaa-cccc-dddd-eeeeeeeeeeff",
"masked_encrypted_extra": json.dumps(
{"credentials_info": {"private_key": "XXXXXXXXXX"}}
),
"extra": {},
"version": "1.0.0",
}
with pytest.raises(ValidationError) as exc_info:
schema.load(config)
error_messages = str(exc_info.value)
assert "Must provide value for masked_encrypted_extra field" in error_messages
assert "$.credentials_info.private_key" in error_messages
def test_import_schema_allows_masked_fields_for_existing_db(
mock_bq_engine: None,
mocker: MockerFixture,
) -> None:
"""
Test that ImportV1DatabaseSchema allows PASSWORD_MASK values when
the DB already exists (UUID match). The reveal will happen later
in import_database().
"""
from superset.databases.schemas import ImportV1DatabaseSchema
mock_session = mocker.patch("superset.databases.schemas.db.session")
mock_existing_db = mocker.MagicMock()
mock_session = mocker.patch("superset.databases.schemas.db.session")
mock_session.query.return_value.filter_by.return_value.first.return_value = (
mock_existing_db
)
schema = ImportV1DatabaseSchema()
config = {
"database_name": "test_db",
"sqlalchemy_uri": "bigquery://test/",
"uuid": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee",
"masked_encrypted_extra": json.dumps(
{"credentials_info": {"private_key": "XXXXXXXXXX"}}
),
"extra": {},
"version": "1.0.0",
}
# Should not raise - masked values are allowed for existing DBs
schema.load(config)