(superset) 07/08: fix(import-datasources): Use "admin" user as default for importing datasources (#27154)

[mi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

michaelsmolina pushed a commit to branch 3.1
in repository https://gitbox.apache.org/repos/asf/superset.git

commit dc493fe7afeab11464c435b09b6da8b6e1caaf32
Author: James O'Claire <ja...@gmail.com>
AuthorDate: Tue Feb 27 23:32:04 2024 +0800

    fix(import-datasources): Use "admin" user as default for importing datasources (#27154)
---
 superset/cli/importexport.py | 39 ++++++++++++++++++++++++---------------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/superset/cli/importexport.py b/superset/cli/importexport.py
index 5b19cc9ffc..56abcc31fe 100755
--- a/superset/cli/importexport.py
+++ b/superset/cli/importexport.py
@@ -29,6 +29,7 @@ from flask.cli import with_appcontext
 from superset import security_manager
 from superset.cli.lib import feature_flags
 from superset.extensions import db
+from superset.utils.core import override_user
 
 logger = logging.getLogger(__name__)
 
@@ -174,26 +175,34 @@ if feature_flags.get("VERSIONED_EXPORT"):
         "-p",
         help="Path to a single ZIP file",
     )
-    def import_datasources(path: str) -> None:
+    @click.option(
+        "--username",
+        "-u",
+        required=False,
+        default="admin",
+        help="Specify the user name to assign datasources to",
+    )
+    def import_datasources(path: str, username: Optional[str] = "admin") -> None:
         """Import datasources from ZIP file"""
         # pylint: disable=import-outside-toplevel
         from superset.commands.dataset.importers.dispatcher import ImportDatasetsCommand
         from superset.commands.importers.v1.utils import get_contents_from_bundle
 
-        if is_zipfile(path):
-            with ZipFile(path) as bundle:
-                contents = get_contents_from_bundle(bundle)
-        else:
-            with open(path) as file:
-                contents = {path: file.read()}
-        try:
-            ImportDatasetsCommand(contents, overwrite=True).run()
-        except Exception:  # pylint: disable=broad-except
-            logger.exception(
-                "There was an error when importing the dataset(s), please check the "
-                "exception traceback in the log"
-            )
-            sys.exit(1)
+        with override_user(user=security_manager.find_user(username=username)):
+            if is_zipfile(path):
+                with ZipFile(path) as bundle:
+                    contents = get_contents_from_bundle(bundle)
+            else:
+                with open(path) as file:
+                    contents = {path: file.read()}
+            try:
+                ImportDatasetsCommand(contents, overwrite=True).run()
+            except Exception:  # pylint: disable=broad-except
+                logger.exception(
+                    "There was an error when importing the dataset(s), please check the "
+                    "exception traceback in the log"
+                )
+                sys.exit(1)
 
 else: