From affa4fb8bd606e3b19c0c2c6f5c3d9bca2100396 Mon Sep 17 00:00:00 2001
From: Gavin <gaprice@lbl.gov>
Date: Thu, 26 Dec 2024 13:00:17 -0800
Subject: [PATCH] Add minimal JAWS client

Will url to config later.
---
 cdmtaskservice/jaws/client.py | 69 +++++++++++++++++++++++++++++++++++
 docker-compose.yaml           |  1 +
 test/jaws/jaws_client_test.py |  7 ++++
 3 files changed, 77 insertions(+)
 create mode 100644 cdmtaskservice/jaws/client.py
 create mode 100644 test/jaws/jaws_client_test.py

diff --git a/cdmtaskservice/jaws/client.py b/cdmtaskservice/jaws/client.py
new file mode 100644
index 0000000..75702f8
--- /dev/null
+++ b/cdmtaskservice/jaws/client.py
@@ -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
+    
+    def __init__(self, url: str, token: str):
+        url = _require_string(url, "url")
+        if not url.endswith("/"):
+            url += "/"
+        self._sess = aiohttp.ClientSession(
+            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:
+            # 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()
+
+    async def _user(self) -> str:
+        res = await self._get("user")
+        return res["uid"]
+    
+    async def status(self, run_id: str) -> dict[str, Any]:
+        """
+        Get the status of a JAWS run.
+        """
+        res = await self._get(
+            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
+    
+    async def close(self):
+        await self._sess.close()
+
+
+def _add_tz(timestr: str) -> datetime.datetime:
+    return datetime.datetime.fromisoformat(timestr).replace(tzinfo=datetime.timezone.utc)
diff --git a/docker-compose.yaml b/docker-compose.yaml
index 6bf748d..2f3785a 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -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
diff --git a/test/jaws/jaws_client_test.py b/test/jaws/jaws_client_test.py
new file mode 100644
index 0000000..bff5774
--- /dev/null
+++ b/test/jaws/jaws_client_test.py
@@ -0,0 +1,7 @@
+# TODO TEST add tests
+
+from cdmtaskservice.jaws import client  # @UnusedImport
+
+
+def test_noop():
+    pass