QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-04-20 00:24:38 +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)

52
tests/unit_tests/importexport/api_test.py
View File

@@ -113,6 +113,58 @@ def test_import_assets(
ssh_tunnel_passwords=None,
ssh_tunnel_private_keys=None,
ssh_tunnel_priv_key_passwords=None,
encrypted_extra_secrets=None,
)
def test_import_assets_with_encrypted_extra_secrets(
mocker: MockerFixture,
client: Any,
full_api_access: None,
) -> None:
"""
Test that encrypted_extra_secrets are passed to ImportAssetsCommand.
"""
mocked_contents = {
"metadata.yaml": (
"version: 1.0.0\ntype: assets\ntimestamp: '2022-01-01T00:00:00+00:00'\n"
),
"databases/example.yaml": "<DATABASE CONTENTS>",
}
ImportAssetsCommand = mocker.patch("superset.importexport.api.ImportAssetsCommand") # noqa: N806
root = Path("assets_export")
buf = BytesIO()
with ZipFile(buf, "w") as bundle:
for path, contents in mocked_contents.items():
with bundle.open(str(root / path), "w") as fp:
fp.write(contents.encode())
buf.seek(0)
secrets = {
"assets_export/databases/example.yaml": {
"$.credentials_info.private_key": "-----BEGIN PRIVATE KEY-----\nKEY\n-----END PRIVATE KEY-----\n", # noqa: E501
}
}
form_data = {
"bundle": (buf, "assets_export.zip"),
"encrypted_extra_secrets": json.dumps(secrets),
}
response = client.post(
"/api/v1/assets/import/", data=form_data, content_type="multipart/form-data"
)
assert response.status_code == 200
assert response.json == {"message": "OK"}
ImportAssetsCommand.assert_called_with(
mocked_contents,
sparse=False,
passwords=None,
ssh_tunnel_passwords=None,
ssh_tunnel_private_keys=None,
ssh_tunnel_priv_key_passwords=None,
encrypted_extra_secrets=secrets,
)

139
tests/unit_tests/utils/json_tests.py
View File

@@ -19,6 +19,7 @@ import math
import uuid
from datetime import date, datetime, time, timedelta
from decimal import Decimal
from typing import Any
from unittest.mock import MagicMock
import numpy as np
@@ -26,6 +27,7 @@ import pandas as pd
import pytest
import pytz
from superset.constants import PASSWORD_MASK
from superset.utils import json
from superset.utils.core import (
zlib_compress,
@@ -264,6 +266,143 @@ def test_json_int_dttm_ser():
json.json_int_dttm_ser(np.datetime64())
@pytest.mark.parametrize(
"payload,path_values,expected_result",
[
(
{
"credentials_info": {
"type": "service_account",
"private_key": "XXXXXXXXXX",
},
},
{"$.credentials_info.private_key": "NEW_KEY"},
{
"credentials_info": {
"type": "service_account",
"private_key": "NEW_KEY",
},
},
),
(
{
"auth_params": {
"privatekey_body": "XXXXXXXXXX",
"privatekey_pass": "XXXXXXXXXX",
},
"other": "value",
},
{
"$.auth_params.privatekey_body": "-----BEGIN PRIVATE KEY-----",
"$.auth_params.privatekey_pass": "passphrase",
},
{
"auth_params": {
"privatekey_body": "-----BEGIN PRIVATE KEY-----",
"privatekey_pass": "passphrase",
},
"other": "value",
},
),
(
{"existing": "value"},
{"$.nonexistent.path": "new_value"},
{"existing": "value"},
),
],
)
def test_set_masked_fields(
payload: dict[str, Any],
path_values: dict[str, Any],
expected_result: dict[str, Any],
) -> None:
"""
Test setting a value at a JSONPath location.
"""
result = json.set_masked_fields(payload, path_values)
assert result == expected_result
@pytest.mark.parametrize(
"payload,sensitive_fields,expected_result",
[
(
{
"credentials_info": {
"type": "service_account",
"private_key": PASSWORD_MASK,
},
},
{"$.credentials_info.private_key", "$.credentials_info.type"},
["$.credentials_info.private_key"],
),
(
{
"credentials_info": {
"private_key": "ACTUAL_KEY",
},
},
{"$.credentials_info.private_key"},
[],
),
(
{
"auth_params": {
"privatekey_body": PASSWORD_MASK,
"privatekey_pass": "actual_pass",
},
"oauth2_client_info": {
"secret": PASSWORD_MASK,
},
},
{
"$.auth_params.privatekey_body",
"$.auth_params.privatekey_pass",
"$.oauth2_client_info.secret",
},
[
"$.auth_params.privatekey_body",
"$.oauth2_client_info.secret",
],
),
(
{
"foo": PASSWORD_MASK,
"service_account_info": PASSWORD_MASK,
},
{"$.*"},
["$.foo", "$.service_account_info"],
),
(
{
"foo": PASSWORD_MASK,
"bar": "actual_value",
},
{"$.*"},
["$.foo"],
),
(
{
"foo": "actual_value",
"bar": "other_value",
},
{"$.*"},
[],
),
],
)
def test_get_masked_fields(
payload: dict[str, Any],
sensitive_fields: set[str],
expected_result: dict[str, Any],
) -> None:
"""
Test that get_masked_fields returns paths where value equals PASSWORD_MASK.
"""
masked = json.get_masked_fields(payload, sensitive_fields)
assert sorted(masked) == sorted(expected_result)
def test_format_timedelta():
assert json.format_timedelta(timedelta(0)) == "0:00:00"
assert json.format_timedelta(timedelta(days=1)) == "1 day, 0:00:00"