Token Storage Adapters#
DynamoDB Token Storage#
The following example demonstrates a token storage adapter which uses AWS DynamoDB as the backend storage mechanism.
Unlike the builtin adapters for JSON and sqlite, there is no capability here for an enumeration of all of the tokens in storage. This is because DynamoDB functions as a key-value store, and can efficiently map keys, but features slow sequential scans for enumeration. The example therefore demonstrates that key-value stores with limited or no capabilities for table scans can be used to implement the token storage interface.
from __future__ import annotations
import argparse
import time
import typing as t
import boto3
import globus_sdk
from globus_sdk.tokenstorage import StorageAdapter
CLIENT_ID = "61338d24-54d5-408f-a10d-66c06b59f6d2"
tablename = "example-globus-tokenstorage"
parser = argparse.ArgumentParser()
parser.add_argument(
"--create",
action="store_true",
)
boto_client = boto3.client("dynamodb")
auth_client = globus_sdk.NativeAppAuthClient(CLIENT_ID)
class DynamoDBStorageAdapter(StorageAdapter):
def __init__(self, client, tablename: str, namespace: str = "DEFAULT") -> None:
"""
:param client: A boto3 DyanmoDB client to use
:param tablename: The name of the dynamodb table to use
:param namespace: A namespace for all keys within this storage.
Setting up explicit namespacing allows for multiple storage
adapters for multiple users or applications to share a table.
"""
self.client = client
self.tablename = tablename
self.namespace = namespace
def _compute_key(self, resource_server: str) -> str:
"""
Compute the 'token_data_id' used for storage and retrieval
vis-a-vis a specific resource_server/namespace combination.
This is defined as a simple delimited string which starts with the
namespace given.
Globus keys tokens by the ``resource_server`` string, but also has
additional context about which user and application were being
used. For the storage adapter, we will need to use namespacing to
separate users.
Consider setting ``namespace`` to a value like a user ID or a
combination of user ID and authentication context.
"""
return f"{self.namespace}:{resource_server}"
def store(self, token_response: globus_sdk.OAuthTokenResponse) -> None:
for resource_server, token_data in token_response.by_resource_server.items():
key = self._compute_key(resource_server)
dynamo_item = {
"token_data_id": {"S": key},
"resource_server": {"S": resource_server},
"access_token": {"S": token_data["access_token"]},
"refresh_token": {"S": token_data["refresh_token"]},
"expires_at_seconds": {"N": str(token_data["expires_at_seconds"])},
"scope": {"S": token_data["scope"]},
}
# avoid setting `refresh_token` if it is null (meaning the
# login flow used access tokens only)
if token_data["refresh_token"] is None:
del dynamo_item["refresh_token"]
self.client.put_item(TableName=self.tablename, Item=dynamo_item)
def get_token_data(self, resource_server: str) -> dict[str, t.Any] | None:
key = self._compute_key(resource_server)
wrapped_item = self.client.get_item(
TableName=self.tablename,
Key={"token_data_id": {"S": key}},
ConsistentRead=True,
)
if "Item" not in wrapped_item:
return None
dynamo_item = wrapped_item["Item"]
return {
"resource_server": dynamo_item["resource_server"]["S"],
"access_token": dynamo_item["access_token"]["S"],
"refresh_token": dynamo_item.get("refresh_token", {"S": None})["S"],
"expires_at_seconds": int(dynamo_item["expires_at_seconds"]["N"]),
"scope": dynamo_item["scope"]["S"],
}
def create_table():
# create a table with a key of "token_data_id"
# this is a nonspecific string key which we will compute
#
# the relationship of "token_data_id" to the token will be explained below
boto_client.create_table(
TableName=tablename,
KeySchema=[{"AttributeName": "token_data_id", "KeyType": "HASH"}],
AttributeDefinitions=[{"AttributeName": "token_data_id", "AttributeType": "S"}],
BillingMode="PROVISIONED",
ProvisionedThroughput={"ReadCapacityUnits": 5, "WriteCapacityUnits": 5},
)
print(f"'{tablename}' create in progress.")
# poll status until the table is "ACTIVE"
print("Waiting for active status (Ctrl+C to cancel)...", end="", flush=True)
status = None
while status != "ACTIVE":
time.sleep(1)
try:
r = boto_client.describe_table(TableName=tablename)
except boto_client.exceptions.ResourceNotFoundException:
continue
print(".", end="", flush=True)
status = r["Table"]["TableStatus"]
print("ok")
def do_login_flow(storage: StorageAdapter):
auth_client.oauth2_start_flow(
requested_scopes=globus_sdk.GroupsClient.scopes.view_my_groups_and_memberships,
refresh_tokens=True,
)
authorize_url = auth_client.oauth2_get_authorize_url()
print(f"Please go to this URL and login:\n\n{authorize_url}\n")
auth_code = input("Please enter the code here: ").strip()
tokens = auth_client.oauth2_exchange_code_for_tokens(auth_code)
storage.store(tokens)
return tokens.by_resource_server[globus_sdk.GroupsClient.resource_server]
def group_list(storage: StorageAdapter):
tokens = storage.get_token_data(globus_sdk.GroupsClient.resource_server)
if tokens is None:
tokens = do_login_flow(storage)
authorizer = globus_sdk.RefreshTokenAuthorizer(
tokens["refresh_token"],
auth_client,
access_token=tokens["access_token"],
expires_at=tokens["expires_at_seconds"],
on_refresh=storage.store,
)
groups_client = globus_sdk.GroupsClient(authorizer=authorizer)
print("ID,Name,Type,Session Enforcement,Roles")
for group in groups_client.get_my_groups():
# parse the group to get data for output
if group.get("enforce_session"):
session_enforcement = "strict"
else:
session_enforcement = "not strict"
roles = ",".join({m["role"] for m in group["my_memberships"]})
print(
",".join(
[
group["id"],
group["name"],
group["group_type"],
session_enforcement,
roles,
]
)
)
if __name__ == "__main__":
args = parser.parse_args()
if args.create:
create_table()
else:
storage = DynamoDBStorageAdapter(boto_client, tablename)
group_list(storage)