drf-yasg打开文档速度很慢的解决方案

背景

在使用drf-yasg的时候，有时候一不小心写了一些代码，比如写serializers的时候写了一些default参数，该参数是会查询数据库的，打开swagger文档的时候就会触发数据库查询，从而慢慢的增加了打开swagger文档的时间

找到耗时的接口

需要安装依赖pip install prettytable方便打印耗时接口

在项目根目录下新增一个drf-yasg的generator

generators.py

import time

import prettytable
from drf_yasg.generators import OpenAPISchemaGenerator


class TimeCostSchemaGenerator(OpenAPISchemaGenerator):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.path_method_cost_map = {}
        self.table = prettytable.PrettyTable(field_names=['cost', 'method', 'url'])

    def get_operation(self, view, path, prefix, method, components, request):
        """计算获取每个接口的schema的时间

        如果接口操作的耗时小于0.01秒，则不记录该接口
        """
        now = time.perf_counter()
        operation = super().get_operation(view, path, prefix, method, components, request)
        cost = round(time.perf_counter() - now, 2)
        if cost:
            self.path_method_cost_map[(path, method)] = cost
        return operation

    def get_schema(self, request=None, public=False):
        """在获取schema的时候打印记录下来的接口和耗时表格"""
        schema = super().get_schema(request, public)
        path_method_cost_items = sorted(self.path_method_cost_map.items(), key=lambda i: i[1], reverse=True)
        for ((path, method), cost) in path_method_cost_items:
            self.table.add_row([cost, method, path])
        print(self.table)
        return schema

执行命令验证

$ python manage.py generate_swagger test.json -g generators.TimeCostSchemaGenerator -o

上述命令会打印出如下格式的输出，表格三列表示接口耗时，请求方法，请求url，按照耗时倒序排列

+-------+--------+---------------------------------------------------------------+
|  cost | method |                              url                              |
+-------+--------+---------------------------------------------------------------+
| 15.36 |  GET   |                /api/v2/**/                                    |
|  3.57 |  GET   |               /api/***                                        |
+-------+--------+---------------------------------------------------------------+

参数解释

• generate_swagger是drf-yasg自定义的命令，安装完drf-yasg之后可以直接使用
• test.json是生成schema文件的路径
• -g generators.TimeCostSchemaGenerator表示指定自定义generator的路径位置，程序执行的时候会按照指定参数导入
• -o表示会覆盖旧的schema文件