feat: add Litestar framework support and related middleware

Author: pythonhubdev

dev-requirements.txt

@@ -21,4 +21,5 @@ pyyaml==6.0.2
 requests-mock==1.12.1
 respx>=0.13.0, <1.0.0
 uvicorn==0.32.0
+litestar==2.16.0
 -e .

supertokens_python/__init__.py

@@ -30,7 +30,7 @@
 
 def init(
     app_info: InputAppInfo,
-    framework: Literal["fastapi", "flask", "django"],
+    framework: Literal["fastapi", "flask", "django", "litestar"],
     supertokens_config: SupertokensConfig,
     recipe_list: List[Callable[[supertokens.AppInfo], RecipeModule]],
     mode: Optional[Literal["asgi", "wsgi"]] = None,

supertokens_python/framework/litestar/framework.py

@@ -0,0 +1,30 @@
+# Copyright (c) 2021, VRAI Labs and/or its affiliates. All rights reserved.
+#
+# This software is licensed under the Apache License, Version 2.0 (the
+# "License") as published by the Apache Software Foundation.
+#
+# You may not use this file except in compliance with the License. You may
+# obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from supertokens_python.framework.types import Framework
+
+if TYPE_CHECKING:
+    from litestar import Request
+
+
+class LitestarFramework(Framework):
+    def wrap_request(self, unwrapped: Request[Any, Any, Any]):
+        from supertokens_python.framework.litestar.litestar_request import (
+            LitestarRequest,
+        )
+
+        return LitestarRequest(unwrapped)

supertokens_python/framework/litestar/litestar_middleware.py

@@ -0,0 +1,98 @@
+from typing import Any
+
+from litestar import Request, Response
+from litestar.datastructures import MutableScopeHeaders
+from litestar.middleware.base import AbstractMiddleware
+from litestar.types import Message, Receive, Scope, Send
+from supertokens_python import Supertokens
+from supertokens_python.exceptions import SuperTokensError
+from supertokens_python.framework.litestar.litestar_request import LitestarRequest
+from supertokens_python.framework.litestar.litestar_response import LitestarResponse
+from supertokens_python.recipe.session import SessionContainer
+from supertokens_python.supertokens import manage_session_post_response
+from supertokens_python.utils import default_user_context
+
+
+class LitestarMiddleware(AbstractMiddleware):
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        st = Supertokens.get_instance()
+        request = Request[Any, Any, Any](scope, receive=receive, send=send)
+        custom_request = LitestarRequest(request)
+        user_context = default_user_context(custom_request)
+
+        try:
+            response = LitestarResponse(Response[Any](content=None))
+            result = await st.middleware(custom_request, response, user_context)
+            if result is not None:
+                # SuperTokens handled the request
+                if hasattr(request.state, "supertokens") and isinstance(
+                    request.state.supertokens, SessionContainer
+                ):
+                    manage_session_post_response(
+                        request.state.supertokens, result, user_context
+                    )
+                # Add cookies using MutableScopeHeaders
+                asgi_response = await result.response.to_asgi_response(
+                    app=None, request=request
+                )
+
+                async def modified_send(message: Message):
+                    if message["type"] == "http.response.start":
+                        mutable_headers = MutableScopeHeaders(message)
+                        for cookie in result.response.cookies:
+                            cookie_value = cookie.to_header().split(": ", 1)[1]
+                            mutable_headers.add("set-cookie", cookie_value)
+                    await send(message)
+
+                await asgi_response(scope, receive, modified_send)
+                return
+            else:
+                # SuperTokens didn’t handle the request; wrap the send function
+                async def send_wrapper(message: Message):
+                    if message["type"] == "http.response.start":
+                        if hasattr(request.state, "supertokens") and isinstance(
+                            request.state.supertokens, SessionContainer
+                        ):
+                            temp_response = Response[Any](content=None)
+                            temp_response.headers = MutableScopeHeaders(message)
+                            litestar_response = LitestarResponse(temp_response)
+                            manage_session_post_response(
+                                request.state.supertokens,
+                                litestar_response,
+                                user_context,
+                            )
+                            mutable_headers = MutableScopeHeaders(message)
+                            for cookie in litestar_response.response.cookies:
+                                cookie_value = cookie.to_header().split(": ", 1)[1]
+                                mutable_headers.add("set-cookie", cookie_value)
+                    await send(message)
+
+                await self.app(scope, receive, send_wrapper)
+                return
+
+        except SuperTokensError as e:
+            response = LitestarResponse(Response[Any](content=None))
+            result = await st.handle_supertokens_error(
+                custom_request, e, response, user_context
+            )
+            if isinstance(result, LitestarResponse):
+                # Add cookies using MutableScopeHeaders
+                asgi_response = await result.response.to_asgi_response(
+                    app=None, request=request
+                )
+
+                async def modified_send(message: Message):
+                    if message["type"] == "http.response.start":
+                        mutable_headers = MutableScopeHeaders(message)
+                        for cookie in result.response.cookies:
+                            cookie_value = cookie.to_header().split(": ", 1)[1]
+                            mutable_headers.add("set-cookie", cookie_value)
+                    await send(message)
+
+                await asgi_response(scope, receive, modified_send)
+                return
+            raise Exception("Should never come here")

supertokens_python/framework/litestar/litestar_request.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2021, VRAI Labs and/or its affiliates. All rights reserved.
+#
+# This software is licensed under the Apache License, Version 2.0 (the
+# "License") as published by the Apache Software Foundation.
+#
+# You may not use this file except in compliance with the License. You may
+# obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+from typing import Any, Dict, Union
+from urllib.parse import parse_qsl
+
+from litestar import Request
+from litestar.exceptions import SerializationException
+from supertokens_python.framework.request import BaseRequest
+from supertokens_python.recipe.session.interfaces import SessionContainer
+
+
+class LitestarRequest(BaseRequest):
+    def __init__(self, request: Request[Any, Any, Any]):
+        super().__init__()
+        self.request = request
+
+    def get_original_url(self) -> str:
+        print(f"The request url is {self.request.url}")
+        print(f"The request url is {self.request.url.from_components()}")
+        return str(self.request.url)
+
+    def get_query_param(
+        self, key: str, default: Union[str, None] = None
+    ) -> Union[str, None]:
+        return self.request.query_params.get(key, default)
+
+    def get_query_params(self) -> Dict[str, Any]:
+        return dict(self.request.query_params.items())  # type: ignore
+
+    async def json(self) -> Union[Any, None]:
+        try:
+            return await self.request.json()
+        except SerializationException:
+            return {}
+
+    def method(self) -> str:
+        return self.request.method
+
+    def get_cookie(self, key: str) -> Union[str, None]:
+        return self.request.cookies.get(key)
+
+    def get_header(self, key: str) -> Union[str, None]:
+        return self.request.headers.get(key, None)
+
+    def get_session(self) -> Union[SessionContainer, None]:
+        return self.request.state.supertokens
+
+    def set_session(self, session: SessionContainer):
+        self.request.state.supertokens = session
+
+    def set_session_as_none(self):
+        self.request.state.supertokens = None
+
+    def get_path(self) -> str:
+        print(f"The request url is {self.request.url}")
+        return self.request.url.path
+
+    async def form_data(self):
+        return dict(parse_qsl((await self.request.body()).decode("utf-8")))

supertokens_python/framework/litestar/litestar_response.py

@@ -0,0 +1,94 @@
+# Copyright (c) 2021, VRAI Labs and/or its affiliates. All rights reserved.
+#
+# This software is licensed under the Apache License, Version 2.0 (the
+# "License") as published by the Apache Software Foundation.
+#
+# You may not use this file except in compliance with the License. You may
+# obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+from math import ceil
+from typing import Any, Dict, Literal, Optional
+
+from supertokens_python.framework.response import BaseResponse
+from supertokens_python.utils import get_timestamp_ms
+
+
+class LitestarResponse(BaseResponse):
+    from litestar import Response
+
+    def __init__(self, response: Response[Any]):
+        super().__init__({})
+        self.response = response
+        self.original = response
+        self.parser_checked = False
+        self.response_sent = False
+        self.status_set = False
+
+    def set_html_content(self, content: str):
+        if not self.response_sent:
+            body = bytes(content, "utf-8")
+            self.set_header("Content-Length", str(len(body)))
+            self.set_header("Content-Type", "text/html")
+            self.response.content = body
+            self.response_sent = True
+
+    def set_cookie(
+        self,
+        key: str,
+        value: str,
+        expires: int,
+        path: str = "/",
+        domain: Optional[str] = None,
+        secure: bool = False,
+        httponly: bool = False,
+        samesite: Literal["lax", "strict", "none"] = "lax",
+    ):
+        self.response.set_cookie(
+            key=key,
+            value=value,
+            expires=ceil((expires - get_timestamp_ms()) / 1000),
+            path=path,
+            domain=domain,
+            secure=secure,
+            httponly=httponly,
+            samesite=samesite,
+        )
+
+    def set_header(self, key: str, value: str):
+        self.response.set_header(key, value)
+
+    def get_header(self, key: str) -> Optional[str]:
+        return self.response.headers.get(key, None)
+
+    def remove_header(self, key: str):
+        del self.response.headers[key]
+
+    def set_status_code(self, status_code: int):
+        if not self.status_set:
+            self.response.status_code = status_code
+            self.status_code = status_code
+            self.status_set = True
+
+    def set_json_content(self, content: Dict[str, Any]):
+        if not self.response_sent:
+            from litestar.serialization import encode_json
+
+            body = encode_json(
+                content,
+            )
+            self.set_header("Content-Type", "application/json; charset=utf-8")
+            self.set_header("Content-Length", str(len(body)))
+            self.response.content = body
+            self.response_sent = True
+
+    def redirect(self, url: str) -> BaseResponse:
+        if not self.response_sent:
+            self.set_header("Location", url)
+            self.set_status_code(302)
+            self.response_sent = True
+        return self

supertokens_python/framework/types.py

@@ -18,13 +18,14 @@
 
 from supertokens_python.framework.request import BaseRequest
 
-frameworks = ["fastapi", "flask", "django"]
+frameworks = ["fastapi", "flask", "django", "litestar"]
 
 
 class FrameworkEnum(Enum):
     FASTAPI = 1
     FLASK = 2
     DJANGO = 3
+    LITESTAR = 4
 
 
 class Framework(ABC):