feat(api-docs): Generate OpenAPI 3.1 docs with drf-spectacular

api/api/apps.py

@@ -5,3 +5,7 @@
 
 class ApiConfig(AppConfig):
     name = "api"
+
+    def ready(self) -> None:
+        # Import openapi extensions to register them with drf-spectacular
+        from api import openapi  # noqa: F401

api/api/openapi.py

@@ -1,98 +1,102 @@
-import inspect
-from typing import Any
-
-from drf_yasg.inspectors import SwaggerAutoSchema  # type: ignore[import-untyped]
-from drf_yasg.openapi import (  # type: ignore[import-untyped]
-    SCHEMA_DEFINITIONS,
-    Response,
-    Schema,
+from typing import Any, Literal
+
+from drf_spectacular import generators, openapi
+from drf_spectacular.extensions import (
+    OpenApiSerializerExtension,
 )
+from drf_spectacular.plumbing import ResolvedComponent, safe_ref
+from drf_spectacular.plumbing import append_meta as append_meta_orig
 from pydantic import BaseModel
-from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
-from pydantic_core import core_schema
+from rest_framework.request import Request
 
 
-class _GenerateJsonSchema(GenerateJsonSchema):
-    def nullable_schema(self, schema: core_schema.NullableSchema) -> JsonSchemaValue:
-        """Generates an OpenAPI 2.0-compatible JSON schema that matches a schema that allows null values.
+def append_meta(schema: dict[str, Any], meta: dict[str, Any]) -> dict[str, Any]:
+    """
+    See https://github.com/tfranzel/drf-spectacular/issues/1480
+    """
+    try:
+        return append_meta_orig(schema, meta)
+    except AssertionError as exc:
+        if str(exc) == "Invalid nullable case":
+            pass
+        else:  # pragma: no cover
+            raise exc
 
-        (The catch is OpenAPI 2.0 does not allow them, but some clients are capable
-        to consume the `x-nullable` annotation.)
+    if any("nullable" in d for d in (schema, meta)) and "oneOf" in schema:
+        schema = schema.copy()
+        meta = meta.copy()
 
-        Args:
-            schema: The core schema.
+        schema.pop("nullable", None)
+        meta.pop("nullable", None)
 
-        Returns:
-            The generated JSON schema.
-        """
-        anyof_schema_value = super().nullable_schema(schema)
-        elem = next(
-            any_of
-            for any_of in anyof_schema_value["anyOf"]
-            if any_of.get("type") != "null"
-        )
-        if type := elem.get("type"):
-            return {"type": type, "x-nullable": True}
-        # Assuming a reference here (which we can not annotate)
-        return elem  # type: ignore[no-any-return]
+        schema["oneOf"].append({"type": "null"})
+
+    if "exclusiveMinimum" in schema and "minimum" in schema:  # pragma: no cover
+        schema["exclusiveMinimum"] = schema.pop("minimum")
+    if "exclusiveMaximum" in schema and "maximum" in schema:  # pragma: no cover
+        schema["exclusiveMaximum"] = schema.pop("maximum")
+
+    return safe_ref({**schema, **meta})
+
+
+openapi.append_meta = append_meta  # type: ignore[attr-defined]
 
 
-class PydanticResponseCapableSwaggerAutoSchema(SwaggerAutoSchema):  # type: ignore[misc]
+class SchemaGenerator(generators.SchemaGenerator):
+    """
+    Adds a `$schema` property to the root schema object.
     """
-    A `SwaggerAutoSchema` subclass that allows to generate view response Swagger docs
-    from a Pydantic model.
 
-    Example usage:
+    def get_schema(
+        self, request: Request | None = None, public: bool = False
+    ) -> dict[str, Any]:
+        schema: dict[str, Any] = super().get_schema(request, public)  # type: ignore[no-untyped-call]
+        schema["$schema"] = "https://spec.openapis.org/oas/3.1/dialect/base"
+        return schema
 
-    ```
-    @drf_yasg.utils.swagger_auto_schema(
-        responses={200: YourPydanticSchema},
-        auto_schema=PydanticResponseCapableSwaggerAutoSchema,
-    )
-    def your_view(): ...
-    ```
 
-    To adapt Pydantic-generated schemas, the following is taken care of:
+class PydanticSchemaExtension(
+    OpenApiSerializerExtension  # type: ignore[no-untyped-call]
+):
+    """
+    An OpenAPI extension that allows drf-spectacular to generate schema documentation
+    from Pydantic models.
 
-    1. Pydantic-generated definitions are unwrapped and added to drf-yasg's global definitions.
-    2. Rather than using `anyOf`, nullable fields are annotated with `x-nullable`.
-    3. As there's no way to annotate a reference, all nested models are assumed to be `x-nullable`.
+    This extension is automatically used when a Pydantic BaseModel subclass is passed
+    as a response type in @extend_schema decorators.
     """
 
-    def get_response_schemas(
+    target_class = "pydantic.BaseModel"
+    match_subclasses = True
+
+    def get_name(
         self,
-        response_serializers: dict[str | int, Any],
-    ) -> dict[str, Response]:
-        result = {}
-
-        definitions = self.components.with_scope(SCHEMA_DEFINITIONS)
-
-        for status_code in list(response_serializers):
-            if inspect.isclass(response_serializers[status_code]) and issubclass(
-                model_cls := response_serializers[status_code], BaseModel
-            ):
-                model_json_schema = model_cls.model_json_schema(
-                    mode="serialization",
-                    schema_generator=_GenerateJsonSchema,
-                    ref_template=f"#/{SCHEMA_DEFINITIONS}/{{model}}",
-                )
+        auto_schema: openapi.AutoSchema | None = None,
+        direction: Literal["request", "response"] | None = None,
+    ) -> str | None:
+        return self.target.__name__  # type: ignore[no-any-return]
 
-                for ref_name, schema_kwargs in model_json_schema.pop("$defs").items():
-                    definitions.setdefault(
-                        ref_name,
-                        maker=lambda: Schema(
-                            **schema_kwargs,
-                            # We can not annotate references with `x-nullable`,
-                            # So just assume all nested models as nullable for now.
-                            x_nullable=True,
-                        ),
-                    )
-
-                result[str(status_code)] = Response(
-                    description=model_cls.__name__,
-                    schema=Schema(**model_json_schema),
-                )
+    def map_serializer(
+        self,
+        auto_schema: openapi.AutoSchema,
+        direction: str,
+    ) -> dict[str, Any]:
+        model_cls: type[BaseModel] = self.target
+
+        model_json_schema = model_cls.model_json_schema(
+            mode="serialization",
+            ref_template="#/components/schemas/{model}",
+        )
 
-                del response_serializers[status_code]
+        # Register nested definitions as components
+        if "$defs" in model_json_schema:
+            for ref_name, schema_kwargs in model_json_schema.pop("$defs").items():
+                component = ResolvedComponent(  # type: ignore[no-untyped-call]
+                    name=ref_name,
+                    type=ResolvedComponent.SCHEMA,
+                    object=ref_name,
+                    schema=schema_kwargs,
+                )
+                auto_schema.registry.register_on_missing(component)
 
-        return {**super().get_response_schemas(response_serializers), **result}
+        return model_json_schema

api/api/urls/v1.py

@@ -1,8 +1,11 @@
 from django.conf import settings
 from django.urls import include, path, re_path
-from drf_yasg import openapi  # type: ignore[import-untyped]
-from drf_yasg.views import get_schema_view  # type: ignore[import-untyped]
-from rest_framework import authentication, permissions, routers
+from drf_spectacular.views import (
+    SpectacularJSONAPIView,
+    SpectacularSwaggerView,
+    SpectacularYAMLAPIView,
+)
+from rest_framework import permissions, routers
 
 from app_analytics.views import SDKAnalyticsFlags, SelfHostedTelemetryAPIView
 from environments.identities.traits.views import SDKTraits
@@ -19,19 +22,6 @@
     else permissions.AllowAny
 )
 
-schema_view = get_schema_view(
-    openapi.Info(
-        title="Flagsmith API",
-        default_version="v1",
-        description="",
-        license=openapi.License(name="BSD License"),
-        contact=openapi.Contact(email="support@flagsmith.com"),
-    ),
-    public=True,
-    permission_classes=[schema_view_permission_class],
-    authentication_classes=[authentication.BasicAuthentication],
-)
-
 traits_router = routers.DefaultRouter()
 traits_router.register(r"", SDKTraits, basename="sdk-traits")
 
@@ -81,14 +71,26 @@
     ),
     re_path("", include("features.versioning.urls", namespace="versioning")),
     # API documentation
-    re_path(
-        r"^swagger(?P<format>\.json|\.yaml)$",
-        schema_view.without_ui(cache_timeout=0),
+    path(
+        "swagger.json",
+        SpectacularJSONAPIView.as_view(
+            permission_classes=[schema_view_permission_class],
+        ),
         name="schema-json",
     ),
-    re_path(
-        r"^docs/$",
-        schema_view.with_ui("swagger", cache_timeout=0),
+    path(
+        "swagger.yaml",
+        SpectacularYAMLAPIView.as_view(
+            permission_classes=[schema_view_permission_class],
+        ),
+        name="schema-yaml",
+    ),
+    path(
+        "docs/",
+        SpectacularSwaggerView.as_view(
+            url_name="v1:schema-json",
+            permission_classes=[schema_view_permission_class],
+        ),
         name="schema-swagger-ui",
     ),
     # Test webhook url

api/app/pagination.py

@@ -1,11 +1,10 @@
 import base64
 import json
 from collections import OrderedDict
+from typing import Any
 
-from drf_yasg import openapi  # type: ignore[import-untyped]
-from drf_yasg.inspectors import PaginatorInspector  # type: ignore[import-untyped]
 from flag_engine.identities.models import IdentityModel
-from rest_framework.pagination import BasePagination, PageNumberPagination
+from rest_framework.pagination import PageNumberPagination
 from rest_framework.response import Response
 
 
@@ -15,56 +14,6 @@ class CustomPagination(PageNumberPagination):
     max_page_size = 999
 
 
-class EdgeIdentityPaginationInspector(PaginatorInspector):  # type: ignore[misc]
-    def get_paginator_parameters(
-        self, paginator: BasePagination
-    ) -> list[openapi.Parameter]:
-        """
-        :param BasePagination paginator: the paginator
-        :rtype: list[openapi.Parameter]
-        """
-        return [
-            openapi.Parameter(
-                "page_size",
-                openapi.IN_QUERY,
-                "Number of results to return per page.",
-                required=False,
-                type=openapi.TYPE_INTEGER,
-            ),
-            openapi.Parameter(
-                "last_evaluated_key",
-                openapi.IN_QUERY,
-                "Used as the starting point for the page",
-                required=False,
-                type=openapi.TYPE_STRING,
-            ),
-        ]
-
-    def get_paginated_response(self, paginator, response_schema):  # type: ignore[no-untyped-def]
-        """
-        :param BasePagination paginator: the paginator
-        :param openapi.Schema response_schema: the response schema that must be paged.
-        :rtype: openapi.Schema
-        """
-
-        return openapi.Schema(
-            type=openapi.TYPE_OBJECT,
-            properties=OrderedDict(
-                (
-                    (
-                        "last_evaluated_key",
-                        openapi.Schema(
-                            type=openapi.TYPE_STRING,
-                            x_nullable=True,
-                        ),
-                    ),
-                    ("results", response_schema),
-                )
-            ),
-            required=["results"],
-        )
-
-
 class EdgeIdentityPagination(CustomPagination):
     max_page_size = 100
     page_size = 100
@@ -105,3 +54,42 @@ def get_paginated_response(self, data) -> Response:  # type: ignore[no-untyped-d
                 ]
             )
         )
+
+    def get_schema_operation_parameters(self, view: Any) -> list[dict[str, Any]]:
+        """
+        Returns the OpenAPI parameters for the pagination.
+        This is used by drf-spectacular to generate the schema.
+        """
+        return [
+            {
+                "name": "page_size",
+                "in": "query",
+                "description": "Number of results to return per page.",
+                "required": False,
+                "schema": {"type": "integer"},
+            },
+            {
+                "name": "last_evaluated_key",
+                "in": "query",
+                "description": "Used as the starting point for the page",
+                "required": False,
+                "schema": {"type": "string"},
+            },
+        ]
+
+    def get_paginated_response_schema(self, schema: dict[str, Any]) -> dict[str, Any]:
+        """
+        Returns the OpenAPI schema for the paginated response.
+        This is used by drf-spectacular to generate the schema.
+        """
+        return {
+            "type": "object",
+            "required": ["results"],
+            "properties": {
+                "last_evaluated_key": {
+                    "type": "string",
+                    "nullable": True,
+                },
+                "results": schema,
+            },
+        }