Merge pull request #15 from code-kern-ai/monitoring-upload

Monitoring upload

Author: Johannes Hötter

refinery/__init__.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 
+from black import Any
 from wasabi import msg
 import pandas as pd
 from refinery import authentication, api_calls, settings, exceptions, util
@@ -8,6 +9,8 @@
 import os.path
 from tqdm import tqdm
 import spacy
+import time
+from refinery import settings
 
 
 class Client:
@@ -148,7 +151,9 @@ def get_record_export(
                     )
 
             else:
-                msg.warn("There are no attributes that can be tokenized in this project.")
+                msg.warn(
+                    "There are no attributes that can be tokenized in this project."
+                )
 
         if download_to is not None:
             df.to_json(download_to, orient="records")
@@ -214,8 +219,54 @@ def post_file_import(
             file_name,
         )
         if success:
-            msg.good(f"Uploaded {path} to your project.")
-            return True
+            msg.good(f"Uploaded {path} to object storage.")
+            upload_task_id = (
+                upload_task_id.split("/")[-1]
+                if "/" in upload_task_id
+                else upload_task_id
+            )
+            self.__monitor_task(upload_task_id)
+
+        else:
+            msg_text = f"Could not upload {path} to your project."
+            msg.fail(msg_text)
+            raise exceptions.FileImportError(msg_text)
+
+    def __monitor_task(self, upload_task_id: str) -> None:
+        do_monitoring = True
+        idx = 0
+        last_progress = 0.0
+        print_success_message = False
+        with tqdm(
+            total=100.00,
+            colour="green",
+            bar_format="{desc}: {percentage:.2f}%|{bar:10}| {n:.2f}/{total_fmt}",
+        ) as pbar:
+            pbar.set_description_str(desc="PENDING", refresh=True)
+            while do_monitoring:
+                idx += 1
+                task = self.__get_task(upload_task_id)
+                task_progress = task.get("progress") if task.get("progress") else 0.0
+                task_state = task.get("state") if task.get("state") else "FAILED"
+                progress = task_progress - last_progress
+                last_progress = task_progress
+                pbar.update(progress)
+                pbar.set_description_str(desc=task_state, refresh=True)
+                if task_state == "DONE" or task_state == "FAILED":
+                    print_success_message = task_state == "DONE"
+                    do_monitoring = False
+                if idx >= 100:
+                    raise exceptions.FileImportError(
+                        "Timeout while upload, please check the upload progress in the UI."
+                    )
+                time.sleep(0.5)
+        if print_success_message:
+            msg.good("File upload successful.")
         else:
-            msg.fail(f"Could not upload {path} to your project.")
-            return False
+            msg.fail("Upload failed. Please look into the UI notification center for more details.")
+
+    def __get_task(self, upload_task_id: str) -> Dict[str, Any]:
+        api_response = api_calls.get_request(
+            settings.get_task(self.project_id, upload_task_id), self.session_token
+        )
+        return api_response

refinery/settings.py

@@ -46,3 +46,7 @@ def get_import_url(project_id: str) -> str:
 
 def get_base_config(project_id: str) -> str:
     return f"{get_project_url(project_id)}/import/base_config"
+
+
+def get_task(project_id: str, task_id: str) -> str:
+    return f"{get_project_url(project_id)}/import/task/{task_id}"