feat: support starlette 1.0+

Author: waketzheng

.github/workflows/ci.yml

@@ -56,17 +56,21 @@ jobs:
         run: make test_sqlite_regexp
         env:
           PYTHONDEVMODE: 1
-      - name: Test FastAPI/Blacksheep/Sanic Examples
+      - name: Test FastAPI/Blacksheep/Sanic/Starlette Examples
         run: |
           PYTHONPATH=$DEST_FASTAPI uv run --frozen tortoise -c config.TORTOISE_ORM migrate
           PYTHONPATH=$DEST_FASTAPI uv run --frozen pytest $PYTEST_ARGS_SEQ $DEST_FASTAPI/_tests.py
           rm -f $DEST_FASTAPI/db.sqlite3
           PYTHONPATH=$DEST_BLACKSHEEP uv run --frozen pytest $PYTEST_ARGS $DEST_BLACKSHEEP/_tests.py
           PYTHONPATH=$DEST_SANIC uv run --frozen pytest $PYTEST_ARGS $DEST_SANIC/_tests.py
+          PYTHONPATH=$DEST_STARLETTE uv run --frozen pytest $PYTEST_ARGS $DEST_STARLETTE/_tests.py
+          uv pip install "starlette<1.0"
+          PYTHONPATH=$DEST_STARLETTE uv run --no-sync pytest $PYTEST_ARGS $DEST_STARLETTE/_tests.py
         env:
           DEST_FASTAPI: examples/fastapi
           DEST_BLACKSHEEP: examples/blacksheep
           DEST_SANIC: examples/sanic
+          DEST_STARLETTE: examples/starlette
           PYTHONDEVMODE: 1
           PYTEST_ARGS: "-n auto --cov=tortoise --cov-append --cov-branch --tb=native -q"
           PYTEST_ARGS_SEQ: "--cov=tortoise --cov-append --cov-branch --tb=native -q"

examples/starlette/_tests.py

@@ -0,0 +1,43 @@
+from pathlib import Path
+
+import pytest
+from asgi_lifespan import LifespanManager
+from httpx import ASGITransport, AsyncClient
+
+try:
+    from main import app
+    from models import Users
+except ImportError:
+    if (cwd := Path.cwd()) == (parent := Path(__file__).parent):
+        dirpath = "."
+    else:
+        dirpath = str(parent.relative_to(cwd))
+    print(f"You may need to explicitly declare python path:\n\nexport PYTHONPATH={dirpath}\n")
+    raise
+
+
+@pytest.fixture(scope="module")
+def anyio_backend() -> str:
+    return "asyncio"
+
+
+@pytest.mark.anyio
+async def test_app() -> None:
+    async with LifespanManager(app):
+        transport = ASGITransport(app=app)
+        # note: you _must_ set `base_url` for relative urls like "/" to work
+        async with AsyncClient(transport=transport, base_url="http://testserver") as client:
+            r = await client.get("/")
+            assert r.status_code == 200
+            assert r.json() == {"users": []}
+            assert await Users.all() == []
+
+            r = await client.post("/user/", json={"username": "Iron"})
+            assert r.status_code == 201
+            assert r.json() == {"user": "Users(id=1, username='Iron')"}
+            assert await Users.get(id=1) == await Users.last()
+
+            r = await client.get("/")
+            assert r.status_code == 200
+            assert r.json() == {"users": ["User 1: Iron"]}
+            assert await Users.all() == [await Users.first()]

examples/starlette/main.py

@@ -1,14 +1,16 @@
+#!/usr/bin/env python
 # pylint: disable=E0401,E0611
 import logging
 from json import JSONDecodeError
+from pathlib import Path
 
 from models import Users
 from starlette.applications import Starlette
 from starlette.exceptions import HTTPException
 from starlette.requests import Request
 from starlette.responses import JSONResponse
+from starlette.routing import Mount, Route
 from starlette.status import HTTP_201_CREATED, HTTP_400_BAD_REQUEST
-from uvicorn.main import run
 
 from tortoise.contrib.starlette import register_tortoise
 
@@ -17,29 +19,39 @@
 app = Starlette()
 
 
-@app.route("/", methods=["GET"])
 async def list_all(_: Request) -> JSONResponse:
     users = await Users.all()
     return JSONResponse({"users": [str(user) for user in users]})
 
 
-@app.route("/user", methods=["POST"])
 async def add_user(request: Request) -> JSONResponse:
     try:
         payload = await request.json()
         username = payload["username"]
     except JSONDecodeError:
-        raise HTTPException(status_code=HTTP_400_BAD_REQUEST, detail="cannot parse request body")
+        raise HTTPException(
+            status_code=HTTP_400_BAD_REQUEST, detail="cannot parse request body"
+        ) from None
     except KeyError:
-        raise HTTPException(status_code=HTTP_400_BAD_REQUEST, detail="username is required")
+        raise HTTPException(
+            status_code=HTTP_400_BAD_REQUEST, detail="username is required"
+        ) from None
 
     user = await Users.create(username=username)
-    return JSONResponse({"user": str(user)}, status_code=HTTP_201_CREATED)
+    return JSONResponse({"user": repr(user)}, status_code=HTTP_201_CREATED)
 
 
+app = Starlette(
+    routes=[
+        Route("/", list_all),
+        Mount("/user", routes=[Route("/", add_user, methods=["POST"])]),
+    ]
+)
 register_tortoise(
     app, db_url="sqlite://:memory:", modules={"models": ["models"]}, generate_schemas=True
 )
 
 if __name__ == "__main__":
-    run(app)
+    import uvicorn
+
+    uvicorn.run("__main__:app", reload=True, reload_dirs=[str(Path(__file__).parent)])

examples/starlette/models.py

@@ -7,3 +7,8 @@ class Users(models.Model):
 
     def __str__(self) -> str:
         return f"User {self.id}: {self.username}"
+
+    def __repr__(self) -> str:
+        fields = sorted(self._meta.db_fields)  # ['id', 'username']
+        values = ", ".join(f"{f}={getattr(self, f)!r}" for f in fields)
+        return f"{self.__class__.__name__}({values})"

tests/contrib/test_starlette.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+
+import pytest
+from asgi_lifespan import LifespanManager
+from httpx import ASGITransport, AsyncClient
+from starlette.applications import Starlette
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+from starlette.routing import Route
+
+from tortoise import fields, models
+from tortoise.context import TortoiseContext, get_current_context
+from tortoise.contrib.starlette import TortoiseContextMiddleware, register_tortoise
+
+
+class StarletteMiddlewareModel(models.Model):
+    id = fields.IntField(primary_key=True)
+
+
+def _register_tortoise(app: Starlette) -> None:
+    register_tortoise(
+        app,
+        db_url="sqlite://:memory:",
+        modules={"models": [__name__]},
+    )
+
+
+def _register_tortoise_with_lifespan(app: Starlette) -> None:
+    register_tortoise(
+        app,
+        db_url="sqlite://:memory:",
+        modules={"models": [__name__]},
+        generate_schemas=True,
+    )
+
+
+async def _context_endpoint(_: Request) -> JSONResponse:
+    ctx = get_current_context()
+    return JSONResponse({"inited": ctx.inited if ctx is not None else False})
+
+
+async def _get(client: AsyncClient) -> dict:
+    response = await client.get("/")
+    assert response.status_code == 200
+    return response.json()
+
+
+@pytest.mark.asyncio
+async def test_starlette_register_tortoise_uses_middleware_without_patching_endpoint() -> None:
+    route = Route("/", _context_endpoint)
+    app = Starlette(routes=[route])
+    original_endpoint = route.endpoint
+
+    _register_tortoise(app)
+
+    assert route.endpoint is original_endpoint
+    assert any(
+        middleware.cls is TortoiseContextMiddleware  # type:ignore[comparison-overlap]
+        for middleware in app.user_middleware
+    )
+
+    transport = ASGITransport(app=app)
+    async with AsyncClient(transport=transport, base_url="http://testserver") as client:
+        assert await _get(client) == {"inited": True}
+
+
+@pytest.mark.asyncio
+async def test_starlette_middleware_covers_routes_added_after_register_tortoise() -> None:
+    app = Starlette()
+    _register_tortoise(app)
+    app.add_route("/", _context_endpoint)
+
+    transport = ASGITransport(app=app)
+    async with AsyncClient(transport=transport, base_url="http://testserver") as client:
+        assert await _get(client) == {"inited": True}
+
+
+@pytest.mark.asyncio
+async def test_starlette_middleware_does_not_reuse_unrelated_global_context() -> None:
+    unrelated_ctx = TortoiseContext()
+    unrelated_ctx.__enter__()
+    await unrelated_ctx.init(
+        db_url="sqlite://:memory:",
+        modules={"models": [__name__]},
+        _enable_global_fallback=True,
+    )
+    unrelated_ctx.__exit__(None, None, None)
+
+    async def endpoint(_: Request) -> JSONResponse:
+        return JSONResponse({"unrelated": get_current_context() is unrelated_ctx})
+
+    try:
+        app = Starlette(routes=[Route("/", endpoint)])
+        _register_tortoise(app)
+
+        transport = ASGITransport(app=app)
+        async with AsyncClient(transport=transport, base_url="http://testserver") as client:
+            assert await _get(client) == {"unrelated": False}
+    finally:
+        await unrelated_ctx.close_connections()
+
+
+@pytest.mark.asyncio
+async def test_starlette_lifespan_context_is_stored_and_cleared() -> None:
+    @asynccontextmanager
+    async def lifespan(app: Starlette):
+        yield
+
+    app = Starlette(routes=[Route("/", _context_endpoint)], lifespan=lifespan)
+    _register_tortoise_with_lifespan(app)
+
+    async with LifespanManager(app):
+        assert hasattr(app.state, "_tortoise_context")
+
+        transport = ASGITransport(app=app)
+        async with AsyncClient(transport=transport, base_url="http://testserver") as client:
+            assert await _get(client) == {"inited": True}
+
+    assert not hasattr(app.state, "_tortoise_context")

tortoise/contrib/starlette/__init__.py

@@ -1,18 +1,53 @@
 from __future__ import annotations
 
-from collections.abc import Iterable
+from collections.abc import AsyncIterator, Iterable, Mapping
+from contextlib import asynccontextmanager
 from types import ModuleType
+from typing import Any, cast
 
+import starlette
 from starlette.applications import Starlette  # pylint: disable=E0401
+from starlette.routing import _DefaultLifespan as StarletteDefaultLifespan
+from starlette.types import ASGIApp, Lifespan, Receive, Scope, Send
 
 from tortoise import Tortoise
+from tortoise.config import TortoiseConfig
 from tortoise.connection import get_connections
+from tortoise.context import TortoiseContext, get_current_context
 from tortoise.log import logger
 
+_TORTOISE_CONTEXT_STATE = "_tortoise_context"
+
+
+class TortoiseContextMiddleware:
+    def __init__(self, app: ASGIApp, config: TortoiseConfig, starlette_app: Starlette) -> None:
+        self.app = app
+        self.config = config
+        self.starlette_app = starlette_app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] not in {"http", "websocket"}:
+            await self.app(scope, receive, send)
+            return
+
+        app_context = getattr(self.starlette_app.state, _TORTOISE_CONTEXT_STATE, None)
+        if app_context is None:
+            async with TortoiseContext() as request_context:
+                await request_context.init(self.config)
+                await self.app(scope, receive, send)
+            return
+
+        if get_current_context() is not app_context:
+            with app_context:
+                await self.app(scope, receive, send)
+            return
+
+        await self.app(scope, receive, send)
+
 
 def register_tortoise(
     app: Starlette,
-    config: dict | None = None,
+    config: dict[str, Any] | TortoiseConfig | None = None,
     config_file: str | None = None,
     db_url: str | None = None,
     modules: dict[str, Iterable[str | ModuleType]] | None = None,
@@ -79,16 +114,45 @@ def register_tortoise(
     ConfigurationError
         For any configuration error
     """
+    typed_config = TortoiseConfig.resolve_args(config, config_file, db_url, modules)
 
-    @app.on_event("startup")
     async def init_orm() -> None:  # pylint: disable=W0612
-        await Tortoise.init(config=config, config_file=config_file, db_url=db_url, modules=modules)
+        ctx = await Tortoise.init(config=typed_config, _enable_global_fallback=True)
+        setattr(app.state, _TORTOISE_CONTEXT_STATE, ctx)
         logger.info("Tortoise-ORM started, %s, %s", get_connections()._get_storage(), Tortoise.apps)
         if generate_schemas:
             logger.info("Tortoise-ORM generating schema")
             await Tortoise.generate_schemas()
 
-    @app.on_event("shutdown")
     async def close_orm() -> None:  # pylint: disable=W0612
         await Tortoise.close_connections()
+        if hasattr(app.state, _TORTOISE_CONTEXT_STATE):
+            delattr(app.state, _TORTOISE_CONTEXT_STATE)
         logger.info("Tortoise-ORM shutdown")
+
+    if starlette.__version__ < "1":
+        if (on_event := getattr(app, "on_event", None)) is not None:
+            on_event("startup")(init_orm)
+            on_event("shutdown")(close_orm)
+        return
+
+    original_lifespan = app.router.lifespan_context
+
+    if generate_schemas or not isinstance(original_lifespan, StarletteDefaultLifespan):
+
+        @asynccontextmanager
+        async def orm_inited_lifespan(app_: Starlette) -> AsyncIterator[Mapping[str, Any] | None]:
+            await init_orm()
+            try:
+                async with original_lifespan(app_) as maybe_state:
+                    yield maybe_state
+            finally:
+                await close_orm()
+
+        app.router.lifespan_context = cast("Lifespan[Any]", orm_inited_lifespan)
+
+    if not any(
+        middleware.cls is TortoiseContextMiddleware  # type:ignore[comparison-overlap]
+        for middleware in app.user_middleware
+    ):
+        app.add_middleware(TortoiseContextMiddleware, config=typed_config, starlette_app=app)