upload datasets

dff/stats/__main__.py

@@ -105,11 +105,12 @@ def main(parsed_args: Optional[argparse.Namespace] = None):
     use_infile = hasattr(parsed_args, "infile") and parsed_args.infile is not None
     use_outfile = hasattr(parsed_args, "outfile") and parsed_args.outfile is not None
     if not use_infile:
-        file = make_zip_config(parsed_args)
+        file, datasets_file = make_zip_config(parsed_args)
     else:
         file = parsed_args.infile
+        datasets_file = None
 
-    import_dashboard(parsed_args, zip_file=str(file))
+    import_dashboard(parsed_args, zip_file=str(file), datasets_file=str(datasets_file) if datasets_file else None)
 
     if not use_infile and not use_outfile:
         file.unlink()

dff/stats/cli.py

@@ -144,7 +144,11 @@
 """
 
 
-def import_dashboard(parsed_args: Optional[argparse.Namespace] = None, zip_file: Optional[str] = None):
+def import_dashboard(
+    parsed_args: Optional[argparse.Namespace] = None,
+    zip_file: Optional[str] = None,
+    datasets_file: Optional[str] = None,
+):
     """
     Import an Apache Superset dashboard to a local instance with specified arguments.
     Before using the command, make sure you have your Superset instance
@@ -163,6 +167,7 @@ def import_dashboard(parsed_args: Optional[argparse.Namespace] = None, zip_file:
     session, headers = get_superset_session(parsed_args, superset_url)
     drop_superset_assets(session, headers, superset_url)
     import_dashboard_url = parse.urljoin(superset_url, "/api/v1/dashboard/import/")
+    import_dataset_url = parse.urljoin(superset_url, "/api/v1/dataset/import/")
     # upload files
     with open(zip_file, "rb") as f:
         response = session.request(
@@ -176,7 +181,25 @@ def import_dashboard(parsed_args: Optional[argparse.Namespace] = None, zip_file:
             files=[("formData", (zip_filename, f, "application/zip"))],
         )
         response.raise_for_status()
+
+    if datasets_file is None:
         logger.info(f"Upload finished with status {response.status_code}.")
+        return
+
+    with open(datasets_file, "rb") as f:
+        dataset_response = session.request(
+            "POST",
+            import_dataset_url,
+            headers=headers,
+            data={
+                "passwords": '{"databases/dff_database.yaml":"' + db_password + '"}',
+                "overwrite": "true",
+            },
+            files=[("formData", (zip_filename, f, "application/zip"))],
+        )
+        dataset_response.raise_for_status()
+
+    logger.info(f"Upload finished with status {dataset_response.status_code}.")
 
 
 def make_zip_config(parsed_args: argparse.Namespace) -> Path:
@@ -257,5 +280,13 @@ def make_zip_config(parsed_args: argparse.Namespace) -> Path:
             raise ValueError(f"Outfile name missing .zip extension: {outfile_name}.")
         logger.info(f"Saving the archive to {outfile_name}.")
         shutil.make_archive(outfile_name[: outfile_name.rindex(".zip")], format="zip", root_dir=temp_config_dir)
+        for path in Path(temp_config_dir).iterdir():
+            if path != database_dir and path != dataset_dir:
+                if path.is_dir():
+                    shutil.rmtree(path)
+
+        shutil.make_archive(
+            "datasets_" + outfile_name[: outfile_name.rindex(".zip")], format="zip", root_dir=temp_config_dir
+        )
 
-    return Path(outfile_name)
+    return Path(outfile_name), Path("datasets_" + outfile_name)