api: optimise pagination in API schema (#6478)

goauthentik · Aug 5, 2023 · 00fae23 · 00fae23
1 parent efc6609
commit 00fae23
Show file tree

Hide file tree

Showing 9 changed files with 177 additions and 2,052 deletions.
diff --git a/authentik/api/pagination.py b/authentik/api/pagination.py
@@ -2,6 +2,43 @@
 from rest_framework import pagination
 from rest_framework.response import Response
 
+PAGINATION_COMPONENT_NAME = "Pagination"
+PAGINATION_SCHEMA = {
+    "type": "object",
+    "properties": {
+        "next": {
+            "type": "number",
+        },
+        "previous": {
+            "type": "number",
+        },
+        "count": {
+            "type": "number",
+        },
+        "current": {
+            "type": "number",
+        },
+        "total_pages": {
+            "type": "number",
+        },
+        "start_index": {
+            "type": "number",
+        },
+        "end_index": {
+            "type": "number",
+        },
+    },
+    "required": [
+        "next",
+        "previous",
+        "count",
+        "current",
+        "total_pages",
+        "start_index",
+        "end_index",
+    ],
+}
+
 
 class Pagination(pagination.PageNumberPagination):
     """Pagination which includes total pages and current page"""
@@ -35,41 +72,7 @@ def get_paginated_response_schema(self, schema):
         return {
             "type": "object",
             "properties": {
-                "pagination": {
-                    "type": "object",
-                    "properties": {
-                        "next": {
-                            "type": "number",
-                        },
-                        "previous": {
-                            "type": "number",
-                        },
-                        "count": {
-                            "type": "number",
-                        },
-                        "current": {
-                            "type": "number",
-                        },
-                        "total_pages": {
-                            "type": "number",
-                        },
-                        "start_index": {
-                            "type": "number",
-                        },
-                        "end_index": {
-                            "type": "number",
-                        },
-                    },
-                    "required": [
-                        "next",
-                        "previous",
-                        "count",
-                        "current",
-                        "total_pages",
-                        "start_index",
-                        "end_index",
-                    ],
-                },
+                "pagination": {"$ref": f"#/components/schemas/{PAGINATION_COMPONENT_NAME}"},
                 "results": schema,
             },
             "required": ["pagination", "results"],

diff --git a/authentik/api/schema.py b/authentik/api/schema.py
@@ -1,5 +1,6 @@
 """Error Response schema, from https://github.com/axnsan12/drf-yasg/issues/224"""
 from django.utils.translation import gettext_lazy as _
+from drf_spectacular.generators import SchemaGenerator
 from drf_spectacular.plumbing import (
     ResolvedComponent,
     build_array_type,
@@ -9,6 +10,8 @@
 from drf_spectacular.settings import spectacular_settings
 from drf_spectacular.types import OpenApiTypes
 
+from authentik.api.pagination import PAGINATION_COMPONENT_NAME, PAGINATION_SCHEMA
+
 
 def build_standard_type(obj, **kwargs):
     """Build a basic type with optional add owns."""
@@ -36,27 +39,30 @@ def build_standard_type(obj, **kwargs):
 )
 
 
-def postprocess_schema_responses(result, generator, **kwargs):  # noqa: W0613
+def create_component(generator: SchemaGenerator, name, schema, type_=ResolvedComponent.SCHEMA):
+    """Register a component and return a reference to it."""
+    component = ResolvedComponent(
+        name=name,
+        type=type_,
+        schema=schema,
+        object=name,
+    )
+    generator.registry.register_on_missing(component)
+    return component
+
+
+def postprocess_schema_responses(result, generator: SchemaGenerator, **kwargs):  # noqa: W0613
     """Workaround to set a default response for endpoints.
     Workaround suggested at
     <https://github.com/tfranzel/drf-spectacular/issues/119#issuecomment-656970357>
     for the missing drf-spectacular feature discussed in
     <https://github.com/tfranzel/drf-spectacular/issues/101>.
     """
 
-    def create_component(name, schema, type_=ResolvedComponent.SCHEMA):
-        """Register a component and return a reference to it."""
-        component = ResolvedComponent(
-            name=name,
-            type=type_,
-            schema=schema,
-            object=name,
-        )
-        generator.registry.register_on_missing(component)
-        return component
+    create_component(generator, PAGINATION_COMPONENT_NAME, PAGINATION_SCHEMA)
 
-    generic_error = create_component("GenericError", GENERIC_ERROR)
-    validation_error = create_component("ValidationError", VALIDATION_ERROR)
+    generic_error = create_component(generator, "GenericError", GENERIC_ERROR)
+    validation_error = create_component(generator, "ValidationError", VALIDATION_ERROR)
 
     for path in result["paths"].values():
         for method in path.values():