Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat(drs): optional session_kwargs for aiohttp.ClientSession #220

Merged
merged 1 commit into from
Aug 15, 2024
Merged

feat(drs): optional session_kwargs for aiohttp.ClientSession #220

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions bento_lib/drs/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,8 +47,7 @@ def decode_drs_uri(drs_uri: str, internal_drs_base_url: Optional[str] = None) ->
parsed_drs_uri = urlparse(drs_uri)

if parsed_drs_uri.scheme != "drs":
print(f"[Bento Lib] Invalid scheme: '{parsed_drs_uri.scheme}'",
file=sys.stderr, flush=True)
print(f"[Bento Lib] Invalid scheme: '{parsed_drs_uri.scheme}'", file=sys.stderr, flush=True)
raise DrsInvalidScheme(f"Encountered invalid DRS scheme: {parsed_drs_uri.scheme}")

drs_base_path = internal_drs_base_url.rstrip("/") if internal_drs_base_url else f"https://{parsed_drs_uri.netloc}"
Expand Down Expand Up @@ -82,12 +81,17 @@ def fetch_drs_record_by_uri(drs_uri: str, internal_drs_base_url: Optional[str] =
return drs_res.json()


async def fetch_drs_record_by_uri_async(drs_uri: str, internal_drs_base_url: Optional[str] = None) -> Optional[dict]:
async def fetch_drs_record_by_uri_async(
drs_uri: str,
internal_drs_base_url: Optional[str] = None,
session_kwargs: dict | None = None,
) -> Optional[dict]:
"""
Given a URI in the format drs://<hostname>/<object-id>, decodes it into an
HTTP URL and asynchronously fetches the object metadata.
:param drs_uri: The URI of the object to fetch.
:param internal_drs_base_url: An optional override hard-coded DRS base URL to use, for container networking etc.
:param session_kwargs: Optional dictionary of parameters to pass to the aiohttp.ClientSession constructor.
:return: The fetched DRS object metadata.
"""

Expand All @@ -97,7 +101,7 @@ async def fetch_drs_record_by_uri_async(drs_uri: str, internal_drs_base_url: Opt
print(f"[Bento Lib] Attempting to fetch {decoded_object_uri}", flush=True)

params = {"internal_path": "true"} if internal_drs_base_url else {}
async with aiohttp.ClientSession() as session:
async with aiohttp.ClientSession(**(session_kwargs or {})) as session:
async with session.get(decoded_object_uri, params=params) as drs_res:
if drs_res.status != 200:
print(f"[Bento Lib] Could not fetch: '{decoded_object_uri}'", file=sys.stderr, flush=True)
Expand Down