This is an automated email from the ASF dual-hosted git repository.
rusackas pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new c9e2c7037e feat: Implement sparse import for ImportAssetsCommand
(#32670)
c9e2c7037e is described below
commit c9e2c7037e81ffe31e83bba70f78f865f52f973d
Author: Paul Rhodes <[email protected]>
AuthorDate: Mon Mar 17 14:44:15 2025 +0000
feat: Implement sparse import for ImportAssetsCommand (#32670)
---
superset/commands/importers/v1/assets.py | 25 +++++++++++++++++++++----
superset/commands/importers/v1/utils.py | 18 +++++++++++++++++-
superset/importexport/api.py | 5 +++++
tests/unit_tests/importexport/api_test.py | 1 +
4 files changed, 44 insertions(+), 5 deletions(-)
diff --git a/superset/commands/importers/v1/assets.py
b/superset/commands/importers/v1/assets.py
index 41895f84d7..3b0442157f 100644
--- a/superset/commands/importers/v1/assets.py
+++ b/superset/commands/importers/v1/assets.py
@@ -34,17 +34,21 @@ from superset.commands.database.importers.v1.utils import
import_database
from superset.commands.dataset.importers.v1.utils import import_dataset
from superset.commands.exceptions import CommandInvalidError, ImportFailedError
from superset.commands.importers.v1.utils import (
+ get_resource_mappings_batched,
load_configs,
load_metadata,
validate_metadata_type,
)
from superset.commands.query.importers.v1.utils import import_saved_query
from superset.commands.utils import update_chart_config_dataset
+from superset.connectors.sqla.models import SqlaTable
from superset.dashboards.schemas import ImportV1DashboardSchema
from superset.databases.schemas import ImportV1DatabaseSchema
from superset.datasets.schemas import ImportV1DatasetSchema
from superset.migrations.shared.native_filters import migrate_dashboard
+from superset.models.core import Database
from superset.models.dashboard import dashboard_slices
+from superset.models.slice import Slice
from superset.queries.saved_queries.schemas import ImportV1SavedQuerySchema
from superset.utils.decorators import on_error, transaction
@@ -79,12 +83,27 @@ class ImportAssetsCommand(BaseCommand):
kwargs.get("ssh_tunnel_priv_key_passwords") or {}
)
self._configs: dict[str, Any] = {}
+ self.sparse = kwargs.get("sparse", False)
# pylint: disable=too-many-locals
@staticmethod
- def _import(configs: dict[str, Any]) -> None: # noqa: C901
+ def _import(configs: dict[str, Any], sparse: bool = False) -> None: #
noqa: C901
# import databases first
database_ids: dict[str, int] = {}
+ dataset_info: dict[str, dict[str, Any]] = {}
+ chart_ids: dict[str, int] = {}
+ if sparse:
+ chart_ids = get_resource_mappings_batched(Slice)
+ database_ids = get_resource_mappings_batched(Database)
+ dataset_info = get_resource_mappings_batched(
+ SqlaTable,
+ value_func=lambda x: {
+ "datasource_id": x.id,
+ "datasource_type": x.datasource_type,
+ "datasource_name": x.datasource_name,
+ },
+ )
+
for file_name, config in configs.items():
if file_name.startswith("databases/"):
database = import_database(config, overwrite=True)
@@ -97,7 +116,6 @@ class ImportAssetsCommand(BaseCommand):
import_saved_query(config, overwrite=True)
# import datasets
- dataset_info: dict[str, dict[str, Any]] = {}
for file_name, config in configs.items():
if file_name.startswith("datasets/"):
config["database_id"] = database_ids[config["database_uuid"]]
@@ -110,7 +128,6 @@ class ImportAssetsCommand(BaseCommand):
# import charts
charts = []
- chart_ids: dict[str, int] = {}
for file_name, config in configs.items():
if file_name.startswith("charts/"):
dataset_dict = dataset_info[config["dataset_uuid"]]
@@ -161,7 +178,7 @@ class ImportAssetsCommand(BaseCommand):
)
def run(self) -> None:
self.validate()
- self._import(self._configs)
+ self._import(self._configs, self.sparse)
def validate(self) -> None:
exceptions: list[ValidationError] = []
diff --git a/superset/commands/importers/v1/utils.py
b/superset/commands/importers/v1/utils.py
index ec5387708f..7d4e590913 100644
--- a/superset/commands/importers/v1/utils.py
+++ b/superset/commands/importers/v1/utils.py
@@ -15,7 +15,7 @@
import logging
from pathlib import Path, PurePosixPath
-from typing import Any, Optional
+from typing import Any, Callable, Dict, Optional, Type
from zipfile import ZipFile
import yaml
@@ -214,3 +214,19 @@ def get_contents_from_bundle(bundle: ZipFile) -> dict[str,
str]:
for file_name in bundle.namelist()
if is_valid_config(file_name)
}
+
+
+def get_resource_mappings_batched(
+ model_class: Type[Any],
+ batch_size: int = 1000,
+ value_func: Callable[[Any], Any] = lambda x: x.id,
+) -> Dict[str, Any]:
+ offset = 0
+ mapping = {}
+ while True:
+ batch =
db.session.query(model_class).limit(batch_size).offset(offset).all()
+ if not batch:
+ break
+ mapping.update({str(x.uuid): value_func(x) for x in batch})
+ offset += batch_size
+ return mapping
diff --git a/superset/importexport/api.py b/superset/importexport/api.py
index 0d3906b1a1..0a6a27dd58 100644
--- a/superset/importexport/api.py
+++ b/superset/importexport/api.py
@@ -147,6 +147,9 @@ class ImportExportRestApi(BaseSupersetApi):
the private_key should be provided in the following
format:
`{"databases/MyDatabase.yaml":
"my_private_key_password"}`.
type: string
+ sparse:
+ description: allow sparse update of resources
+ type: boolean
responses:
200:
description: Assets import result
@@ -177,6 +180,7 @@ class ImportExportRestApi(BaseSupersetApi):
if not contents:
raise NoValidFilesFoundError()
+ sparse = request.form.get("sparse") == "true"
passwords = (
json.loads(request.form["passwords"])
@@ -201,6 +205,7 @@ class ImportExportRestApi(BaseSupersetApi):
command = ImportAssetsCommand(
contents,
+ sparse=sparse,
passwords=passwords,
ssh_tunnel_passwords=ssh_tunnel_passwords,
ssh_tunnel_private_keys=ssh_tunnel_private_keys,
diff --git a/tests/unit_tests/importexport/api_test.py
b/tests/unit_tests/importexport/api_test.py
index 0c125c2ccf..1458631e5d 100644
--- a/tests/unit_tests/importexport/api_test.py
+++ b/tests/unit_tests/importexport/api_test.py
@@ -108,6 +108,7 @@ def test_import_assets(
passwords = {"assets_export/databases/imported_database.yaml": "SECRET"}
ImportAssetsCommand.assert_called_with(
mocked_contents,
+ sparse=False,
passwords=passwords,
ssh_tunnel_passwords=None,
ssh_tunnel_private_keys=None,
