feat(api): add ability to import also secure/encrypted extra for databases

superset/commands/importers/v1/__init__.py

@@ -48,6 +48,7 @@
     def __init__(self, contents: dict[str, str], *args: Any, **kwargs: Any):
         self.contents = contents
         self.passwords: dict[str, str] = kwargs.get("passwords") or {}
+        self.encrypted_extras: dict[str, Schema] = kwargs.get("encrypted_extras") or {}
         self.ssh_tunnel_passwords: dict[str, str] = (
             kwargs.get("ssh_tunnel_passwords") or {}
         )
@@ -96,6 +97,7 @@
             self.contents,
             self.schemas,
             self.passwords,
+            self.encrypted_extras,
             exceptions,
             self.ssh_tunnel_passwords,
             self.ssh_tunnel_private_keys,

superset/commands/importers/v1/assets.py

@@ -69,6 +69,7 @@
     def __init__(self, contents: dict[str, str], *args: Any, **kwargs: Any):
         self.contents = contents
         self.passwords: dict[str, str] = kwargs.get("passwords") or {}
+        self.encrypted_extras: dict[str, Schema] = kwargs.get("encrypted_extras") or {}
         self.ssh_tunnel_passwords: dict[str, str] = (
             kwargs.get("ssh_tunnel_passwords") or {}
         )
@@ -178,6 +179,7 @@
             self.contents,
             self.schemas,
             self.passwords,
+            self.encrypted_extras,
             exceptions,
             self.ssh_tunnel_passwords,
             self.ssh_tunnel_private_keys,

superset/commands/importers/v1/utils.py

@@ -100,6 +100,7 @@
     contents: dict[str, str],
     schemas: dict[str, Schema],
     passwords: dict[str, str],
+    encrypted_extra: dict[str, Schema],
     exceptions: list[ValidationError],
     ssh_tunnel_passwords: dict[str, str],
     ssh_tunnel_private_keys: dict[str, str],
@@ -112,6 +113,13 @@
         str(uuid): password
         for uuid, password in db.session.query(Database.uuid, Database.password).all()
     }
+    # load existing encrypted extras so we can apply the password validation
+    db_encrypted_extras: dict[str, str] = {
+        str(uuid): encrypted_extra
+        for uuid, encrypted_extra in db.session.query(
+            Database.uuid, Database.encrypted_extra
+        ).all()
+    }
     # load existing ssh_tunnels so we can apply the password validation
     db_ssh_tunnel_passwords: dict[str, str] = {
         str(uuid): password
@@ -148,6 +156,12 @@
                 elif prefix == "databases" and config["uuid"] in db_passwords:
                     config["password"] = db_passwords[config["uuid"]]
 
+                # populate encrypted_extras from the request or from existing DBs
+                if file_name in encrypted_extra:
+                    config["encrypted_extra"] = encrypted_extra[file_name]
+                elif prefix == "databases" and config["uuid"] in db_encrypted_extras:
+                    config["encrypted_extra"] = db_encrypted_extras[config["uuid"]]
+
                 # populate ssh_tunnel_passwords from the request or from existing DBs
                 if file_name in ssh_tunnel_passwords:
                     config["ssh_tunnel"]["password"] = ssh_tunnel_passwords[file_name]

superset/databases/api.py

@@ -1526,6 +1526,14 @@
                     overwrite:
                       description: overwrite existing databases?
                       type: boolean
+                    encrypted_extras:
+                      description: >-
+                        JSON map of encrypted_extras for each featured database
+                        in the ZIP file. If the ZIP includes a database config in the
+                        path `databases/MyDatabase.yaml`, the encrypted_extra should be
+                        provided in the following format:
+                        `{"databases/MyDatabase.yaml": {"key": "value"}}`.
+                      type: string
                     ssh_tunnel_passwords:
                       description: >-
                         JSON map of passwords for each ssh_tunnel associated to a
@@ -1586,6 +1594,11 @@
             else None
         )
         overwrite = request.form.get("overwrite") == "true"
+        encrypted_extras = (
+            json.loads(request.form["encrypted_extras"])
+            if "encrypted_extras" in request.form
+            else None
+        )
         ssh_tunnel_passwords = (
             json.loads(request.form["ssh_tunnel_passwords"])
             if "ssh_tunnel_passwords" in request.form
@@ -1606,6 +1619,7 @@
             contents,
             passwords=passwords,
             overwrite=overwrite,
+            encrypted_extras=encrypted_extras,
             ssh_tunnel_passwords=ssh_tunnel_passwords,
             ssh_tunnel_private_keys=ssh_tunnel_private_keys,
             ssh_tunnel_priv_key_passwords=ssh_tunnel_priv_key_passwords,

superset/databases/schemas.py

@@ -864,6 +864,7 @@ def fix_allow_csv_upload(
     allow_csv_upload = fields.Boolean()
     impersonate_user = fields.Boolean()
     extra = fields.Nested(ImportV1DatabaseExtraSchema)
+    encrypted_extra = fields.String(allow_none=True)
     uuid = fields.UUID(required=True)
     version = fields.String(required=True)
     is_managed_externally = fields.Boolean(allow_none=True, dump_default=False)

superset/importexport/api.py

@@ -123,6 +123,14 @@ def import_(self) -> Response:
                         in the following format:
                         `{"databases/MyDatabase.yaml": "my_password"}`.
                       type: string
+                    encrypted_extras:
+                      description: >-
+                        JSON map of encrypted_extras for each featured database
+                        in the ZIP file. If the ZIP includes a database config in the
+                        path `databases/MyDatabase.yaml`, the encrypted_extra should be
+                        provided in the following format:
+                        `{"databases/MyDatabase.yaml": {"key": "value"}}`.
+                      type: string
                     ssh_tunnel_passwords:
                       description: >-
                         JSON map of passwords for each ssh_tunnel associated to a
@@ -183,6 +191,11 @@ def import_(self) -> Response:
             if "passwords" in request.form
             else None
         )
+        encrypted_extras = (
+            json.loads(request.form["encrypted_extras"])
+            if "encrypted_extras" in request.form
+            else None
+        )
         ssh_tunnel_passwords = (
             json.loads(request.form["ssh_tunnel_passwords"])
             if "ssh_tunnel_passwords" in request.form
@@ -202,6 +215,7 @@ def import_(self) -> Response:
         command = ImportAssetsCommand(
             contents,
             passwords=passwords,
+            encrypted_extras=encrypted_extras,
             ssh_tunnel_passwords=ssh_tunnel_passwords,
             ssh_tunnel_private_keys=ssh_tunnel_private_keys,
             ssh_tunnel_priv_key_passwords=ssh_tunnel_priv_key_passwords,

tests/integration_tests/databases/api_tests.py

@@ -2532,6 +2532,54 @@ def test_import_database_masked_password_provided(self, mock_add_permissions):
         db.session.delete(database)
         db.session.commit()
 
+    @mock.patch("superset.commands.database.importers.v1.utils.add_permissions")
+    def test_import_database_encrypted_extra_provided(self, mock_add_permissions):
+        """
+        Database API: Test import database with masked password provided
+        """
+        self.login(ADMIN_USERNAME)
+        uri = "api/v1/database/import/"
+
+        masked_database_config = database_config.copy()
+        masked_database_config["sqlalchemy_uri"] = (
+            "vertica+vertica_python://host:5433/dbname?ssl=1"
+        )
+
+        buf = BytesIO()
+        with ZipFile(buf, "w") as bundle:
+            with bundle.open("database_export/metadata.yaml", "w") as fp:
+                fp.write(yaml.safe_dump(database_metadata_config).encode())
+            with bundle.open(
+                "database_export/databases/imported_database.yaml", "w"
+            ) as fp:
+                fp.write(yaml.safe_dump(masked_database_config).encode())
+        buf.seek(0)
+
+        dummy_encrypted_extra = {"key": "value"}
+        form_data = {
+            "formData": (buf, "database_export.zip"),
+            "encrypted_extras": json.dumps(
+                {"databases/imported_database.yaml": json.dumps(dummy_encrypted_extra)}
+            ),
+        }
+        rv = self.client.post(uri, data=form_data, content_type="multipart/form-data")
+        response = json.loads(rv.data.decode("utf-8"))
+
+        assert rv.status_code == 200
+        assert response == {"message": "OK"}
+
+        database = (
+            db.session.query(Database).filter_by(uuid=database_config["uuid"]).one()
+        )
+        assert database.database_name == "imported_database"
+        assert (
+            database.sqlalchemy_uri == "vertica+vertica_python://host:5433/dbname?ssl=1"
+        )
+        assert json.loads(database.encrypted_extra) == dummy_encrypted_extra  # noqa: S105
+
+        db.session.delete(database)
+        db.session.commit()
+
     @mock.patch("superset.databases.schemas.is_feature_enabled")
     @mock.patch("superset.commands.database.importers.v1.utils.add_permissions")
     def test_import_database_masked_ssh_tunnel_password(

tests/unit_tests/importexport/api_test.py

@@ -109,6 +109,7 @@ def test_import_assets(
     ImportAssetsCommand.assert_called_with(
         mocked_contents,
         passwords=passwords,
+        encrypted_extras=None,
         ssh_tunnel_passwords=None,
         ssh_tunnel_private_keys=None,
         ssh_tunnel_priv_key_passwords=None,