From 65f4d7511256733b15e857a4026bc0e3166a5c01 Mon Sep 17 00:00:00 2001
From: Martin Molinero <martin.molinero1@gmail.com>
Date: Mon, 8 Apr 2024 10:28:19 -0300
Subject: [PATCH] WIP

---
 lean/commands/cloud/object_store/get.py       | 30 +++++++++++--------
 .../commands/cloud/object_store/properties.py |  2 +-
 lean/components/api/object_store_client.py    |  6 ++--
 3 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/lean/commands/cloud/object_store/get.py b/lean/commands/cloud/object_store/get.py
index fee35d3a..852c1def 100644
--- a/lean/commands/cloud/object_store/get.py
+++ b/lean/commands/cloud/object_store/get.py
@@ -10,15 +10,17 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from uuid import uuid4
+from os import path, getcwd, unlink, mkdir
 
-from click import command, argument
+from click import command, option
 from lean.click import LeanCommand
 from lean.container import container
 
 
 @command(cls=LeanCommand)
-@argument("key", type=str, multiple=True, help=f"The desired key to fetch")
-@argument("destination_folder", type=str, default="",
+@option("--key", type=str, multiple=True, help=f"The desired key to fetch")
+@option("--destination_folder", type=str, default="",
           help=f"The destination folder to download the object store values,"
                f" if not provided will use to current directory")
 def get(key: [str], destination_folder: str):
@@ -30,24 +32,28 @@ def get(key: [str], destination_folder: str):
     api_client = container.api_client
     logger = container.logger
 
-    logger.debug(f"Fetch object store download url")
-    url = api_client.object_store.get(key, organization_id)
+    logger.info(f"Fetching object store download url")
+    url = api_client.object_store.get(key, organization_id, logger)
 
-    logger.debug(f"Start downloading: {url}")
     progress = logger.progress(suffix="{task.percentage:0.0f}% ({task.completed:,.0f}/{task.total:,.0f})")
     progress_task = progress.add_task("", total=1)
 
-    from uuid import uuid4
-    temp_file = f"{str(uuid4())}.zip"
-    api_client.data.download_url(url, temp_file, lambda advance: progress.update(progress_task, advance=advance))
-
     if not destination_folder:
-        from os import getcwd
         destination_folder = getcwd()
 
-    logger.debug(f"Unzipping object store keys values into: '{destination_folder}'")
+    if not path.exists(destination_folder):
+        mkdir(destination_folder)
+
+    temp_file = path.join(destination_folder, f"{str(uuid4())}.zip")
+    logger.info(f"Start downloading: {url} into {temp_file}")
+    api_client.data.download_url(url, temp_file, lambda advance: progress.update(progress_task, advance=advance))
+
+    logger.info(f"Unzipping object store keys values into: '{destination_folder}'")
     from zipfile import ZipFile
     with ZipFile(temp_file, 'r') as zip_ref:
         zip_ref.extractall(destination_folder)
 
+    logger.info(f"Deleting temp file: '{temp_file}'")
+    unlink(temp_file)
+
 
diff --git a/lean/commands/cloud/object_store/properties.py b/lean/commands/cloud/object_store/properties.py
index fe321fb7..a001510b 100644
--- a/lean/commands/cloud/object_store/properties.py
+++ b/lean/commands/cloud/object_store/properties.py
@@ -18,7 +18,7 @@
 
 @command(cls=LeanCommand)
 @argument("key", type=str)
-def properties(key: str) -> str:
+def properties(key: str):
     """
     Get a value properties from the organization's cloud object store.
 
diff --git a/lean/components/api/object_store_client.py b/lean/components/api/object_store_client.py
index e1af31c3..5f6ca455 100644
--- a/lean/components/api/object_store_client.py
+++ b/lean/components/api/object_store_client.py
@@ -12,7 +12,7 @@
 # limitations under the License.
 
 from lean.components.api.api_client import *
-
+from lean.components.util.logger import Logger
 
 class ObjectStoreClient:
     """The ObjectStoreClient class contains methods to interact with object/* API endpoints."""
@@ -38,11 +38,12 @@ def properties(self, key: str, organization_id: str) -> str:
 
         return self._api.post("object/properties", payload)
 
-    def get(self, keys: [str], organization_id: str) -> str:
+    def get(self, keys: [str], organization_id: str, logger: Logger) -> str:
         """Will fetch an url to download the requested keys
 
         :param keys: keys of the object to retrieve
         :param organization_id: the id of the organization who's object store to retrieve from
+        :param logger: the logger instance to use
         :return: the url to download the requested keys
         """
 
@@ -58,6 +59,7 @@ def get(self, keys: [str], organization_id: str) -> str:
 
         job_id = data["jobId"]
         if job_id:
+            logger.info(f"Waiting for files to be ready for download...")
             from time import sleep
             payload = {
                 "organizationId": organization_id,