feat: update Litestar framework integration with new request and response modules

Author: pythonhubdev

examples/with-litestar/with-thirdpartyemailpassword/main.py

@@ -12,7 +12,6 @@
     get_all_cors_headers,
     init,
 )
-from supertokens_python.framework.litestar.middleware import LitestarMiddleware
 from supertokens_python.recipe import (
     dashboard,
     emailverification,
@@ -179,9 +178,6 @@ def f_405(_, __: Exception):
     route_handlers=[
         get_session_info,
     ],
-    middleware=[
-        LitestarMiddleware(),
-    ],
     cors_config=cors,
     exception_handlers={
         Exception: f_405,

supertokens_python/framework/litestar/__init__.py

@@ -14,12 +14,13 @@
 from typing import Any
 
 from litestar import Request
+
 from supertokens_python.framework.types import Framework
 
 
 class LitestarFramework(Framework):
     def wrap_request(self, unwrapped: Request[Any, Any, Any]):
-        from supertokens_python.framework.litestar.request import (
+        from supertokens_python.framework.litestar.litestar_request import (
             LitestarRequest,
         )
 

supertokens_python/framework/litestar/framework.py

@@ -0,0 +1,124 @@
+# 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, Union, cast
+
+
+def get_middleware():
+    from litestar import Request as LitestarRequestClass
+    from litestar import Response as LitestarResponseClass
+    from starlette.responses import Response as StarletteResponse
+    from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+    from supertokens_python import Supertokens
+    from supertokens_python.exceptions import SuperTokensError
+    from supertokens_python.framework import BaseResponse
+    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 ASGIMiddleware:
+        def __init__(self, app: ASGIApp) -> None:
+            self.app = app
+
+        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 = LitestarRequestClass[Any, Any, Any](
+                cast(Any, scope), receive=cast(Any, receive)
+            )
+            custom_request = LitestarRequest(request)
+            user_context = default_user_context(custom_request)
+
+            try:
+                response = LitestarResponse(LitestarResponseClass[Any](content={}))
+                result: Union[BaseResponse, None] = await st.middleware(
+                    custom_request, response, user_context
+                )
+                if result is None:
+
+                    async def send_wrapper(message: Message):
+                        if message["type"] == "http.response.start":
+                            if hasattr(request.state, "supertokens") and isinstance(
+                                request.state.supertokens, SessionContainer
+                            ):
+                                starlette_response = StarletteResponse()
+                                starlette_response.raw_headers = message["headers"]
+                                response = LitestarResponse(
+                                    cast(LitestarResponseClass[Any], starlette_response)
+                                )
+                                manage_session_post_response(
+                                    request.state.supertokens, response, user_context
+                                )
+                                message["headers"] = starlette_response.raw_headers
+
+                        await send(message)
+
+                    await self.app(scope, receive, send_wrapper)
+                    return
+
+                if hasattr(request.state, "supertokens") and isinstance(
+                    request.state.supertokens, SessionContainer
+                ):
+                    manage_session_post_response(
+                        request.state.supertokens, result, user_context
+                    )
+
+                if isinstance(result, LitestarResponse):
+                    resp_any = cast(Any, result.response)
+                    asgi_response = resp_any.to_asgi_response(app=None, request=request)
+                    await asgi_response(scope, receive, send)
+                    return
+
+                return
+
+            except SuperTokensError as e:
+                response = LitestarResponse(LitestarResponseClass[Any](content={}))
+                result: Union[BaseResponse, None] = await st.handle_supertokens_error(
+                    LitestarRequest(request), e, response, user_context
+                )
+                if isinstance(result, LitestarResponse):
+                    resp_any = cast(Any, result.response)
+                    asgi_response = resp_any.to_asgi_response(app=None, request=request)
+                    await asgi_response(scope, receive, send)
+                    return
+
+                async def send_wrapper(message: Message):
+                    if message["type"] == "http.response.start":
+                        if hasattr(request.state, "supertokens") and isinstance(
+                            request.state.supertokens, SessionContainer
+                        ):
+                            starlette_response = StarletteResponse()
+                            starlette_response.raw_headers = message["headers"]
+                            response = LitestarResponse(
+                                cast(LitestarResponseClass[Any], starlette_response)
+                            )
+                            manage_session_post_response(
+                                request.state.supertokens, response, user_context
+                            )
+                            message["headers"] = starlette_response.raw_headers
+                    await send(message)
+
+                await self.app(scope, receive, send_wrapper)
+
+    return ASGIMiddleware

supertokens_python/framework/litestar/litestar_middleware.py

@@ -12,10 +12,11 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 import json
-from typing import Any, Union
+from typing import Any, Dict, Union
 from urllib.parse import parse_qsl
 
 from litestar import Request
+
 from supertokens_python.framework.request import BaseRequest
 from supertokens_python.recipe.session.interfaces import SessionContainer
 
@@ -33,10 +34,10 @@ def get_query_param(
     ) -> Union[str, None]:
         return self.request.query_params.get(key, default)
 
-    def get_query_params(self) -> dict[str, Any]:
+    def get_query_params(self) -> Dict[str, Any]:
         return dict(self.request.query_params.items())  # type: ignore
 
-    async def json(self) -> dict[str, Any]:
+    async def json(self) -> Dict[str, Any]:
         """
         Read the entire ASGI stream and JSON-decode it,
         sidestepping Litestar’s internal max-body-size logic.

…ens_python/framework/litestar/request.py …n/framework/litestar/litestar_request.pysupertokens_python/framework/litestar/request.py renamed to supertokens_python/framework/litestar/litestar_request.py supertokens_python/framework/litestar/request.py renamed to supertokens_python/framework/litestar/litestar_request.py

@@ -16,6 +16,7 @@
 
 from litestar import Response
 from litestar.serialization import encode_json
+
 from supertokens_python.framework.response import BaseResponse
 from supertokens_python.utils import get_timestamp_ms
 
@@ -60,7 +61,7 @@ def set_cookie(
         )
 
     def set_header(self, key: str, value: str):
-        self.response.set_header(key, value)
+        self.response.headers[key] = value
 
     def get_header(self, key: str) -> Optional[str]:
         return self.response.headers.get(key, None)

…ns_python/framework/litestar/response.py …/framework/litestar/litestar_response.pysupertokens_python/framework/litestar/response.py renamed to supertokens_python/framework/litestar/litestar_response.py supertokens_python/framework/litestar/response.py renamed to supertokens_python/framework/litestar/litestar_response.py

@@ -18,8 +18,8 @@
 
 from supertokens_python import Supertokens
 from supertokens_python.exceptions import SuperTokensError
-from supertokens_python.framework.litestar.request import LitestarRequest
-from supertokens_python.framework.litestar.response import LitestarResponse
+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, SessionRecipe
 from supertokens_python.recipe.session.interfaces import SessionClaimValidator
 from supertokens_python.types import MaybeAwaitable

supertokens_python/framework/litestar/middleware.py

@@ -21,7 +21,7 @@
 from pytest import fixture, mark
 from supertokens_python import InputAppInfo, SupertokensConfig, init
 from supertokens_python.framework import BaseRequest
-from supertokens_python.framework.litestar.middleware import LitestarMiddleware
+from supertokens_python.framework.litestar.litestar_middleware import get_middleware
 from supertokens_python.querier import Querier
 from supertokens_python.recipe import emailpassword, session, thirdparty
 from supertokens_python.recipe.dashboard import DashboardRecipe, InputOverrideConfig
@@ -169,12 +169,9 @@ async def _create_throw(request: Request[Any, Any, Any]) -> None:
             _create,
             _create_throw,
         ],
-        middleware=[
-            LitestarMiddleware(),
-        ],
     )
-
-    return TestClient(app)
+    return TestClient(get_middleware()(app))
+    # return TestClient(app)
 
 
 def apis_override_session(param: APIInterface):
@@ -575,11 +572,10 @@ def test_litestar_root_path(litestar_root_path: str):
         litestar_root_path = litestar_root_path[1:]
     app = Litestar(
         path=litestar_root_path,
-        middleware=[
-            LitestarMiddleware(),
-        ],
     )
-    test_client = TestClient(app)
+
+    test_client = TestClient(get_middleware()(app))
+    # test_client = TestClient(app)
 
     response = test_client.get(
         f"{litestar_root_path}/auth/signup/email/exists?email=test@example.com"