FastAPI 示例¶

from tortoise import fields, models
from tortoise.contrib.pydantic import pydantic_model_creator


class Users(models.Model):
    """
    The User model
    """

    id = fields.IntField(primary_key=True)
    #: This is a username
    username = fields.CharField(max_length=20, unique=True)
    name = fields.CharField(max_length=50, null=True)
    family_name = fields.CharField(max_length=50, null=True)
    category = fields.CharField(max_length=30, default="misc")
    password_hash = fields.CharField(max_length=128, null=True)
    created_at = fields.DatetimeField(auto_now_add=True)
    modified_at = fields.DatetimeField(auto_now=True)

    def full_name(self) -> str:
        """
        Returns the best name
        """
        if self.name or self.family_name:
            return f"{self.name or ''} {self.family_name or ''}".strip()
        return self.username

    class PydanticMeta:
        computed = ["full_name"]
        exclude = ["password_hash"]


User_Pydantic = pydantic_model_creator(Users, name="User")
UserIn_Pydantic = pydantic_model_creator(Users, name="UserIn", exclude_readonly=True)

# mypy: no-disallow-untyped-decorators
# pylint: disable=E0611,E0401
import datetime
from typing import AsyncGenerator

import pytest
import pytz
from asgi_lifespan import LifespanManager
from httpx import ASGITransport, AsyncClient
from main import app, app_east
from models import Users


@pytest.fixture(scope="module")
def anyio_backend() -> str:
    return "asyncio"


@pytest.fixture(scope="module")
async def client() -> AsyncGenerator[AsyncClient, None]:
    async with LifespanManager(app):
        transport = ASGITransport(app=app)
        async with AsyncClient(transport=transport, base_url="http://test") as c:
            yield c


@pytest.mark.anyio
async def test_create_user(client: AsyncClient) -> None:  # nosec
    response = await client.post("/users", json={"username": "admin"})
    assert response.status_code == 200, response.text
    data = response.json()
    assert data["username"] == "admin"
    assert "id" in data
    user_id = data["id"]

    user_obj = await Users.get(id=user_id)
    assert user_obj.id == user_id


@pytest.fixture(scope="module")
async def client_east() -> AsyncGenerator[AsyncClient, None]:
    async with LifespanManager(app_east):
        transport = ASGITransport(app=app_east)
        async with AsyncClient(transport=transport, base_url="http://test") as c:
            yield c


@pytest.mark.anyio
async def test_create_user_east(client_east: AsyncClient) -> None:  # nosec
    response = await client_east.post("/users_east", json={"username": "admin"})
    assert response.status_code == 200, response.text
    data = response.json()
    assert data["username"] == "admin"
    assert "id" in data
    user_id = data["id"]

    user_obj = await Users.get(id=user_id)
    assert user_obj.id == user_id

    # Verify that the time zone is East 8.
    created_at = user_obj.created_at

    # Asia/Shanghai timezone
    asia_tz = pytz.timezone("Asia/Shanghai")
    asia_now = datetime.datetime.now(pytz.utc).astimezone(asia_tz)
    assert created_at.hour - asia_now.hour == 0

    # UTC timezone
    utc_tz = pytz.timezone("UTC")
    utc_now = datetime.datetime.now(pytz.utc).astimezone(utc_tz)
    assert created_at.hour - utc_now.hour == 8

# pylint: disable=E0611,E0401
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, AsyncGenerator, List

from fastapi import FastAPI, HTTPException
from models import Users
from pydantic import BaseModel

from tortoise.contrib.fastapi import RegisterTortoise
from tortoise.contrib.pydantic import PydanticModel

if TYPE_CHECKING:  # pragma: nocoverage

    class UserIn_Pydantic(Users, PydanticModel):  # type:ignore[misc]
        pass

    class User_Pydantic(Users, PydanticModel):  # type:ignore[misc]
        pass

else:
    from models import User_Pydantic, UserIn_Pydantic


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
    # app startup
    async with RegisterTortoise(
        app,
        db_url="sqlite://:memory:",
        modules={"models": ["models"]},
        generate_schemas=True,
        add_exception_handlers=True,
    ):
        # db connected
        yield
        # app teardown
    # db connections closed


@asynccontextmanager
async def lifespan_east(app: FastAPI) -> AsyncGenerator[None, None]:
    # app startup
    async with RegisterTortoise(
        app,
        db_url="sqlite://:memory:",
        modules={"models": ["models"]},
        generate_schemas=True,
        add_exception_handlers=True,
        use_tz=False,
        timezone="Asia/Shanghai",
    ):
        # db connected
        yield
        # app teardown
    # db connections closed


app = FastAPI(title="Tortoise ORM FastAPI example", lifespan=lifespan)
app_east = FastAPI(title="Tortoise ORM FastAPI example", lifespan=lifespan_east)


class Status(BaseModel):
    message: str


@app.get("/users", response_model=List[User_Pydantic])
async def get_users():
    return await User_Pydantic.from_queryset(Users.all())


@app.post("/users", response_model=User_Pydantic)
async def create_user(user: UserIn_Pydantic):
    user_obj = await Users.create(**user.model_dump(exclude_unset=True))
    return await User_Pydantic.from_tortoise_orm(user_obj)


@app.get("/user/{user_id}", response_model=User_Pydantic)
async def get_user(user_id: int):
    return await User_Pydantic.from_queryset_single(Users.get(id=user_id))


@app.put("/user/{user_id}", response_model=User_Pydantic)
async def update_user(user_id: int, user: UserIn_Pydantic):
    await Users.filter(id=user_id).update(**user.model_dump(exclude_unset=True))
    return await User_Pydantic.from_queryset_single(Users.get(id=user_id))


@app.delete("/user/{user_id}", response_model=Status)
async def delete_user(user_id: int):
    deleted_count = await Users.filter(id=user_id).delete()
    if not deleted_count:
        raise HTTPException(status_code=404, detail=f"User {user_id} not found")
    return Status(message=f"Deleted user {user_id}")


############################ East 8 ############################
@app_east.get("/users_east", response_model=List[User_Pydantic])
async def get_users_east():
    return await User_Pydantic.from_queryset(Users.all())


@app_east.post("/users_east", response_model=User_Pydantic)
async def create_user_east(user: UserIn_Pydantic):
    user_obj = await Users.create(**user.model_dump(exclude_unset=True))
    return await User_Pydantic.from_tortoise_orm(user_obj)


@app_east.get("/user_east/{user_id}", response_model=User_Pydantic)
async def get_user_east(user_id: int):
    return await User_Pydantic.from_queryset_single(Users.get(id=user_id))


@app_east.put("/user_east/{user_id}", response_model=User_Pydantic)
async def update_user_east(user_id: int, user: UserIn_Pydantic):
    await Users.filter(id=user_id).update(**user.model_dump(exclude_unset=True))
    return await User_Pydantic.from_queryset_single(Users.get(id=user_id))


@app_east.delete("/user_east/{user_id}", response_model=Status)
async def delete_user_east(user_id: int):
    deleted_count = await Users.filter(id=user_id).delete()
    if not deleted_count:
        raise HTTPException(status_code=404, detail=f"User {user_id} not found")
    return Status(message=f"Deleted user {user_id}")