66 lines
1.7 KiB
Python
66 lines
1.7 KiB
Python
from contextlib import asynccontextmanager
|
|
from random import randint
|
|
from typing import AsyncGenerator
|
|
|
|
from fastapi import FastAPI
|
|
from tortoise import Tortoise, generate_config
|
|
from tortoise.contrib.fastapi import RegisterTortoise
|
|
|
|
from models import RawData
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan_test(app: FastAPI) -> AsyncGenerator[None, None]:
|
|
config = generate_config(
|
|
"sqlite://:memory:",
|
|
app_modules={"models": ["models"]},
|
|
testing=True,
|
|
connection_label="models",
|
|
)
|
|
async with RegisterTortoise(
|
|
app=app,
|
|
config=config,
|
|
generate_schemas=True,
|
|
add_exception_handlers=True,
|
|
_create_db=True,
|
|
):
|
|
# db connected
|
|
yield
|
|
# app teardown
|
|
# db connections closed
|
|
await Tortoise._drop_databases()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
|
|
if getattr(app.state, "testing", None):
|
|
async with lifespan_test(app) as _:
|
|
yield
|
|
else:
|
|
config = generate_config(
|
|
"sqlite:///db.sqlite3",
|
|
app_modules={"models": ["models"]},
|
|
connection_label="models",
|
|
)
|
|
async with RegisterTortoise(
|
|
app=app,
|
|
config=config,
|
|
):
|
|
# db connected
|
|
yield
|
|
# app teardown
|
|
# db connections closed
|
|
await Tortoise.close_connections()
|
|
|
|
|
|
app = FastAPI(lifespan=lifespan)
|
|
|
|
|
|
@app.get("/")
|
|
async def root():
|
|
cnt = await RawData.all().count()
|
|
# 生成一个范围内的随机整数，随机返回一篇文章
|
|
r = randint(1, cnt)
|
|
record = await RawData.get(id=r)
|
|
return {"title": record.title, "content": record.content}
|