diff --git a/refinery/__init__.py b/refinery/__init__.py index 036d372..6e70e9a 100644 --- a/refinery/__init__.py +++ b/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 diff --git a/refinery/settings.py b/refinery/settings.py index 0adcb50..64c7c5a 100644 --- a/refinery/settings.py +++ b/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}"