1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135
|
from typing import Any, Dict, Literal, Optional
from uuid import UUID, uuid4
from pydantic import BaseModel, EmailStr, SecretStr
from litestar import Litestar, Request, get, post
from litestar.connection import ASGIConnection
from litestar.exceptions import NotAuthorizedException
from litestar.middleware.session.server_side import ServerSideSessionBackend, ServerSideSessionConfig
from litestar.openapi.config import OpenAPIConfig
from litestar.security.session_auth import SessionAuth
from litestar.stores.memory import MemoryStore
# Let's assume we have a User model that is a pydantic model.
# This though is not required - we need some sort of user class -
# but it can be any arbitrary value, e.g. an SQLAlchemy model,
# a representation of a MongoDB etc.
class User(BaseModel):
id: UUID
name: str
email: EmailStr
# we also have pydantic types for two different
# kinds of POST request bodies: one for creating
# a user, e.g. "sign-up", and the other for logging
# an existing user in.
class UserCreatePayload(BaseModel):
name: str
email: EmailStr
password: SecretStr
class UserLoginPayload(BaseModel):
email: EmailStr
password: SecretStr
MOCK_DB: Dict[str, User] = {}
memory_store = MemoryStore()
# The SessionAuth class requires a handler callable
# that takes the session dictionary, and returns the
# 'User' instance correlating to it.
#
# The session dictionary itself is a value the user decides
# upon. So for example, it might be a simple dictionary
# that holds a user id, for example: { "id": "abcd123" }
#
# Note: The callable can be either sync or async - both will work.
async def retrieve_user_handler(
session: Dict[str, Any], connection: "ASGIConnection[Any, Any, Any, Any]"
) -> Optional[User]:
return MOCK_DB.get(user_id) if (user_id := session.get("user_id")) else None
@post("/login")
async def login(data: UserLoginPayload, request: "Request[Any, Any, Any]") -> User:
# we received log-in data via post.
# our login handler should retrieve from persistence (a db etc.)
# the user data and verify that the login details
# are correct. If we are using passwords, we should check that
# the password hashes match etc. We will simply assume that we
# have done all of that we now have a user value:
user_id = await memory_store.get(data.email)
if not user_id:
raise NotAuthorizedException
user_id = user_id.decode("utf-8")
# once verified we can create a session.
# to do this we simply need to call the Litestar
# 'Request.set_session' method, which accepts either dictionaries
# or pydantic models. In our case, we can simply record a
# simple dictionary with the user ID value:
request.set_session({"user_id": user_id})
# you can do whatever we want here. In this case, we will simply return the user data:
return MOCK_DB[user_id]
@post("/signup")
async def signup(data: UserCreatePayload, request: Request[Any, Any, Any]) -> User:
# this is similar to the login handler, except here we should
# insert into persistence - after doing whatever extra
# validation we might require. We will assume that this is done,
# and we now have a user instance with an assigned ID value:
user = User(name=data.name, email=data.email, id=uuid4())
await memory_store.set(data.email, str(user.id))
MOCK_DB[str(user.id)] = user
# we are creating a session the same as we do in the
# 'login_handler' above:
request.set_session({"user_id": str(user.id)})
# and again, you can add whatever logic is required here, we
# will simply return the user:
return user
# the endpoint below requires the user to be already authenticated
# to be able to access it.
@get("/user", sync_to_thread=False)
def get_user(request: Request[User, Dict[Literal["user_id"], str], Any]) -> Any:
# because this route requires authentication, we can access
# `request.user`, which is the authenticated user returned
# by the 'retrieve_user_handler' function we passed to SessionAuth.
return request.user
# We add the session security schema to the OpenAPI config.
openapi_config = OpenAPIConfig(
title="My API",
version="1.0.0",
)
session_auth = SessionAuth[User, ServerSideSessionBackend](
retrieve_user_handler=retrieve_user_handler,
# we must pass a config for a session backend.
# all session backends are supported
session_backend_config=ServerSideSessionConfig(),
# exclude any URLs that should not have authentication.
# We exclude the documentation URLs, signup and login.
exclude=["/login", "/signup", "/schema"],
)
# We initialize the app instance, passing to it the 'session_auth.on_app_init' and the 'openapi_config'.
app = Litestar(
route_handlers=[login, signup, get_user],
on_app_init=[session_auth.on_app_init],
openapi_config=openapi_config,
)
|