add support to new upload file endpoints

Qiskit · Dec 17, 2024 · 8e9f7a2 · 8e9f7a2
1 parent 47c3064
commit 8e9f7a2
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 4 deletions.
diff --git a/client/qiskit_serverless/core/clients/serverless_client.py b/client/qiskit_serverless/core/clients/serverless_client.py
@@ -433,9 +433,15 @@ def file_delete(self, file: str, provider: Optional[str] = None):
         """Deletes file uploaded or produced by the programs,"""
         return self._files_client.delete(file, provider)
 
-    def file_upload(self, file: str, provider: Optional[str] = None):
+    def file_upload(
+        self, file: str, function: QiskitFunction, provider: Optional[str] = None
+    ):
+        """Upload file."""
+        return self._files_client.upload(file, function, provider)
+
+    def provider_file_upload(self, file: str, function: QiskitFunction, provider: str):
         """Upload file."""
-        return self._files_client.upload(file, provider)
+        return self._files_client.provider_upload(file, function, provider)
 
 
 class IBMServerlessClient(ServerlessClient):

diff --git a/client/qiskit_serverless/core/files.py b/client/qiskit_serverless/core/files.py
@@ -130,15 +130,37 @@ def provider_download(
                 os.path.join(self._files_url, "provider", "download"),
             )
 
-    def upload(self, file: str, provider: Optional[str] = None) -> Optional[str]:
+    def upload(
+        self, file: str, function: QiskitFunction, provider: Optional[str] = None
+    ) -> Optional[str]:
         """Uploads file."""
         tracer = trace.get_tracer("client.tracer")
         with tracer.start_as_current_span("files.upload"):
             with open(file, "rb") as f:
                 with requests.post(
                     os.path.join(self._files_url, "upload"),
                     files={"file": f},
-                    data={"provider": provider},
+                    data={"provider": provider, "function": function.title},
+                    stream=True,
+                    headers={"Authorization": f"Bearer {self._token}"},
+                    timeout=REQUESTS_STREAMING_TIMEOUT,
+                ) as req:
+                    if req.ok:
+                        return req.text
+                    return "Upload failed"
+            return "Can not open file"
+
+    def provider_upload(
+        self, file: str, function: QiskitFunction, provider: str
+    ) -> Optional[str]:
+        """Uploads file to provider/function file storage."""
+        tracer = trace.get_tracer("client.tracer")
+        with tracer.start_as_current_span("files.upload"):
+            with open(file, "rb") as f:
+                with requests.post(
+                    os.path.join(self._files_url, "upload"),
+                    files={"file": f},
+                    data={"provider": provider, "function": function.title},
                     stream=True,
                     headers={"Authorization": f"Bearer {self._token}"},
                     timeout=REQUESTS_STREAMING_TIMEOUT,