使用drf-spectacular为不同API模块定制全局请求头

在Django REST框架项目中，我们经常需要为不同的API模块设置不同的请求头要求。本文将介绍如何利用drf-spectacular这个强大的OpenAPI 3.0文档生成工具，为内部API和公共API分别配置不同的请求头规范。

需求场景分析

在一个典型的Django REST框架项目中，我们可能会有两类API端点：

1. 内部API：需要特定的认证头信息
2. 公共API：不需要额外的认证头

直接修改DEFAULT_SCHEMA_CLASS会影响所有端点，这不是我们想要的结果。我们需要一种更精细的控制方式。

解决方案：利用认证类扩展

drf-spectacular提供了一个优雅的解决方案——通过OpenAPIAuthenticationExtension来为不同的认证类指定不同的请求头。

实现步骤

1. 创建自定义认证类： 即使内部API和公共API使用相同的认证机制，我们也可以创建空子类来区分它们：

from rest_framework.authentication import BasicAuthentication

class InternalAPIAuth(BasicAuthentication):
    pass

class PublicAPIAuth(BasicAuthentication):
    pass

2. 实现认证扩展： 为内部API认证类创建扩展，指定需要的头信息：

from drf_spectacular.extensions import OpenApiAuthenticationExtension

class InternalAuthExtension(OpenApiAuthenticationExtension):
    target_class = 'path.to.InternalAPIAuth'
    name = 'internalAuth'
    
    def get_security_definition(self, auto_schema):
        return {
            'type': 'apiKey',
            'in': 'header',
            'name': 'X-Custom-Header1',
            'description': '内部API专用头信息1'
        }, {
            'type': 'apiKey',
            'in': 'header',
            'name': 'X-Custom-Header2',
            'description': '内部API专用头信息2'
        }

3. 应用到视图： 在内部API视图上使用自定义认证类：

from rest_framework.views import APIView

class InternalView(APIView):
    authentication_classes = [InternalAPIAuth]
    # ... 其他视图代码

替代方案：后处理钩子

如果上述方法不能满足需求，还可以使用后处理钩子手动修改生成的schema：

def custom_postprocessing_hook(result, generator, request, public):
    # 根据路径或其他条件判断是否为内部API
    if 'internal' in request.path:
        for path_item in result['paths'].values():
            for operation in path_item.values():
                if 'parameters' not in operation:
                    operation['parameters'] = []
                operation['parameters'].extend([
                    {
                        'name': 'X-Custom-Header1',
                        'in': 'header',
                        'required': True,
                        'schema': {'type': 'string'}
                    },
                    {
                        'name': 'X-Custom-Header2',
                        'in': 'header',
                        'required': True,
                        'schema': {'type': 'string'}
                    }
                ])
    return result

然后在设置中注册这个钩子：

SPECTACULAR_SETTINGS = {
    'POSTPROCESSING_HOOKS': [
        'path.to.custom_postprocessing_hook',
    ]
}

最佳实践建议

1. 优先使用认证扩展：这种方式更符合DRF的设计哲学，与认证系统紧密集成
2. 保持一致性：确保文档中的头信息要求与实际API行为一致
3. 明确区分：使用清晰的命名区分不同API类型，便于维护
4. 文档说明：在API文档中清楚地说明不同端点的头信息要求

通过这种方式，我们可以灵活地为不同类型的API端点配置不同的请求头规范，同时保持代码的整洁和可维护性。