如何基于FastAPI+PostgreSQL实现多级商品分类树接口？

问题描述 #

使用FastAPI和PostgreSQL开发商品分类树接口，目前仅能获取父分类及其一级子分类，无法处理子分类下的多级嵌套子分类，需要修复该问题。

已定义的Pydantic模型 #

class CategoryTree(BaseModel):
    id: int
    name: str
    children: list[CategoryTree] | None = None

class ListCategoryTree(RootModel):
    root: list[CategoryTree]

数据库分类表结构 #

现有问题代码 #

def rec(cat, all_categories):
    if cat.parent is None:
        all_categories.append(CategoryTree(
            id=cat.id,
            name=cat.name,
            children=[]
        ))
    else:
        for parent_cat in all_categories:
            if parent_cat.id == cat.parent:
                parent_cat.children.append(cat)

async def get_category_tree() -> ListCategoryTree:
    categories = await category_pg_service.get_categories()  # request to db
    all_categories = []
    for cat in categories:
        rec(cat, all_categories)
    return ListCategoryTree(root=all_categories)  # pydantic model (final result)

@router.get('', response_model=ListCategoryTree)
async def category_tree() -> ListCategoryTree:
    return await get_category_tree()

问题原因 #

原代码的递归函数仅在根节点列表all_categories中查找父节点，对于二级及以下的子分类，它们的父节点存在于根节点的children列表中，而非直接在all_categories里，因此无法匹配到父节点，导致多级子分类无法被正确添加到树结构中。

解决方案 #

通过建立分类ID到CategoryTree对象的映射字典，快速定位父节点，无论父节点处于树的哪一层，都能将当前分类添加到父节点的子列表中，最终收集所有根节点即可完成多级分类树的构建。

优化后的代码实现 #

async def get_category_tree() -> ListCategoryTree:
    categories = await category_pg_service.get_categories()
    
    # 构建ID到CategoryTree对象的映射，方便快速查找父节点
    category_map = {}
    for cat in categories:
        category_tree = CategoryTree(id=cat.id, name=cat.name, children=[])
        category_map[cat.id] = category_tree
    
    # 遍历所有分类，将当前分类添加到对应父节点的子列表中
    root_categories = []
    for cat in categories:
        current_node = category_map[cat.id]
        if cat.parent is None:
            # 父节点为空，作为根节点加入列表
            root_categories.append(current_node)
        else:
            # 找到父节点并添加当前节点到其子列表
            parent_node = category_map.get(cat.parent)
            if parent_node:
                parent_node.children.append(current_node)
    
    return ListCategoryTree(root=root_categories)

@router.get('', response_model=ListCategoryTree)
async def category_tree() -> ListCategoryTree:
    return await get_category_tree()

代码说明 #

1. 映射字典构建：先将所有数据库返回的分类转换为CategoryTree对象，并存入以分类ID为键的字典中，实现O(1)时间复杂度的父节点查找。
2. 树结构组装：遍历每个分类，若为根分类（parent为None）则直接加入根列表；否则通过字典找到父节点，将当前节点添加到父节点的children列表中，自动处理多级嵌套。
3. 返回结果：最后将根列表传入ListCategoryTree模型返回，即可得到完整的多级分类树。

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