|
1 | | -from typing import Any, Literal |
| 1 | +from typing import TYPE_CHECKING, Any, Literal |
2 | 2 |
|
3 | 3 | from drf_spectacular import generators, openapi |
4 | 4 | from drf_spectacular.extensions import ( |
@@ -47,11 +47,91 @@ class SchemaGenerator(generators.SchemaGenerator): |
47 | 47 | Adds a `$schema` property to the root schema object. |
48 | 48 | """ |
49 | 49 |
|
| 50 | + if TYPE_CHECKING: |
| 51 | + # Parent class has no type_hints, this is used to ignore the type error upstream |
| 52 | + def __init__(self, **kwargs: Any) -> None: |
| 53 | + super().__init__(**kwargs) # type: ignore[no-untyped-call] |
| 54 | + |
50 | 55 | def get_schema( |
51 | 56 | self, request: Request | None = None, public: bool = False |
52 | 57 | ) -> dict[str, Any]: |
53 | 58 | schema: dict[str, Any] = super().get_schema(request, public) # type: ignore[no-untyped-call] |
54 | | - schema["$schema"] = "https://spec.openapis.org/oas/3.1/dialect/base" |
| 59 | + return { |
| 60 | + "$schema": "https://spec.openapis.org/oas/3.1/dialect/base", |
| 61 | + **schema, |
| 62 | + } |
| 63 | + |
| 64 | + |
| 65 | +class MCPSchemaGenerator(SchemaGenerator): |
| 66 | + """ |
| 67 | + Schema generator that filters to only include operations tagged with "mcp". |
| 68 | +
|
| 69 | + Uses x-gram extension for Gram-native tool naming and descriptions. |
| 70 | + Gram reads x-gram directly from the spec. |
| 71 | + """ |
| 72 | + |
| 73 | + MCP_TAG = "mcp" |
| 74 | + MCP_SERVER_URL = "https://api.flagsmith.com" |
| 75 | + |
| 76 | + def get_schema( |
| 77 | + self, request: Request | None = None, public: bool = False |
| 78 | + ) -> dict[str, Any]: |
| 79 | + schema = super().get_schema(request, public) |
| 80 | + schema["paths"] = self._filter_paths(schema.get("paths", {})) |
| 81 | + schema = self._update_security_for_mcp(schema) |
| 82 | + schema.pop("$schema", None) |
| 83 | + info = schema.pop("info").copy() |
| 84 | + info["title"] = "mcp_openapi" |
| 85 | + return { |
| 86 | + "openapi": schema.pop("openapi"), |
| 87 | + "info": info, |
| 88 | + "servers": [{"url": self.MCP_SERVER_URL}], |
| 89 | + **schema, |
| 90 | + } |
| 91 | + |
| 92 | + def _filter_paths(self, paths: dict[str, Any]) -> dict[str, Any]: |
| 93 | + """Filter paths to only include operations tagged with 'mcp'.""" |
| 94 | + filtered_paths: dict[str, Any] = {} |
| 95 | + |
| 96 | + for path, path_item in paths.items(): |
| 97 | + filtered_operations: dict[str, Any] = {} |
| 98 | + has_any_mcp_tag = False |
| 99 | + |
| 100 | + for method, operation in path_item.items(): |
| 101 | + if not isinstance(operation, dict): |
| 102 | + filtered_operations[method] = operation |
| 103 | + continue |
| 104 | + |
| 105 | + tags = operation.get("tags", []) |
| 106 | + if self.MCP_TAG in tags: |
| 107 | + filtered_operations[method] = self._transform_for_mcp(operation) |
| 108 | + has_any_mcp_tag = True |
| 109 | + |
| 110 | + if has_any_mcp_tag: |
| 111 | + filtered_paths[path] = filtered_operations |
| 112 | + |
| 113 | + return filtered_paths |
| 114 | + |
| 115 | + def _transform_for_mcp(self, operation: dict[str, Any]) -> dict[str, Any]: |
| 116 | + """Apply MCP-specific transformations to an operation.""" |
| 117 | + operation = operation.copy() |
| 118 | + # Remove operation-level security (use global MCP security instead) |
| 119 | + operation.pop("security", None) |
| 120 | + return operation |
| 121 | + |
| 122 | + def _update_security_for_mcp(self, schema: dict[str, Any]) -> dict[str, Any]: |
| 123 | + """Update security schemes for MCP (Organisation API Key).""" |
| 124 | + schema = schema.copy() |
| 125 | + schema["components"] = schema.get("components", {}).copy() |
| 126 | + schema["components"]["securitySchemes"] = { |
| 127 | + "TOKEN_AUTH": { |
| 128 | + "type": "apiKey", |
| 129 | + "in": "header", |
| 130 | + "name": "Authorization", |
| 131 | + "description": "Organisation API Key. Format: Api-Key <key>", |
| 132 | + }, |
| 133 | + } |
| 134 | + schema["security"] = [{"TOKEN_AUTH": []}] |
55 | 135 | return schema |
56 | 136 |
|
57 | 137 |
|
|
0 commit comments