Fix download export

Author: dshabin

src/superannotate/lib/app/interface/sdk_interface.py

@@ -2153,50 +2153,13 @@ def download_export(
     """
     project_name, folder_name = extract_project_folder(project)
     export_name = export["name"] if isinstance(export, dict) else export
-
-    if to_s3_bucket:
-        with tempfile.TemporaryDirectory() as tmp:
-            response = controller.download_export(
-                project_name=project_name,
-                export_name=export_name,
-                folder_path=tmp,
-                extract_zip_contents=extract_zip_contents,
-                to_s3_bucket=to_s3_bucket,
-            )
-            downloaded_folder_path = response.data
-            if to_s3_bucket:
-                to_s3_bucket = boto3.Session().resource("s3").Bucket(to_s3_bucket)
-                files_to_upload = []
-                for file in Path(downloaded_folder_path).rglob("*.*"):
-                    files_to_upload.append(file)
-
-                def _upload_file_to_s3(to_s3_bucket, path, s3_key) -> None:
-                    controller.upload_file_to_s3(
-                        to_s3_bucket=to_s3_bucket, path=path, s3_key=s3_key
-                    )
-
-                with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
-                    results = []
-                    for path in files_to_upload:
-                        s3_key = f"{folder_path}/{path.name}"
-                        results.append(
-                            executor.submit(
-                                _upload_file_to_s3, to_s3_bucket, str(path), s3_key
-                            )
-                        )
-                    for future in concurrent.futures.as_completed(results):
-                        future.result()
-                logger.info(
-                    "Exported to AWS %s/%s", to_s3_bucket.name, str(folder_path)
-                )
-    else:
-        controller.download_export(
-            project_name=project_name,
-            export_name=export_name,
-            folder_path=folder_path,
-            extract_zip_contents=extract_zip_contents,
-            to_s3_bucket=to_s3_bucket,
-        )
+    controller.download_export(
+        project_name=project_name,
+        export_name=export_name,
+        folder_path=folder_path,
+        extract_zip_contents=extract_zip_contents,
+        to_s3_bucket=to_s3_bucket,
+    )
 
 
 @Trackable

src/superannotate/lib/core/usecases.py

@@ -21,6 +21,7 @@
 import numpy as np
 import pandas as pd
 import requests
+import tempfile
 from boto3.exceptions import Boto3Error
 from lib.app.analytics.common import aggregate_annotations_as_df
 from lib.app.analytics.common import consensus_plot
@@ -3966,51 +3967,76 @@ def validate_project_type(self):
                 constances.LIMITED_FUNCTIONS[self._project.project_type]
             )
 
-    def execute(self):
-        if self.is_valid():
-            exports = self._service.get_exports(
-                team_id=self._project.team_id, project_id=self._project.uuid
-            )
-            export_id = None
-            for export in exports:
-                if export["name"] == self._export_name:
-                    export_id = export["id"]
-                    break
-            if not export_id:
-                raise AppException("Export not found.")
+    def upload_to_s3_from_folder(self, folder_path: str):
+        to_s3_bucket = boto3.Session().resource("s3").Bucket(self._to_s3_bucket)
+        files_to_upload = list(Path(folder_path).rglob("*.*"))
 
-            while True:
-                export = self._service.get_export(
-                    team_id=self._project.team_id,
-                    project_id=self._project.uuid,
-                    export_id=export_id,
+        def _upload_file_to_s3(_to_s3_bucket, _path, _s3_key) -> None:
+            _to_s3_bucket.upload_file(_path, _s3_key)
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
+            results = []
+            for path in files_to_upload:
+                s3_key = f"{self._folder_path}/{path.name}"
+                results.append(
+                    executor.submit(
+                        _upload_file_to_s3, to_s3_bucket, str(path), s3_key
+                    )
                 )
 
-                if export["status"] == ExportStatus.IN_PROGRESS.value:
-                    logger.info("Waiting 5 seconds for export to finish on server.")
-                    time.sleep(5)
-                    continue
-                if export["status"] == ExportStatus.ERROR.value:
-                    raise AppException("Couldn't download export.")
-                    pass
-                break
-
-            filename = Path(export["path"]).name
-            filepath = Path(self._folder_path) / filename
-            with requests.get(export["download"], stream=True) as r:
-                r.raise_for_status()
-                with open(filepath, "wb") as f:
-                    for chunk in r.iter_content(chunk_size=8192):
-                        f.write(chunk)
-            if self._extract_zip_contents:
-                with zipfile.ZipFile(filepath, "r") as f:
-                    f.extractall(self._folder_path)
-                Path.unlink(filepath)
-                logger.info(f"Extracted {filepath} to folder {self._folder_path}")
-            else:
-                logger.info(f"Downloaded export ID {export['id']} to {filepath}")
+    def download_to_local_storage(self, destination: str):
+        exports = self._service.get_exports(
+            team_id=self._project.team_id, project_id=self._project.uuid
+        )
+        export = next(filter(lambda i: i["name"] == self._export_name, exports), None)
+        export = self._service.get_export(
+            team_id=self._project.team_id,
+            project_id=self._project.uuid,
+            export_id=export["id"],
+        )
+        if not export:
+            raise AppException("Export not found.")
+        export_status = export["status"]
+
+        while export_status != ExportStatus.COMPLETE.value:
+            logger.info("Waiting 5 seconds for export to finish on server.")
+            time.sleep(5)
+
+            export = self._service.get_export(
+                team_id=self._project.team_id,
+                project_id=self._project.uuid,
+                export_id=export["id"],
+            )
+            export_status = export["status"]
+            if export_status in (ExportStatus.ERROR.value, ExportStatus.CANCELED.value):
+                raise AppException("Couldn't download export.")
+
+        filename = Path(export["path"]).name
+        filepath = Path(destination) / filename
+        with requests.get(export["download"], stream=True) as response:
+            response.raise_for_status()
+            with open(filepath, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+        if self._extract_zip_contents:
+            with zipfile.ZipFile(filepath, "r") as f:
+                f.extractall(destination)
+            Path.unlink(filepath)
+        return export["id"], filepath, destination
 
-            self._response.data = self._folder_path
+    def execute(self):
+        if self.is_valid():
+            if self._to_s3_bucket:
+                with tempfile.TemporaryDirectory() as tmp:
+                    self.download_to_local_storage(tmp)
+                    self.upload_to_s3_from_folder(tmp)
+                logger.info(f"Exported to AWS {self._to_s3_bucket}/{self._folder_path}")
+            else:
+                export_id, filepath, destination = self.download_to_local_storage(self._folder_path)
+                if self._extract_zip_contents:
+                    logger.info(f"Extracted {filepath} to folder {destination}")
+                else:
+                    logger.info(f"Downloaded export ID {export_id} to {filepath}")
         return self._response