FastAPI & pdfkit 实现 PDF 文件下载: 3种高效方案

2024-12-15 11:54:23

使用 FastAPI 和 pdfkit 实现 PDF 文件下载

当使用 FastAPI 和 pdfkit 构建 HTML 转 PDF 的 API 时，一个常见的问题是如何让用户下载生成的 PDF 文件，而不是将其保存在服务器本地。本文将探讨这个问题，并提供几种解决方案。

问题分析

最初的代码片段展示了一个基础的 FastAPI 应用，它接收一个 URL，使用 pdfkit 将其转换为 PDF 文件，并将其保存到服务器的文件系统中。这种方法在本地测试时可行，但在部署到服务器后，用户无法直接访问服务器文件系统上的文件。为了解决这个问题，需要提供一种机制，让用户可以通过 HTTP 请求下载这个 PDF 文件。

解决方案

以下是几种解决 FastAPI 中 PDF 文件下载问题的方案：

1. 直接返回 PDF 文件内容

这种方法将 PDF 文件读取到内存中，然后直接作为 HTTP 响应返回。这种方式简单直接，适用于文件不大的情况。

原理: pdfkit 可以将 HTML 转换为字节流，FastAPI 可以直接返回字节流作为响应。设置响应头 Content-Disposition 为 attachment; filename=your_file_name.pdf 可以指示浏览器下载文件，并指定默认文件名。
步骤:
1. 使用 pdfkit.from_url 将 URL 转换为 PDF 字节流。
2. 使用 FastAPI 的 StreamingResponse 将字节流作为响应返回。
3. 设置响应头 Content-Disposition 指定文件名。

代码示例:

from typing import Optional
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import pdfkit
import io

app = FastAPI()

@app.post("/htmltopdf/{url}")
def convert_url(url:str):
  pdf_content = pdfkit.from_url(url, False)
  return StreamingResponse(io.BytesIO(pdf_content), media_type="application/pdf",headers={"Content-Disposition": f"attachment; filename=converted.pdf"})

安全建议: 对于非常大的 HTML 页面，这种方法可能导致内存占用过高。需要监控服务器的内存使用情况，并设置响应超时和文件大小限制。

2. 临时文件与 StaticFiles

这种方法将 PDF 文件保存为临时文件，然后使用 FastAPI 的 StaticFiles 功能提供文件下载。

原理: pdfkit 可以将 HTML 转换为文件并保存在磁盘上。FastAPI 可以配置静态文件路径，以便客户端可以访问。利用 Python 的 tempfile 模块创建临时文件可以避免文件冲突和磁盘空间占用问题。
步骤:
1. 使用 tempfile.NamedTemporaryFile 创建一个临时文件。
2. 使用 pdfkit.from_url 将 URL 转换为 PDF 文件并保存到临时文件中。
3. 使用 FastAPI 的 StaticFiles 挂载一个目录，并将临时文件所在的目录设置为静态文件目录。
4. 返回一个包含文件下载链接的响应。

代码示例:

from typing import Optional
from fastapi import FastAPI
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
import pdfkit
import tempfile
import os

app = FastAPI()

# 创建一个临时目录用于存储PDF文件
temp_dir = tempfile.gettempdir()
app.mount("/pdf", StaticFiles(directory=temp_dir), name="pdf")

@app.post("/htmltopdf/{url}")
def convert_url(url:str):
  with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf', dir=temp_dir) as tmp_file:
      pdfkit.from_url(url, tmp_file.name)
      file_path = tmp_file.name
      return FileResponse(path=file_path, filename='converted.pdf', media_type='application/pdf', headers={"Content-Disposition": "attachment; filename=converted.pdf"})

安全建议:
* 需要定期清理临时目录中的文件，避免磁盘空间不足。可以使用定时任务或者在文件下载完成后删除文件。
* 合理设置 StaticFiles 的挂载路径，避免暴露服务器的其他文件。
* 严格控制文件名，避免路径穿越攻击。

3. 后台任务与轮询

对于耗时较长的 HTML 转 PDF 任务，可以将转换过程放在后台任务中进行，并提供一个轮询接口供客户端查询转换状态和下载文件。

原理: 使用 FastAPI 的后台任务功能将 PDF 转换任务放到后台执行，避免阻塞主线程。客户端通过轮询接口查询任务状态，当任务完成后返回文件下载链接。
步骤:
1. 使用 FastAPI 的 BackgroundTasks 处理后台任务。
2. 定义一个后台任务函数，用于执行 PDF 转换并将结果保存到文件。
3. 定义一个轮询接口，返回任务状态和文件下载链接（如果任务已完成）。
4. 在主接口中触发后台任务，并返回一个任务 ID。

代码示例:
* 由于涉及 FastAPI 的 BackgroundTasks、多线程、文件管理等多个部分，这个例子仅供演示大致框架，完整实现需要仔细考虑代码逻辑，这里不展开。

from typing import Optional, Dict
from fastapi import FastAPI, BackgroundTasks
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
import pdfkit
import tempfile
import os
import uuid
import time

app = FastAPI()

# 存储任务状态和结果
tasks: Dict[str, Dict] = {}

# 临时目录用于存储PDF文件
temp_dir = tempfile.gettempdir()
app.mount("/pdf", StaticFiles(directory=temp_dir), name="pdf")

def convert_and_save_pdf(url: str, task_id: str):
    try:
      with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf', dir=temp_dir) as tmp_file:
          pdfkit.from_url(url, tmp_file.name)
          tasks[task_id]['file_path'] = tmp_file.name
          tasks[task_id]['status'] = 'completed'
    except Exception as e:
      tasks[task_id]['status'] = 'failed'
      tasks[task_id]['error'] = str(e)

@app.post("/htmltopdf/")
def convert_url(url:str, background_tasks: BackgroundTasks):
  task_id = str(uuid.uuid4())
  tasks[task_id] = {'status': 'pending', 'file_path': None}
  background_tasks.add_task(convert_and_save_pdf, url, task_id)
  return {'task_id': task_id, 'status': 'pending'}

@app.get("/tasks/{task_id}")
def get_task_status(task_id: str):
  task = tasks.get(task_id)
  if not task:
     return {'status': 'not found'}
  if task['status'] == 'completed':
    return FileResponse(path=task['file_path'], filename='converted.pdf', media_type='application/pdf', headers={"Content-Disposition": "attachment; filename=converted.pdf"})
  return task