Fastapi 请求结束后删除临时文件

背景

最近要用 Fastapi 框架简单实现一个文件下载接口，需要临时创建一个目录，将需要下载的文件打包为 zip 文件，返回给前端后，需要把这个 zip 文件删除。一开始我想到了两个思路：1. 维护一个定时器，定时去删除。2. 创建临时文件夹，接口返回后自动清理掉这个文件夹。

对比当然选择第二个方案，因为实现比较简单。

简单 demo 实现如下：

import os
import tempfile
import zipfile

import uvicorn
from fastapi import FastAPI
from fastapi.responses import FileResponse

app = FastAPI()


@app.get("/file/download")
async def file_download():
    base_dir = os.path.dirname(__file__)

    # Create temp dir
    tmp_dir_obj = tempfile.TemporaryDirectory(dir=base_dir)
    tmp_dirname = tmp_dir_obj.name

    # File to be zipped
    file_path = f'{base_dir}/data/test.txt'

    # Zip file
    zip_file_path = f'{tmp_dirname}/test.zip'
    with zipfile.ZipFile(zip_file_path, 'w', zipfile.ZIP_DEFLATED) as fp:
        fp.write(filename=file_path, arcname='test.txt')

    zip_filename = os.path.basename(zip_file_path)

    return FileResponse(path=zip_file_path, filename=zip_filename)


if __name__ == '__main__':
    uvicorn.run(app, port=8000)

可是当我使用 python 内置方法创建了临时文件夹，运行后却报以下错误：

image

找不到 zip 文件，因为 Fastapi 的请求返回是异步处理的，当前进程结束后，文件就被清理掉了，所以找不到这个文件对象。

于是乎我陷入了沉思。

因为实在不想用第一种实现（维护定时器还是觉得太麻烦了），我想了另一个思路：每次调这个下载接口，都会先去删除 zip 文件目录，这样可以保证服务器上最多只是保存一份 zip 文件。

但这种方式还是感觉很不优雅。因为上传文件时，要对保存上传文件的目录大小做限制，当超过设定值时要把最早的文件删除掉，以保证新的文件可以正常上传成功。上传文件目录和 zip 文件目录是在同一个目录下的，所以每次判断大小时，还要额外再判断下这个 zip 文件的大小。

解决

后面无意中发现了 Fastapi 是可以在某个请求结束后，后台执行任务的，主要借助了 starlette 这个包。具体实现如下：

import os
import zipfile

import uvicorn
from fastapi import FastAPI
from fastapi.responses import FileResponse
from starlette.background import BackgroundTask

app = FastAPI()


def rm_file(file_path):
    if os.path.exists(file_path):
        os.remove(file_path)


@app.get("/file/download")
async def file_download():
    base_dir = os.path.dirname(__file__)

    # File to be zipped
    file_path = f'{base_dir}/data/test.txt'

    # Zip file
    zip_file_path = f'{base_dir}/test.zip'
    with zipfile.ZipFile(zip_file_path, 'w', zipfile.ZIP_DEFLATED) as fp:
        fp.write(filename=file_path, arcname='test.txt')

    zip_filename = os.path.basename(zip_file_path)

    # Add background task
    task = BackgroundTask(rm_file, file_path=zip_file_path)

    return FileResponse(path=zip_file_path, filename=zip_filename, background=task)


if __name__ == '__main__':
    uvicorn.run(app, port=8000)

这样就可以完美清理掉临时生成的 zip 文件了。

参考

https://www.starlette.io/background/