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.

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

Add minimal JAWS client #152

Merged
merged 1 commit into from
Jan 7, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
69 changes: 69 additions & 0 deletions cdmtaskservice/jaws/client.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
"""
An minimal async client for the JAWS central server.
"""

import aiohttp
import datetime
import logging
from typing import Any

from cdmtaskservice.arg_checkers import require_string as _require_string


class JAWSClient:
""" The JAWS client. """

@classmethod
async def create(self, url: str, token: str):
"""
Initialize the client.

url - the JAWS Central URL.
token - the JAWS token.
"""
cli = JAWSClient(url, token)
user = await cli._user() # test connection & token
logging.getLogger(__name__).info(f"Initialized JAWS client with user {user}")
return cli

Check warning on line 27 in cdmtaskservice/jaws/client.py

View check run for this annotation

Codecov / codecov/patch

cdmtaskservice/jaws/client.py#L24-L27

Added lines #L24 - L27 were not covered by tests

def __init__(self, url: str, token: str):
url = _require_string(url, "url")
if not url.endswith("/"):
url += "/"
self._sess = aiohttp.ClientSession(

Check warning on line 33 in cdmtaskservice/jaws/client.py

View check run for this annotation

Codecov / codecov/patch

cdmtaskservice/jaws/client.py#L30-L33

Added lines #L30 - L33 were not covered by tests
base_url=url,
headers={"Authorization": f"Bearer {_require_string(token, 'token')}"}
)

async def _get(self, url, params=None) -> dict[str, Any]:
async with self._sess.get(url, params=params) as res:

Check warning on line 39 in cdmtaskservice/jaws/client.py

View check run for this annotation

Codecov / codecov/patch

cdmtaskservice/jaws/client.py#L39

Added line #L39 was not covered by tests
# Any jaws errors would be 500 errors since should just be querying known jobs, so
# don't worry too much about exceptions. Expand later if needed
# May need to add retries
# May need to to add some sort of down notification or detection
res.raise_for_status()
if res.ok: # assume here that if we get a 2XX we get JSON. Fix if necessary
return await res.json()

Check warning on line 46 in cdmtaskservice/jaws/client.py

View check run for this annotation

Codecov / codecov/patch

cdmtaskservice/jaws/client.py#L44-L46

Added lines #L44 - L46 were not covered by tests

async def _user(self) -> str:
res = await self._get("user")
return res["uid"]

Check warning on line 50 in cdmtaskservice/jaws/client.py

View check run for this annotation

Codecov / codecov/patch

cdmtaskservice/jaws/client.py#L49-L50

Added lines #L49 - L50 were not covered by tests

async def status(self, run_id: str) -> dict[str, Any]:
"""
Get the status of a JAWS run.
"""
res = await self._get(

Check warning on line 56 in cdmtaskservice/jaws/client.py

View check run for this annotation

Codecov / codecov/patch

cdmtaskservice/jaws/client.py#L56

Added line #L56 was not covered by tests
f"run/{_require_string(run_id, 'run_id')}",
params={"verbose": "true", "local_tz": "UTC"}
)
res["submitted"] = _add_tz(res["submitted"])
res["updated"] = _add_tz(res["updated"])
return res

Check warning on line 62 in cdmtaskservice/jaws/client.py

View check run for this annotation

Codecov / codecov/patch

cdmtaskservice/jaws/client.py#L60-L62

Added lines #L60 - L62 were not covered by tests

async def close(self):
await self._sess.close()

Check warning on line 65 in cdmtaskservice/jaws/client.py

View check run for this annotation

Codecov / codecov/patch

cdmtaskservice/jaws/client.py#L65

Added line #L65 was not covered by tests


def _add_tz(timestr: str) -> datetime.datetime:
return datetime.datetime.fromisoformat(timestr).replace(tzinfo=datetime.timezone.utc)

Check warning on line 69 in cdmtaskservice/jaws/client.py

View check run for this annotation

Codecov / codecov/patch

cdmtaskservice/jaws/client.py#L69

Added line #L69 was not covered by tests
1 change: 1 addition & 0 deletions docker-compose.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ services:
- KBCTS_SFAPI_CRED_PATH=/creds/sfapi_creds
- KBCTS_SFAPI_USER=cdm_ts
- KBCTS_NERSC_REMOTE_CODE_DIR=/global/cfs/cdirs/kbase/cdm_task_service
- KBCTS_JAWS_URL=https://jaws-api.jgi.doe.gov/api/v2
# Don't commit your token to github please
- KBCTS_JAWS_TOKEN=tokengoeshere
- KBCTS_JAWS_GROUP=kbase
Expand Down
7 changes: 7 additions & 0 deletions test/jaws/jaws_client_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
# TODO TEST add tests

from cdmtaskservice.jaws import client # @UnusedImport


def test_noop():
pass
Loading