Merge pull request #9 from DACCS-Climate/add-jupyterlab-functions

add jupyterlab functions

Author: mishaschwartz

README.md

@@ -105,4 +105,49 @@ Weaver URL is https://pavics.ouranos.ca/weaver/
 >>> # A MarbleService object is returned that can be used wherever a string can be used
 >>> print(f"Weaver URL is {MarbleClient()['PAVICS'].weaver}")
 Weaver URL is https://pavics.ouranos.ca/weaver/
+```
+
+## Jupyterlab functionality
+
+When running in a Marble Jupyterlab environment, the client can take advantage of various environment variables and 
+Jupyter's API to provide some additional functionality. 
+
+> [!WARNING]
+> Calling any of the methods described below outside a Marble Jupyterlab environment will raise a 
+> `JupyterEnvironmentError`.
+
+Get the node your notebook/script is currently running on:
+
+```python
+>>> client = MarbleClient()
+>>> client.this_node
+<marble_client.node.MarbleNode at 0x10c129990>
+```
+
+Add session cookies to a `requests.Session` object. This means that any request made with that session variable will
+be made as if you were logged in to the current Marble node. This is the recommended way to access protected resources
+programmatically in your scripts:
+
+```python
+>>> client = MarbleClient()
+>>> session = requests.Session
+>>> client.this_session(session)
+>>> session.cookies.get_dict()
+{...} # session cookiejar now includes cookies
+```
+
+You can also use the `this_session` method to create a new `requests.Sesssion` object:
+
+```python
+>>> client = MarbleClient()
+>>> session = client.this_session()
+>>> session.cookies.get_dict()
+{...} # now includes session cookies
+```
+
+You can now make a request to a protected resource on the current node using this session object. You will be able to 
+access the resource if you have permission:
+
+```python
+>>> session.get(f"{client.this_node.url}/some/protected/subpath")
 ```

marble_client/client.py

@@ -4,18 +4,36 @@
 import os
 import shutil
 import warnings
+from functools import wraps, cache
 from typing import Optional
+from urllib.parse import urlparse
 
 import dateutil.parser
 import requests
 
 from marble_client.constants import CACHE_FNAME, CACHE_META_FNAME, NODE_REGISTRY_URL
-from marble_client.exceptions import UnknownNodeError
+from marble_client.exceptions import UnknownNodeError, JupyterEnvironmentError
 from marble_client.node import MarbleNode
 
 __all__ = ["MarbleClient"]
 
 
+def check_jupyterlab(f):
+    """
+    Wraps the function f by first checking if the current script is running in a
+    Marble Jupyterlab environment and raising a JupyterEnvironmentError if not.
+
+    This is used as a pre-check for functions that only work in a Marble Jupyterlab
+    environment.
+    """
+    @wraps(f)
+    def wrapper(*args, **kwargs):
+        if os.getenv("PAVICS_HOST_URL"):
+            return f(*args, **kwargs)
+        raise JupyterEnvironmentError("Not in a Marble jupyterlab environment")
+    return wrapper
+
+
 class MarbleClient:
     def __init__(self, fallback: Optional[bool] = True) -> None:
         """Constructor method
@@ -51,6 +69,38 @@ def __init__(self, fallback: Optional[bool] = True) -> None:
     def nodes(self) -> dict[str, MarbleNode]:
         return self._nodes
 
+    @property
+    @cache
+    @check_jupyterlab
+    def this_node(self) -> MarbleNode:
+        """
+        Return the node where this script is currently running.
+
+        Note that this function only works in a Marble Jupyterlab environment.
+        """
+        host_url = urlparse(os.getenv("PAVICS_HOST_URL"))
+        for node in self.nodes.values():
+            if urlparse(node.url).hostname == host_url.hostname:
+                return node
+        raise UnknownNodeError(f"No node found in the registry with the url {host_url}")
+
+    @check_jupyterlab
+    def this_session(self, session: Optional[requests.Session]) -> requests.Session:
+        """
+        Add the login session cookies of the user who is currently logged in to the session object.
+        If a session object is not passed as an argument to this function, create a new session
+        object as well.
+
+        Note that this function only works in a Marble Jupyterlab environment.
+        """
+        if session is None:
+            session = requests.Session()
+        r = requests.get(f"{os.getenv('JUPYTERHUB_API_URL')}/users/{os.getenv('JUPYTERHUB_USER')}",
+                         headers={"Authorization": f"token {os.getenv('JUPYTERHUB_API_TOKEN')}"})
+        for name, value in r.json().get("auth_state", {}).get("magpie_cookies", {}).items():
+            session.cookies.set(name, value)
+        return session
+
     def __getitem__(self, node: str) -> MarbleNode:
         try:
             return self.nodes[node]

marble_client/exceptions.py

@@ -8,3 +8,7 @@ class ServiceNotAvailableError(MarbleBaseError):
 
 class UnknownNodeError(MarbleBaseError):
     pass
+
+
+class JupyterEnvironmentError(MarbleBaseError):
+    pass