Skip to content

Commit

Permalink
add support to new download files endpoints (#1550)
Browse files Browse the repository at this point in the history
  • Loading branch information
korgan00 authored Dec 13, 2024
1 parent 1d0803e commit ac33aab
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 24 deletions.
16 changes: 14 additions & 2 deletions client/qiskit_serverless/core/clients/serverless_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
84 changes: 62 additions & 22 deletions client/qiskit_serverless/core/files.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down

0 comments on commit ac33aab

Please sign in to comment.