如何在Flask中流式传输大CSV文件，避免全量加载至内存？

问题背景 #

现有Flask导出CSV的端点在处理500万条数据时崩溃，原代码将所有数据先写入内存中的StringIO再一次性返回，导致内存占用过高。原代码如下：

import io, csv
from flask import Flask, Response

app = Flask(__name__)

def get_report_rows():
    # 模拟百万级数据库行
    for i in range(5_000_000):
        yield (i, f"name-{i}", i % 100)

@app.get("/export")
def export():
    rows = get_report_rows()

    output = io.StringIO()                   
    writer = csv.writer(output)
    writer.writerow(["id", "name", "score"])
    writer.writerows(rows)
    output.seek(0)

    return Response(
        output.getvalue(),
        mimetype="text/csv",
        headers={"Content-Disposition": "attachment; filename=report.csv"}
    )

流式导出实现 #

利用Flask的Response支持生成器作为响应体的特性，实现逐行输出CSV内容，无需一次性加载全部数据到内存：

import csv
from io import StringIO
from flask import Flask, Response

app = Flask(__name__)

def get_report_rows():
    # 实际场景替换为数据库查询迭代器（如SQLAlchemy的yield_per）
    for i in range(5_000_000):
        yield (i, f"name-{i}", i % 100)

def csv_stream():
    # 初始化CSV写入器，用StringIO作为临时缓冲
    buffer = StringIO()
    writer = csv.writer(buffer)
    
    # 输出表头
    writer.writerow(["id", "name", "score"])
    yield buffer.getvalue()
    buffer.seek(0)
    buffer.truncate()
    
    # 逐行输出数据
    for row in get_report_rows():
        writer.writerow(row)
        yield buffer.getvalue()
        buffer.seek(0)
        buffer.truncate()

@app.get("/export")
def export():
    return Response(
        csv_stream(),
        mimetype="text/csv",
        headers={
            "Content-Disposition": "attachment; filename=report.csv",
            "Transfer-Encoding": "chunked"
        }
    )

核心优化点 #

• 分块传输：生成器逐次产出内容，Flask自动启用分块编码，避免一次性占用大量内存。
• CSV格式安全：保留csv.writer处理字段转义（如含逗号、引号的字段），避免手动拼接导致格式错误。
• 内存高效：每次仅在内存中存储一行CSV内容，处理百万级数据时内存占用极低。

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