Sync monorepo state at "[python sdk] Add support for caching access token across client instances" (#107)

Syncing from
userclouds/userclouds@d444b579997500a1ff5124d961d68cbab6279cc4

Author: uc-build-user

CHANGELOG.md

@@ -1,10 +1,11 @@
 # Changelog
 
-## UNPUBLISHED
+## 1.8.0 - UNPUBLISHED
 
 - Breaking change: Add "ending_before" argument to all paginated methods, add pagination to ExecuteAccessor, change "starting_after" and "ending_before" to be str instead of uuid
+- Add the ability to cache the access token globally (in process) and share it across clients instances. Pass the optional `use_global_cache_for_token=true` to the client in order to enable.
 
-## 1.7.0
+## 1.7.0 - 09-05-2024
 
 - Update userstore sample to exercise partial update columns
 - Add methods for creating, retrieving, updating, and deleting ColumnDataTypes
@@ -13,7 +14,7 @@
 - Add ColumnDataType resource IDs for native ColumnDataTypes
 - Update userstore_sample.py to interact with ColumnDataTypes
 
-## 1.6.1
+## 1.6.1 - Not published
 
 - Add SDK method for data import via ExecuteMutator
 

requirements-dev.txt

@@ -1,4 +1,4 @@
-black>=24.2.0
+black>=24.4.2
 ruff
 pyupgrade
 isort

src/usercloudssdk/asyncclient.py

@@ -2,14 +2,11 @@
 
 import asyncio
 import base64
-import time
 import urllib.parse
 import uuid
 from dataclasses import asdict
 from pathlib import Path
 
-import jwt
-
 from . import ucjson
 from .client_helpers import _SDK_VERSION, _id_from_identical_conflict, _read_env
 from .constants import _JSON_CONTENT_TYPE, AuthnType, Region
@@ -37,6 +34,7 @@
     UpdateColumnRetentionDurationsRequest,
     UserResponse,
 )
+from .token import cache_token, get_cached_token, is_token_expiring
 from .uchttpclient import create_default_uc_http_async_client
 
 
@@ -60,6 +58,7 @@ def __init__(
         client_secret: str,
         client_factory=create_default_uc_http_async_client,
         session_name: str | None = None,
+        use_global_cache_for_token: bool = False,
         **kwargs,
     ):
         self._authorization = base64.b64encode(
@@ -68,9 +67,9 @@ def __init__(
                 "ISO-8859-1",
             )
         ).decode("ascii")
-
         self._client = client_factory(base_url=url, **kwargs)
         self._access_token: str | None = None  # lazy loaded
+        self._use_global_cache_for_token = use_global_cache_for_token
         self._access_token_lock = asyncio.Lock()
         base_ua = f"UserClouds Python SDK v{_SDK_VERSION}"
         self._common_headers = {
@@ -1126,6 +1125,10 @@ async def CheckDataImportStatusAsync(self, import_id: uuid.UUID) -> dict:
     # Access Token Helpers
 
     async def _get_access_token_async(self) -> str:
+        if self._use_global_cache_for_token:
+            token = get_cached_token(self._authorization)
+            if token:
+                return token
         # Encode the client ID and client secret
         headers = {
             "Authorization": f"Basic {self._authorization}",
@@ -1142,7 +1145,10 @@ async def _get_access_token_async(self) -> str:
         if resp.status_code >= 400:
             raise UserCloudsSDKError.from_response(resp)
         json_data = ucjson.loads(resp.text)
-        return json_data.get("access_token")
+        token = json_data["access_token"]
+        if self._use_global_cache_for_token:
+            cache_token(self._authorization, token)
+        return token
 
     async def _refresh_access_token_if_needed_async(self) -> None:
         if self._access_token is None:
@@ -1151,23 +1157,9 @@ async def _refresh_access_token_if_needed_async(self) -> None:
                     self._access_token = await self._get_access_token_async()
                     return
 
-        # TODO: this takes advantage of an implementation detail that we use JWTs for
-        # access tokens, but we should probably either expose an endpoint to verify
-        # expiration time, or expect to retry requests with a well-formed error, or
-        # change our bearer token format in time.
-        if (
-            jwt.decode(self._access_token, options={"verify_signature": False}).get(
-                "exp"
-            )
-            < time.time()
-        ):
+        if is_token_expiring(self._access_token):
             async with self._access_token_lock:
-                if (
-                    jwt.decode(
-                        self._access_token, options={"verify_signature": False}
-                    ).get("exp")
-                    < time.time()
-                ):
+                if is_token_expiring(self._access_token):
                     self._access_token = await self._get_access_token_async()
 
     # Request Helpers

src/usercloudssdk/client.py

@@ -1,14 +1,11 @@
 from __future__ import annotations
 
 import base64
-import time
 import urllib.parse
 import uuid
 from dataclasses import asdict
 from pathlib import Path
 
-import jwt
-
 from . import ucjson
 from .client_helpers import _SDK_VERSION, _id_from_identical_conflict, _read_env
 from .constants import _JSON_CONTENT_TYPE, AuthnType, Region
@@ -36,6 +33,7 @@
     UpdateColumnRetentionDurationsRequest,
     UserResponse,
 )
+from .token import cache_token, get_cached_token, is_token_expiring
 from .uchttpclient import create_default_uc_http_client
 
 
@@ -59,6 +57,7 @@ def __init__(
         client_secret: str,
         client_factory=create_default_uc_http_client,
         session_name: str | None = None,
+        use_global_cache_for_token: bool = False,
         **kwargs,
     ):
         self._authorization = base64.b64encode(
@@ -70,6 +69,7 @@ def __init__(
 
         self._client = client_factory(base_url=url, **kwargs)
         self._access_token: str | None = None  # lazy loaded
+        self._use_global_cache_for_token = use_global_cache_for_token
         base_ua = f"UserClouds Python SDK v{_SDK_VERSION}"
         self._common_headers = {
             "User-Agent": f"{base_ua} [{session_name}]" if session_name else base_ua,
@@ -1082,6 +1082,10 @@ def CheckDataImportStatus(self, import_id: uuid.UUID) -> dict:
     # Access Token Helpers
 
     def _get_access_token(self) -> str:
+        if self._use_global_cache_for_token:
+            token = get_cached_token(self._authorization)
+            if token:
+                return token
         # Encode the client ID and client secret
         headers = {
             "Authorization": f"Basic {self._authorization}",
@@ -1096,23 +1100,17 @@ def _get_access_token(self) -> str:
         if resp.status_code >= 400:
             raise UserCloudsSDKError.from_response(resp)
         json_data = ucjson.loads(resp.text)
-        return json_data.get("access_token")
+        token = json_data.get("access_token")
+        if self._use_global_cache_for_token:
+            cache_token(self._authorization, token)
+        return token
 
     def _refresh_access_token_if_needed(self) -> None:
         if self._access_token is None:
             self._access_token = self._get_access_token()
             return
 
-        # TODO: this takes advantage of an implementation detail that we use JWTs for
-        # access tokens, but we should probably either expose an endpoint to verify
-        # expiration time, or expect to retry requests with a well-formed error, or
-        # change our bearer token format in time.
-        if (
-            jwt.decode(self._access_token, options={"verify_signature": False}).get(
-                "exp"
-            )
-            < time.time()
-        ):
+        if is_token_expiring(self._access_token):
             self._access_token = self._get_access_token()
 
     # Request Helpers

src/usercloudssdk/token.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import time
+
+import jwt
+
+_token_cache: dict[str, str] = {}
+
+
+def get_cached_token(authorization: str) -> str | None:
+    global _token_cache
+    return _token_cache.get(authorization, None)
+
+
+def cache_token(authorization: str, token: str) -> None:
+    if not token:
+        return
+    global _token_cache
+    _token_cache[authorization] = token
+
+
+def is_token_expiring(token: str | None) -> bool:
+    if not token:
+        return True
+    # TODO: this takes advantage of an implementation detail that we use JWTs for
+    # access tokens, but we should probably either expose an endpoint to verify
+    # expiration time, or expect to retry requests with a well-formed error, or
+    # change our bearer token format in time.
+    decoded_token = jwt.decode(token, options={"verify_signature": False})
+    expiration_time = decoded_token.get("exp")
+    if not expiration_time:
+        return True
+    return expiration_time > time.time()

tools/genasyncclient.py

@@ -8,14 +8,11 @@
 
 import asyncio
 import base64
-import time
 import urllib.parse
 import uuid
 from dataclasses import asdict
 from pathlib import Path
 
-import jwt
-
 from . import ucjson
 from .client_helpers import _SDK_VERSION, _id_from_identical_conflict, _read_env
 from .constants import _JSON_CONTENT_TYPE, AuthnType, Region
@@ -43,6 +40,7 @@
     UpdateColumnRetentionDurationsRequest,
     UserResponse,
 )
+from .token import cache_token, get_cached_token, is_token_expiring
 from .uchttpclient import create_default_uc_http_async_client
 
 
@@ -64,6 +62,7 @@ def __init__(
         client_secret: str,
         client_factory=create_default_uc_http_async_client,
         session_name: str | None = None,
+        use_global_cache_for_token: bool = False,
         **kwargs,
     ):
         self._authorization = base64.b64encode(
@@ -72,9 +71,9 @@ def __init__(
                 "ISO-8859-1",
             )
         ).decode("ascii")
-
         self._client = client_factory(base_url=url, **kwargs)
         self._access_token: str | None = None  # lazy loaded
+        self._use_global_cache_for_token = use_global_cache_for_token
         self._access_token_lock = asyncio.Lock()
         base_ua = f"UserClouds Python SDK v{_SDK_VERSION}"
         self._common_headers = {
@@ -85,6 +84,10 @@ def __init__(
 
 non_auto_generated_footer = """
     async def _get_access_token_async(self) -> str:
+        if self._use_global_cache_for_token:
+            token = get_cached_token(self._authorization)
+            if token:
+                return token
         # Encode the client ID and client secret
         headers = {
             "Authorization": f"Basic {self._authorization}",
@@ -101,7 +104,10 @@ async def _get_access_token_async(self) -> str:
         if resp.status_code >= 400:
             raise UserCloudsSDKError.from_response(resp)
         json_data = ucjson.loads(resp.text)
-        return json_data.get("access_token")
+        token = json_data["access_token"]
+        if self._use_global_cache_for_token:
+            cache_token(self._authorization, token)
+        return token
 
     async def _refresh_access_token_if_needed_async(self) -> None:
         if self._access_token is None:
@@ -110,23 +116,9 @@ async def _refresh_access_token_if_needed_async(self) -> None:
                     self._access_token = await self._get_access_token_async()
                     return
 
-        # TODO: this takes advantage of an implementation detail that we use JWTs for
-        # access tokens, but we should probably either expose an endpoint to verify
-        # expiration time, or expect to retry requests with a well-formed error, or
-        # change our bearer token format in time.
-        if (
-            jwt.decode(self._access_token, options={"verify_signature": False}).get(
-                "exp"
-            )
-            < time.time()
-        ):
+        if is_token_expiring(self._access_token):
             async with self._access_token_lock:
-                if (
-                    jwt.decode(
-                        self._access_token, options={"verify_signature": False}
-                    ).get("exp")
-                    < time.time()
-                ):
+                if is_token_expiring(self._access_token):
                     self._access_token = await self._get_access_token_async()
 
     # Request Helpers