基于FastAPI+SQLModel构建GraphQL API的嵌套数据返回问题

在基于FastAPI+SQLModel构建的REST API中，GET /heroes/接口返回HeroReadWithTeam模型时，会自动触发关联的团队信息查询——这是因为SQLModel的Relationship默认懒加载，FastAPI在转换response_model时访问了team属性，从而自动发起了额外查询。但切换到Strawberry构建GraphQL API时，无法自动实现这种嵌套关联数据的返回，需要手动处理数据加载逻辑。

原REST API代码 #

from typing import List, Optional

from fastapi import Depends, FastAPI, HTTPException, Query
from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select


class TeamBase(SQLModel):
    name: str = Field(index=True)
    headquarters: str


class Team(TeamBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)

    heroes: List["Hero"] = Relationship(back_populates="team")


class TeamCreate(TeamBase):
    pass


class TeamRead(TeamBase):
    id: int

class HeroBase(SQLModel):
    name: str = Field(index=True)
    secret_name: str
    age: Optional[int] = Field(default=None, index=True)

    team_id: Optional[int] = Field(default=None, foreign_key="team.id")


class Hero(HeroBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)

    team: Optional[Team] = Relationship(back_populates="heroes")


class HeroRead(HeroBase):
    id: int


class HeroCreate(HeroBase):
    pass

class HeroReadWithTeam(HeroRead):
    team: Optional[TeamRead] = None


class TeamReadWithHeroes(TeamRead):
    heroes: List[HeroRead] = []


DATABASE_URL = 'postgresql+psycopg2://postgres:password@db:5432/blog_db'

engine = create_engine(DATABASE_URL, echo=True)


def create_db_and_tables():
    SQLModel.metadata.create_all(engine)


def get_session():
    with Session(engine) as session:
        yield session


app = FastAPI()


@app.on_event("startup")
def on_startup():
    create_db_and_tables()


@app.post("/heroes/", response_model=HeroRead)
def create_hero(*, session: Session = Depends(get_session), hero: HeroCreate):
    db_hero = Hero.from_orm(hero)
    session.add(db_hero)
    session.commit()
    session.refresh(db_hero)
    return db_hero


@app.get("/heroes/", response_model=List[HeroReadWithTeam])
def read_heroes(
    *,
    session: Session = Depends(get_session),
    offset: int = 0,
    limit: int = Query(default=100, lte=100),
):
    heroes = session.exec(select(Hero).offset(offset).limit(limit)).all()
    print(heroes)
    return heroes


@app.post("/teams/", response_model=TeamRead)
def create_team(*, session: Session = Depends(get_session), team: TeamCreate):
    db_team = Team.from_orm(team)
    session.add(db_team)
    session.commit()
    session.refresh(db_team)
    return db_team


@app.get("/teams/", response_model=List[TeamRead])
def read_teams(
    *,
    session: Session = Depends(get_session),
    offset: int = 0,
    limit: int = Query(default=100, lte=100),
):
    teams = session.exec(select(Team).offset(offset).limit(limit)).all()
    
    return teams

解决方案 #

核心思路 #

GraphQL不会像FastAPI response_model那样自动触发懒加载，需要主动预加载关联数据，并定义对应的GraphQL类型来支持嵌套返回。具体步骤如下：

1. 安装Strawberry依赖

   pip install strawberry-graphql[fastapi]
   

2. 定义GraphQL类型
   对应SQLModel的Hero和Team模型，创建Strawberry的类型，包含嵌套字段：

   import strawberry
   from typing import List, Optional
   
   @strawberry.type
   class TeamType:
       id: int
       name: str
       headquarters: str
       heroes: Optional[List["HeroType"]] = None
   
   @strawberry.type
   class HeroType:
       id: int
       name: str
       secret_name: str
       age: Optional[int]
       team: Optional[TeamType] = None
   
   # 解决前向引用问题
   TeamType.heroes = strawberry.field(default_factory=list)
   

3. 预加载关联数据
   在查询时使用SQLModel的selectinload（或joinedload）预加载关联的团队数据，避免N+1查询：

   from sqlmodel import select, selectinload
   
   def get_heroes_with_teams(session: Session, offset: int = 0, limit: int = 100):
       # 预加载Hero关联的Team数据
       heroes = session.exec(
           select(Hero)
           .offset(offset)
           .limit(limit)
           .options(selectinload(Hero.team))
       ).all()
       return heroes
   

4. 实现GraphQL查询解析器
   创建Strawberry的查询类型，并实现解析器，将SQLModel实例转换为GraphQL类型：

   @strawberry.type
   class Query:
       @strawberry.field
       def heroes(self, offset: int = 0, limit: int = 100, session: Session = Depends(get_session)) -> List[HeroType]:
           heroes = get_heroes_with_teams(session, offset, limit)
           # 将SQLModel实例转换为GraphQL类型
           return [
               HeroType(
                   id=hero.id,
                   name=hero.name,
                   secret_name=hero.secret_name,
                   age=hero.age,
                   team=TeamType(
                       id=hero.team.id,
                       name=hero.team.name,
                       headquarters=hero.team.headquarters
                   ) if hero.team else None
               )
               for hero in heroes
           ]
   

5. 挂载GraphQL路由到FastAPI
   将Strawberry的GraphQL接口挂载到FastAPI应用：

   from strawberry.fastapi import GraphQLRouter
   
   schema = strawberry.Schema(query=Query)
   graphql_app = GraphQLRouter(schema)
   app.include_router(graphql_app, prefix="/graphql")
   

完整整合代码 #

from typing import List, Optional

from fastapi import Depends, FastAPI, Query
from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select, selectinload
import strawberry
from strawberry.fastapi import GraphQLRouter


# SQLModel 模型定义
class TeamBase(SQLModel):
    name: str = Field(index=True)
    headquarters: str


class Team(TeamBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    heroes: List["Hero"] = Relationship(back_populates="team")


class TeamCreate(TeamBase):
    pass


class HeroBase(SQLModel):
    name: str = Field(index=True)
    secret_name: str
    age: Optional[int] = Field(default=None, index=True)
    team_id: Optional[int] = Field(default=None, foreign_key="team.id")


class Hero(HeroBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    team: Optional[Team] = Relationship(back_populates="heroes")


# GraphQL 类型定义
@strawberry.type
class TeamType:
    id: int
    name: str
    headquarters: str
    heroes: Optional[List["HeroType"]] = None


@strawberry.type
class HeroType:
    id: int
    name: str
    secret_name: str
    age: Optional[int]
    team: Optional[TeamType] = None


TeamType.heroes = strawberry.field(default_factory=list)


# 数据库配置
DATABASE_URL = 'postgresql+psycopg2://postgres:password@db:5432/blog_db'
engine = create_engine(DATABASE_URL, echo=True)


def create_db_and_tables():
    SQLModel.metadata.create_all(engine)


def get_session():
    with Session(engine) as session:
        yield session


# FastAPI 应用初始化
app = FastAPI()


@app.on_event("startup")
def on_startup():
    create_db_and_tables()


# REST API 接口（保留原接口）
@app.post("/heroes/")
def create_hero(*, session: Session = Depends(get_session), hero: HeroBase):
    db_hero = Hero.from_orm(hero)
    session.add(db_hero)
    session.commit()
    session.refresh(db_hero)
    return db_hero


@app.get("/heroes/")
def read_heroes(*, session: Session = Depends(get_session), offset: int = 0, limit: int = Query(default=100, lte=100)):
    heroes = session.exec(select(Hero).offset(offset).limit(limit)).all()
    return heroes


@app.post("/teams/")
def create_team(*, session: Session = Depends(get_session), team: TeamCreate):
    db_team = Team.from_orm(team)
    session.add(db_team)
    session.commit()
    session.refresh(db_team)
    return db_team


# GraphQL 查询解析器
def get_heroes_with_teams(session: Session, offset: int = 0, limit: int = 100):
    return session.exec(
        select(Hero)
        .offset(offset)
        .limit(limit)
        .options(selectinload(Hero.team))
    ).all()


@strawberry.type
class Query:
    @strawberry.field
    def heroes(self, offset: int = 0, limit: int = 100, session: Session = Depends(get_session)) -> List[HeroType]:
        heroes = get_heroes_with_teams(session, offset, limit)
        return [
            HeroType(
                id=hero.id,
                name=hero.name,
                secret_name=hero.secret_name,
                age=hero.age,
                team=TeamType(
                    id=hero.team.id,
                    name=hero.team.name,
                    headquarters=hero.team.headquarters
                ) if hero.team else None
            )
            for hero in heroes
        ]

    @strawberry.field
    def teams(self, offset: int = 0, limit: int = 100, session: Session = Depends(get_session)) -> List[TeamType]:
        teams = session.exec(
            select(Team)
            .offset(offset)
            .limit(limit)
            .options(selectinload(Team.heroes))
        ).all()
        return [
            TeamType(
                id=team.id,
                name=team.name,
                headquarters=team.headquarters,
                heroes=[
                    HeroType(
                        id=hero.id,
                        name=hero.name,
                        secret_name=hero.secret_name,
                        age=hero.age,
                        team=None
                    ) for hero in team.heroes
                ]
            )
            for team in teams
        ]


# 挂载GraphQL路由
schema = strawberry.Schema(query=Query)
graphql_app = GraphQLRouter(schema)
app.include_router(graphql_app, prefix="/graphql")

关键说明 #

• 预加载优化：使用selectinload可以一次性加载所有关联的团队数据，避免每个英雄都发起一次团队查询（N+1问题），如果需要更高效的关联查询，也可以使用joinedload实现JOIN查询。
• 类型转换：需要手动将SQLModel的实例转换为Strawberry的GraphQL类型，也可以使用Strawberry的自动转换工具（如strawberry.experimental.pydantic.convert_pydantic_model）简化转换逻辑。
• 循环引用处理：在TeamType和HeroType的嵌套中，可根据需求决定是否保留双向引用，GraphQL本身支持循环类型，但实际返回时需注意避免序列化问题。

内容的提问来源于stack exchange，提问作者gontxomde