diff --git a/superset/cli/importexport.py b/superset/cli/importexport.py index 54322aaec2b9..50c8aa78a7a6 100755 --- a/superset/cli/importexport.py +++ b/superset/cli/importexport.py @@ -28,6 +28,7 @@ from superset import security_manager from superset.extensions import db +from superset.utils.core import override_user logger = logging.getLogger(__name__) @@ -170,26 +171,34 @@ def import_dashboards(path: str, username: Optional[str]) -> None: "-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) @click.command()