Merge branch 'cern_dlaas_notebook' into 'master'

SDC-307 Cern dlaas notebook

See merge request astron-sdc/esap-userprofile-python-client!8

rucio_cli/__init__.py

+from rucio_cli.rucio_connector import rucio_connector

rucio_cli/rucio_connector.py

+import requests
+import json
+import io
+import getpass
+import pandas as pd
+
+from typing import Union, Optional
+
+class rucio_connector:
+
+    name = "rucio"
+    archive = "rucio"
+
+    def basket_item_to_pandas(
+            self, basket_item: Union[dict, pd.Series], validate: bool = True
+    ) -> Optional[pd.Series]:
+        """Convert an item from the shopping basket into a `pd.Series` with
+        optional validation.
+
+        Parameters
+        ----------
+        basket_item : Union[dict, pd.Series]
+            A single item from a retrieved shopping basket - either a raw `dict`
+            or a converted `pd.Series`.
+        validate : bool
+            If `True`, check that the data in the shopping item conforms with
+            the expected format before attempting the conversion.
+
+        Returns
+        -------
+        Optional[pd.Series]
+            `pd.Series` containing the data encoded in the shopping item or
+            `NoneType`.
+
+        """
+        if validate:
+            item_data = self.validate_basket_item(basket_item, return_loaded=True)
+        else:
+            item_data = json.loads(basket_item["item_data"])
+        if item_data:
+            return pd.Series(item_data)
+        return None
+
+
+    def validate_basket_item(
+            self, basket_item: Union[dict, pd.Series], return_loaded: bool = False
+    ) -> Union[dict, bool, None]:
+        """Check that the data in the shopping item conforms with
+        the expected format
+
+        Parameters
+        ----------
+        basket_item : Union[dict, pd.Series]
+            A single item from a retrieved shopping basket - either a raw `dict`
+            or a converted `pd.Series`.
+        return_loaded : bool
+            If `True`, and validation succeeds return the extracted shopping item
+            as `dict`, otherwise return `True` if validation succeeds and `None`
+            otherwise.
+
+        Returns
+        -------
+        Union[dict, bool, None]
+            If `return_loaded` is `True`, return a `dict` containing the data
+            encoded in the shopping item when validation succeeds.
+            Otherwise if `return_loaded` is `True` validation succeeds.
+            If validation fails return `None`.
+
+        """
+        item_data = json.loads(basket_item["item_data"])
+        if "archive" in item_data and item_data["archive"] == self.archive:
+            if return_loaded:
+                return item_data
+            else:
+                return True
+        return None

setup.py

@@ -8,7 +8,7 @@ setuptools.setup(
     version="0.0.4",
     author="Hugh Dickinson",
     author_email="hugh.dickinson@open.ac.uk",
-    description="Python client for ESAP Data Discovery Shoipping Basket",
+    description="Python client for ESAP Data Discovery Shopping Basket",
     long_description=long_description,
     long_description_content_type="text/markdown",
     url="https://git.astron.nl/astron-sdc/esap-userprofile-python-client",

shopping_client/shopping_client.py

-import requests
+import base64
+import getpass
 import json
+import logging
+import time
 import urllib.parse
-import getpass
+from os import getenv
+from typing import Optional, Union
+from warnings import warn
+
 import pandas as pd
+import requests
 
-from typing import Union, Optional
-from warnings import warn
+logger = logging.getLogger(__name__)
 
 
 class shopping_client:
 
     endpoint = "esap-api/accounts/user-profiles/"
+    audience = "rucio"  # Audience used by ESAP, might be configurable later
 
     def __init__(
         self,
-        token: str = None,
+        token: Optional[str] = None,
         host: str = "http://localhost:5555/",
         connectors: list = [],
     ):
@@ -38,7 +45,10 @@ class shopping_client:
         self.basket = None
 
     def get_basket(
-        self, convert_to_pandas: bool = False, reload: bool = False, filter_archives: bool = False
+        self,
+        convert_to_pandas: bool = False,
+        reload: bool = False,
+        filter_archives: bool = False,
     ) -> Union[list, pd.DataFrame, None]:
         """Retrieve the shopping basket for a user.
         Prompts for access token if one was not supplied to constructor.
@@ -87,8 +97,22 @@ class shopping_client:
 
         return self.basket
 
+    def _is_valid_token(self, token: Optional[str]) -> bool:
+        """Checks expiry of the token"""
+
+        if token is None:
+            return False
+
+        try:
+            data = token.split(".")[1]
+            padded = data + "=" * divmod(len(data), 4)[1]
+            payload = json.loads(base64.urlsafe_b64decode(padded))
+            return payload["exp"] > int(time.time()) + 10
+        except KeyError:
+            raise RuntimeError("Invalid JWT format")
+
     def _request_header(self):
-        while self.token is None:
+        while not self._is_valid_token(self.token):
             self._get_token()
 
         return dict(Accept="application/json", Authorization=f"Bearer {self.token}")
@@ -102,12 +126,14 @@ class shopping_client:
                 item_data = json.loads(item["item_data"])
 
                 for connector in self.connectors:
-                    if "archive" in item_data and item_data["archive"] == connector.archive:
+                    if (
+                        "archive" in item_data
+                        and item_data["archive"] == connector.archive
+                    ):
                         filtered_items.append(item)
 
         return filtered_items
 
-
     def _basket_to_pandas(self):
         if len(self.connectors):
             converted_basket = {
@@ -131,4 +157,32 @@ class shopping_client:
         return self.basket
 
     def _get_token(self):
+
+        # Generic JH token method using authstate
+        jh_api_uri = getenv("JUPYTERHUB_API_URL")
+        jh_api_token = getenv("JUPYTERHUB_API_TOKEN")
+        # Fallback to older rucio file
+        token_fn = getenv("RUCIO_OIDC_FILE_NAME")
+
+        try:
+            if all((jh_api_token, jh_api_uri)):
+                res = requests.get(
+                    f"{jh_api_uri}/user",
+                    headers={"Authorization": f"token {jh_api_token}"},
+                )
+
+                self.token = res.json()["auth_state"]["exchanged_tokens"][self.audience]
+
+        except KeyError:
+            logger.warning("JupyterHub without Authstate enabled")
+
+        # Try to get token from Rucio OIDC file (when running in CERN DLaaS notebook)
+        if self.token is None and token_fn is not None:
+            with open(token_fn) as token_file:
+                self.token = token_file.readline()
+
+        elif self.token is None:
             self.token = getpass.getpass("Enter your ESAP access token:")
+
+        if self.token is None:
+            raise RuntimeError("No token found!")