Skip to content

背景

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

找到耗时的接口

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

在项目根目录下新增一个drf-yasggenerator

generators.py

python
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
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
$ python manage.py generate_swagger test.json -g generators.TimeCostSchemaGenerator -o

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

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

参数解释

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

Last updated:

Released under the MIT License.