Handle Access Token Refresh (#11)

meierj-wayscript · web-flow · commit 3efa29bd233c · 2021-11-15T12:49:16.000-05:00
* Document how to use context library

* Refresh access token on a 401
diff --git a/README.md b/README.md
@@ -1,5 +1,13 @@
 # wayscript-python
 
+## Context
+
+```
+from wayscript import context
+
+event = context.get_event()
+```
+
 ## Triggers
 
 ### HTTP Triggers
diff --git a/src/wayscript/settings.py b/src/wayscript/settings.py
@@ -3,7 +3,7 @@
 """
 import os
 
-
+AUTH_ROUTE = "auth"
 PROCESSES_ROUTE = "processes"
 LAIRS_ROUTE = "lairs"
 WEBHOOKS = "webhooks"
@@ -13,6 +13,7 @@
 
 
 ROUTES = {
+    "auth": {"refresh": f"{AUTH_ROUTE}/refresh"},
     "lairs": {"detail": f"{LAIRS_ROUTE}/$id"},
     "processes": { "detail_expanded": f"{PROCESSES_ROUTE}/$id/detail"},
     "webhooks": {"http_trigger_response": f"{WEBHOOKS}/http-trigger/response/$id"},
diff --git a/src/wayscript/utils.py b/src/wayscript/utils.py
@@ -1,3 +1,4 @@
+import functools
 import os
 import string
 
@@ -11,13 +12,36 @@ def get_process_execution_user_token():
     token = os.environ.get("WAYSCRIPT_EXECUTION_USER_TOKEN")
     return token
 
+def set_process_execution_user_token(value: str):
+    os.environ["WAYSCRIPT_EXECUTION_USER_TOKEN"] = value
+    return value
 
 def get_process_id():
     """Return uuid of current container execution"""
     process_id = os.environ["WS_PROCESS_ID"]
     return process_id
 
 
+def get_refresh_token():
+    refresh_token = os.environ["WAYSCRIPT_EXECUTION_USER_REFRESH_TOKEN"]
+    return refresh_token
+
+
+def retry_on_401_wrapper(f):
+
+    @functools.wraps(f)
+    def _retry_on_401(client, *args, **kwargs):
+
+        response = f(client, *args, **kwargs)
+        # refresh credentials and retry on 401
+        if response.status_code == 401:
+            client._refresh_access_token()
+            response = f(client, *args, **kwargs)
+        return response
+
+    return _retry_on_401
+
+
 class WayScriptClient:
 
     def __init__(self, *args, **kwargs):
@@ -29,36 +53,53 @@ def __init__(self, *args, **kwargs):
     
     def _get_url(self, subpath: str, route: str, template_args: dict=None):
         """Generate an url"""
-        subpath_template = string.Template(settings.ROUTES[subpath][route])
-        subpath = subpath_template.substitute(**template_args)
+        subpath_template_str = settings.ROUTES[subpath][route]
+        subpath_template = string.Template(subpath_template_str)
+        full_subpath = subpath_template.substitute(**template_args) if template_args else subpath_template_str
 
-        url = f"{settings.WAYSCRIPT_ORIGIN}/{subpath}"
+        url = f"{settings.WAYSCRIPT_ORIGIN}/{full_subpath}"
         return url
 
+    def _refresh_access_token(self):
+        """Refresh access token and update environment"""
+        url = self._get_url(subpath="auth", route="refresh")
+        refresh_token = get_refresh_token()
+        payload = {"refresh_token": refresh_token}
+        response = self.session.post(url, json=payload)
+        response.raise_for_status()
+        access_token = response.json()["access_token"]
+        set_process_execution_user_token(access_token)
+        self.session.headers["authorization"] = f"Bearer {access_token}"
+
+    @retry_on_401_wrapper
     def get_process_detail_expanded(self, _id: str):
         """Request process expanded detail endpoint"""
         url = self._get_url(subpath="processes", route="detail_expanded", template_args={"id": _id})
         response = self.session.get(url)
         return response
 
+    @retry_on_401_wrapper
     def get_workspace_integration_detail(self, _id: str):
         """Request a workspace-integrations detail"""
         url = self._get_url(subpath="workspace-integrations", route="detail", template_args={"id": _id})
         response = self.session.get(url)
         return response
 
+    @retry_on_401_wrapper
     def get_lair_detail(self, _id: str):
         """Request lair detail"""
         url = self._get_url(subpath="lairs", route="detail", template_args={"id": _id})
         response = self.session.get(url)
         return response
 
+    @retry_on_401_wrapper
     def get_workspace_detail(self, _id: str):
         """Request workspace detail"""
         url = self._get_url(subpath="workspaces", route="detail", template_args={"id": _id})
         response = self.session.get(url)
         return response
 
+    @retry_on_401_wrapper
     def post_webhook_http_trigger_response(self, _id: str, payload: dict=None):
         """
         Post an http trigger response
