A Simple Way To Manage Sessions With AWS Lambda & DynamoDB In Python

Motivation & code examples for how to use DynamoDB as a session store for AWS Lambda function — in a python environment.

Tl;dr

I created the session-lambda python package for managing sessions on AWS lambda with DynamoDB — https://github.com/roy-pstr/session-lambda

Motivation

AWS Lambdas are stateless functions running in a managed cloud environment. The key here is stateless. It means that two subsequent lambda calls have no dependency. This is by design, and due to the fact there is no shared memory between lambda calls (at least not on purpose).

With that said, there are use cases where you would want those lambdas to have a state. For example when running a backend server with user sessions or for applications that need to have a state — like a chat app where you want to access the recent chat history. Let's see how this can be done.

Why DynamoDB?

For a lambda to have a shared state between subsequent calls it needs to have a shared memory of some kind. Then those lambdas will need a way to access this memory and pull the data they want — the session data. But, with restrictions. The lambda should gain access only to the session data of the client that invoked the call. The best way to do this is with key/value storage — DynamoDB!

I chose DynamoDB over Reddis because DynamoDB is a fully serverless service, while Reddis has no native serverless service in the AWS ecosystem, yet.

Session-Lambda Python Package

session-lambda is a simple lambda handler wrapper in python that takes care of the session management part! You provide the table name in DynamoDB and the wrapper handles everything else for you.

from lambda_session import init_session_store, session, get_session_data

@session(key_name='session-id', update=False, ttl=15*60)
def lambda_handler(event, context):
  session_data = get_session_data()
  ...

Its key features are:

• Session id generation (both server-side and client-side).
• Get the session data from the table so you can access it from within the handler function.
• TTL feature (time to live) which is important if you don’t want your session data to live forever in the DynamoDB table.
• Both mutable and immutable session data modes.

Code Example

Clone example repository

git clone git@github.com:roy-pstr/session-lambda-example.git

cd session-lambda-example

Install python requirements

python3 -m venv .venv
source ./.venv/bin/activate
pip install -r requirements.txt

Create a local store

Store refers to a key/value database where we keep the sessions data

For simplicity, we will work with a local JSON file as the sessions’ store.

echo '{}' >> local_store.json

A brief look at the code files

1. lambda_function.py — This is where the handler implementation is. Pay attention to the @session, get_session_data() and set_sessions_data() . The handler parses the user message out of the event body and saves it to the session data.

• @session is a python decorator that pulls the session data (if a session-id is provided and the session exists) from the store and injects it into the handler function.
• get_session_data() is our way to fetch the session data into a local variable inside the handler.
• set_session_data() is our way to put the session data in the store. If the session does not exist in the store it will create it, and if it already exists it can update its’ value, depending on the update flag value.

2. invoke.py — This is a simple script for local invoking of the lambda handler. It gets two arguments: message [session-id], the second is optional. Then the message is passed as the body value, and the session-id is passed inside the headers.

Invoke without session-id

python invoke.py "hello"
python invoke.py "hello"

• Call the script twice
• The store content should be (with different session ids)

{
    "zpXmm-bkGQCwKwKGmwltJYBbumx_RZfHZrmfCegD1FQ": null,
    "xtNGpXFFPl8tgHr9CN8dWlWrF81FbeATx_f2CApIijs": null
}

• Two sessions were created in the store, but there is no session data. Only null values.

Invoke lambda with session-id

• Let's try again to invoke the lambda, but this time with a session-id (this session-id is passed as a header to the lambda handler request)

python invoke.py "xtNGpXFFPl8tgHr9CN8dWlWrF81FbeATx_f2CApIijs" "hello"

• And now you should see it in the local store

{
    "zpXmm-bkGQCwKwKGmwltJYBbumx_RZfHZrmfCegD1FQ": null,
    "xtNGpXFFPl8tgHr9CN8dWlWrF81FbeATx_f2CApIijs": "hello"
}

• Nice! now you can access this session data in any lambda call using the same session-id.
• But what if you want to update this session data? For that, you can use the update flag. By default this behavior is disabled and calling set_session_data() does not override existing value.

Invoke the lambda and update session data

• Enable update flag

@session(update=True)
def lambda_handler(event, context):
  ...

• Invoke the lambda with new data

python invoke.py "-xtNGpXFFPl8tgHr9CN8dWlWrF81FbeATx_f2CApIijs" "hello world"

• Now check the store, you should see:

{
    "zpXmm-bkGQCwKwKGmwltJYBbumx_RZfHZrmfCegD1FQ": null,
    "xtNGpXFFPl8tgHr9CN8dWlWrF81FbeATx_f2CApIijs": "hello world"
}

Server-side vs client-side generated session ids

• In the above example, the server was in charge of generating the initial session-id value and passing it to the client. Then the client uses it to access the session data.
• There is another use case, where the client wants to generate its’ own session-id. Let's take a look at how it can be done.
• The client generated the next key: client-session-id
• Now, use it the same way as before:

python invoke.py "client-session-id" "hello"

• You should see in the store

{
    "zpXmm-bkGQCwKwKGmwltJYBbumx_RZfHZrmfCegD1FQ": null,
    "xtNGpXFFPl8tgHr9CN8dWlWrF81FbeATx_f2CApIijs": "hello world",
    "client-session-id": "hello"
}

• The difference here is that when passing a session-id that does not exist in store, we assume it is on purpose and because the client wants to create the id himself. In this case, a new session will be created in the store, with the corresponding id value.

Use DynamoDB table as the session store

• OK. so for the real thing you better use DynamoDB. To do so, just remove the following line from the code:

init_session_store(file_path='local_store.json')

• Now invoke the function again

python invoke.py "-dqMD2t1AiiHf95K58q5bg1qectK0XQOeFe57cJEOlc" "hello world"

• You should get an error message with:

session_lambda.session.SessionStoreNotSet: Store not set

• This is because we did not point the session manager to any store.
• To point to a DynamoDB table store you use SESSION_LAMBDA_DYNAMODB_TABLE_NAME environment variable or use the init_session_store(dynamodb_table_name=TABLE_NAME) like this.
• Your DynamoDB table should have a key attribute named key of type string . session-id will be stored under this attribute, and the session data will be stored under value attribute.
• That’s it! Now your sessions will be saved to your DynamoDB table.

Time to live feature

It is recommended to use the TTL feature for your table. Enable this from your AWS console and set the attribute name to ttl then you can use it with the session decorator like this: @session(ttl=60*15) . ttl units are seconds.

Conclusion

session-lambda can help you manage sessions for lambda functions without the overhead of creating the behavior between the function and the session store (DynamoDB).

The session-lambda package was just released and any feedback and suggestions for improve will be super welcome! Enjoy.

References

GitHub - roy-pstr/session-lambda: A simple way to manage sessions for AWS Lambdas