如何优化Python生成器流式输出SQL数据为CSV的实现？

我正尝试实现一个REST API，将SQL数据库中的记录以CSV格式返回。希望使用import csv完成转换，并借助generator实现结果流式输出。

我已尝试如下代码：

def get_data(query) -> Generator[str, None, None]:
    with pyodbc.connect(connStr) as conn:
        with conn.cursor() as cursor:
            cursor.execute(query)
            while 1:
                row = cursor.fetchone()
                if not row: break
                data = io.StringIO()
                csv.writer(data).writerow(row)
                yield data.getvalue()

该代码能运行，但过于冗长，且会创建大量临时StringIO和writer对象，求更优实现方案。

补充说明：我在FastAPI端点中调用get_data，代码如下：

@app.get("person/csv")
async def person_csv():
    rows = get_data("SELECT * FROM Person")
    response = StreamingResponse(rows, media_type="text/csv")
    response.headers["Content-Disposition"] = "attachment; filename=export.csv"
    return response

我认为这是返回任意数量记录且不受内存限制的唯一方式。

可以通过复用单个csv.writer和StringIO对象来减少资源消耗，同时简化代码逻辑。核心思路是在generator初始化时创建一次writer和缓冲区，之后每次写入行后获取内容并重置缓冲区，避免重复创建对象。

优化后的代码 #

import io
import csv
from typing import Generator
import pyodbc

def get_data(query) -> Generator[str, None, None]:
    with pyodbc.connect(connStr) as conn:
        with conn.cursor() as cursor:
            cursor.execute(query)
            # 获取表头（可选，根据需求添加）
            headers = [column[0] for column in cursor.description]
            # 初始化缓冲区和CSV writer，复用对象
            output = io.StringIO()
            writer = csv.writer(output)
            
            # 先写入表头（如果需要的话）
            writer.writerow(headers)
            yield output.getvalue()
            output.truncate(0)
            output.seek(0)
            
            # 批量获取行，提升效率（可选，比如每次fetchmany(1000)）
            while rows := cursor.fetchmany(1000):
                writer.writerows(rows)
                yield output.getvalue()
                output.truncate(0)
                output.seek(0)

关键优化点 #

• 复用对象：只创建一次StringIO和csv.writer，避免循环内重复初始化，减少内存开销和对象创建耗时。
• 批量读取：使用fetchmany()替代fetchone()，减少数据库交互次数，提升整体性能。
• 可选表头：通过cursor.description获取列名，自动生成CSV表头，符合CSV规范。
• 缓冲区重置：每次yield后通过truncate(0)和seek(0)清空并重置缓冲区，确保下一次写入正确。

FastAPI端点适配 #

你的FastAPI代码无需大改，直接使用优化后的get_data即可，注意补全路径前的斜杠以符合FastAPI规范：

@app.get("/person/csv")
async def person_csv():
    rows = get_data("SELECT * FROM Person")
    response = StreamingResponse(rows, media_type="text/csv")
    response.headers["Content-Disposition"] = "attachment; filename=export.csv"
    return response

额外说明 #

流式输出确实是处理大量数据时避免内存溢出的最优方案，你的思路是正确的。如果需要处理特殊字符或自定义CSV格式，可以在初始化csv.writer时传入delimiter、quotechar等参数进行配置。

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