PsycologyAPI/main.py
2024-11-29 17:01:19 +08:00

73 lines
1.9 KiB
Python

from random import randint
from fastapi import FastAPI
import os
from contextlib import asynccontextmanager
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,
generate_schemas=True,
add_exception_handlers=True,
_create_db=True,
):
# 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}