From ac33aab93c0a08ccf2ef6f1d69e6184580a2fbb5 Mon Sep 17 00:00:00 2001
From: Goyo <korgan00@gmail.com>
Date: Fri, 13 Dec 2024 17:12:25 +0100
Subject: [PATCH] add support to new download files endpoints (#1550)

---
 .../core/clients/serverless_client.py         | 16 +++-
 client/qiskit_serverless/core/files.py        | 84 ++++++++++++++-----
 2 files changed, 76 insertions(+), 24 deletions(-)

diff --git a/client/qiskit_serverless/core/clients/serverless_client.py b/client/qiskit_serverless/core/clients/serverless_client.py
index efec48476..5a4459f60 100644
--- a/client/qiskit_serverless/core/clients/serverless_client.py
+++ b/client/qiskit_serverless/core/clients/serverless_client.py
@@ -408,13 +408,25 @@ def provider_files(self, function: QiskitFunction) -> List[str]:
     def file_download(
         self,
         file: str,
+        function: QiskitFunction,
         target_name: Optional[str] = None,
         download_location: str = "./",
-        provider: Optional[str] = None,
     ):
         """Download file."""
         return self._files_client.download(
-            file, download_location, target_name, provider
+            file, download_location, function, target_name
+        )
+
+    def provider_file_download(
+        self,
+        file: str,
+        function: QiskitFunction,
+        target_name: Optional[str] = None,
+        download_location: str = "./",
+    ):
+        """Download file."""
+        return self._files_client.provider_download(
+            file, download_location, function, target_name
         )
 
     def file_delete(self, file: str, provider: Optional[str] = None):
diff --git a/client/qiskit_serverless/core/files.py b/client/qiskit_serverless/core/files.py
index 7d9677f88..26ecd1720 100644
--- a/client/qiskit_serverless/core/files.py
+++ b/client/qiskit_serverless/core/files.py
@@ -58,37 +58,77 @@ def __init__(self, host: str, token: str, version: str):
         self._token = token
         self._files_url = os.path.join(self.host, "api", self.version, "files")
 
+    def _download_with_url(  # pylint:  disable=too-many-positional-arguments
+        self,
+        file: str,
+        download_location: str,
+        function: QiskitFunction,
+        url: str,
+        target_name: Optional[str] = None,
+    ) -> Optional[str]:
+        """Auxiliar function to download a file using an url."""
+        with requests.get(
+            url,
+            params={
+                "file": file,
+                "provider": function.provider,
+                "function": function.title,
+            },
+            stream=True,
+            headers={"Authorization": f"Bearer {self._token}"},
+            timeout=REQUESTS_STREAMING_TIMEOUT,
+        ) as req:
+            req.raise_for_status()
+
+            total_size_in_bytes = int(req.headers.get("content-length", 0))
+            chunk_size = 8192
+            progress_bar = tqdm(total=total_size_in_bytes, unit="iB", unit_scale=True)
+            file_name = target_name or f"downloaded_{str(uuid.uuid4())[:8]}_{file}"
+            with open(os.path.join(download_location, file_name), "wb") as f:
+                for chunk in req.iter_content(chunk_size=chunk_size):
+                    progress_bar.update(len(chunk))
+                    f.write(chunk)
+            progress_bar.close()
+            return file_name
+
     def download(
         self,
         file: str,
         download_location: str,
+        function: QiskitFunction,
         target_name: Optional[str] = None,
-        provider: Optional[str] = None,
     ) -> Optional[str]:
-        """Downloads file."""
+        """Downloads user file."""
         tracer = trace.get_tracer("client.tracer")
         with tracer.start_as_current_span("files.download"):
-            with requests.get(
+            return self._download_with_url(
+                file,
+                download_location,
+                function,
+                target_name,
                 os.path.join(self._files_url, "download"),
-                params={"file": file, "provider": provider},
-                stream=True,
-                headers={"Authorization": f"Bearer {self._token}"},
-                timeout=REQUESTS_STREAMING_TIMEOUT,
-            ) as req:
-                req.raise_for_status()
-
-                total_size_in_bytes = int(req.headers.get("content-length", 0))
-                chunk_size = 8192
-                progress_bar = tqdm(
-                    total=total_size_in_bytes, unit="iB", unit_scale=True
-                )
-                file_name = target_name or f"downloaded_{str(uuid.uuid4())[:8]}_{file}"
-                with open(os.path.join(download_location, file_name), "wb") as f:
-                    for chunk in req.iter_content(chunk_size=chunk_size):
-                        progress_bar.update(len(chunk))
-                        f.write(chunk)
-                progress_bar.close()
-                return file_name
+            )
+
+    def provider_download(
+        self,
+        file: str,
+        download_location: str,
+        function: QiskitFunction,
+        target_name: Optional[str] = None,
+    ) -> Optional[str]:
+        """Downloads provider file."""
+        if not function.provider:
+            raise QiskitServerlessException("`function` doesn't have a provider.")
+
+        tracer = trace.get_tracer("client.tracer")
+        with tracer.start_as_current_span("files.provider_download"):
+            return self._download_with_url(
+                file,
+                download_location,
+                function,
+                target_name,
+                os.path.join(self._files_url, "provider", "download"),
+            )
 
     def upload(self, file: str, provider: Optional[str] = None) -> Optional[str]:
         """Uploads file."""