Add support for OpenAPI schema transformers (#56093)

Author: captainsafia

src/OpenApi/perf/Microbenchmarks/TransformersBenchmark.cs

@@ -4,6 +4,7 @@
 using BenchmarkDotNet.Attributes;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Routing;
+using Microsoft.OpenApi.Any;
 using Microsoft.OpenApi.Models;
 
 namespace Microsoft.AspNetCore.OpenApi.Microbenchmarks;
@@ -64,6 +65,28 @@ public void DocumentTransformerAsDelegate_Delegate()
         _documentService = CreateDocumentService(_builder, _options);
     }
 
+    [GlobalSetup(Target = nameof(SchemaTransformer))]
+    public void SchemaTransformer_Setup()
+    {
+        _builder.MapPost("/", (Todo todo) => todo);
+        for (var i = 0; i <= TransformerCount; i++)
+        {
+            _options.UseSchemaTransformer((schema, context, token) =>
+            {
+                if (context.Type == typeof(Todo) && context.ParameterDescription != null)
+                {
+                    schema.Extensions["x-my-extension"] = new OpenApiString(context.ParameterDescription.Name);
+                }
+                else
+                {
+                    schema.Extensions["x-my-extension"] = new OpenApiString("response");
+                }
+                return Task.CompletedTask;
+            });
+        }
+        _documentService = CreateDocumentService(_builder, _options);
+    }
+
     [Benchmark]
     public async Task OperationTransformerAsDelegate()
     {
@@ -82,6 +105,12 @@ public async Task DocumentTransformerAsDelegate()
         await _documentService.GetOpenApiDocumentAsync();
     }
 
+    [Benchmark]
+    public async Task SchemaTransformer()
+    {
+        await _documentService.GetOpenApiDocumentAsync();
+    }
+
     private class ActivatedTransformer : IOpenApiDocumentTransformer
     {
         public Task TransformAsync(OpenApiDocument document, OpenApiDocumentTransformerContext context, CancellationToken cancellationToken)

src/OpenApi/src/PublicAPI.Unshipped.txt

@@ -11,9 +11,20 @@ Microsoft.AspNetCore.OpenApi.OpenApiOptions.OpenApiVersion.set -> void
 Microsoft.AspNetCore.OpenApi.OpenApiOptions.ShouldInclude.get -> System.Func<Microsoft.AspNetCore.Mvc.ApiExplorer.ApiDescription!, bool>!
 Microsoft.AspNetCore.OpenApi.OpenApiOptions.ShouldInclude.set -> void
 Microsoft.AspNetCore.OpenApi.OpenApiOptions.UseOperationTransformer(System.Func<Microsoft.OpenApi.Models.OpenApiOperation!, Microsoft.AspNetCore.OpenApi.OpenApiOperationTransformerContext!, System.Threading.CancellationToken, System.Threading.Tasks.Task!>! transformer) -> Microsoft.AspNetCore.OpenApi.OpenApiOptions!
+Microsoft.AspNetCore.OpenApi.OpenApiOptions.UseSchemaTransformer(System.Func<Microsoft.OpenApi.Models.OpenApiSchema!, Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext!, System.Threading.CancellationToken, System.Threading.Tasks.Task!>! transformer) -> Microsoft.AspNetCore.OpenApi.OpenApiOptions!
 Microsoft.AspNetCore.OpenApi.OpenApiOptions.UseTransformer(Microsoft.AspNetCore.OpenApi.IOpenApiDocumentTransformer! transformer) -> Microsoft.AspNetCore.OpenApi.OpenApiOptions!
 Microsoft.AspNetCore.OpenApi.OpenApiOptions.UseTransformer(System.Func<Microsoft.OpenApi.Models.OpenApiDocument!, Microsoft.AspNetCore.OpenApi.OpenApiDocumentTransformerContext!, System.Threading.CancellationToken, System.Threading.Tasks.Task!>! transformer) -> Microsoft.AspNetCore.OpenApi.OpenApiOptions!
 Microsoft.AspNetCore.OpenApi.OpenApiOptions.UseTransformer<TTransformerType>() -> Microsoft.AspNetCore.OpenApi.OpenApiOptions!
+Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext
+Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext.ApplicationServices.get -> System.IServiceProvider!
+Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext.ApplicationServices.init -> void
+Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext.DocumentName.get -> string!
+Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext.DocumentName.init -> void
+Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext.OpenApiSchemaTransformerContext() -> void
+Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext.ParameterDescription.get -> Microsoft.AspNetCore.Mvc.ApiExplorer.ApiParameterDescription?
+Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext.ParameterDescription.init -> void
+Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext.Type.get -> System.Type!
+Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext.Type.init -> void
 Microsoft.Extensions.DependencyInjection.OpenApiServiceCollectionExtensions
 static Microsoft.AspNetCore.Builder.OpenApiEndpointRouteBuilderExtensions.MapOpenApi(this Microsoft.AspNetCore.Routing.IEndpointRouteBuilder! endpoints, string! pattern = "/openapi/{documentName}.json") -> Microsoft.AspNetCore.Builder.IEndpointConventionBuilder!
 static Microsoft.Extensions.DependencyInjection.OpenApiServiceCollectionExtensions.AddOpenApi(this Microsoft.Extensions.DependencyInjection.IServiceCollection! services) -> Microsoft.Extensions.DependencyInjection.IServiceCollection!

src/OpenApi/src/Services/OpenApiDocumentService.cs

@@ -55,7 +55,7 @@ public async Task<OpenApiDocument> GetOpenApiDocumentAsync(CancellationToken can
         var document = new OpenApiDocument
         {
             Info = GetOpenApiInfo(),
-            Paths = GetOpenApiPaths(capturedTags),
+            Paths = await GetOpenApiPathsAsync(capturedTags, cancellationToken),
             Tags = [.. capturedTags]
         };
         await ApplyTransformersAsync(document, cancellationToken);
@@ -99,7 +99,7 @@ internal OpenApiInfo GetOpenApiInfo()
     /// the object to support filtering each
     /// description instance into its appropriate document.
     /// </remarks>
-    private OpenApiPaths GetOpenApiPaths(HashSet<OpenApiTag> capturedTags)
+    private async Task<OpenApiPaths> GetOpenApiPathsAsync(HashSet<OpenApiTag> capturedTags, CancellationToken cancellationToken)
     {
         var descriptionsByPath = apiDescriptionGroupCollectionProvider.ApiDescriptionGroups.Items
             .SelectMany(group => group.Items)
@@ -109,17 +109,17 @@ private OpenApiPaths GetOpenApiPaths(HashSet<OpenApiTag> capturedTags)
         foreach (var descriptions in descriptionsByPath)
         {
             Debug.Assert(descriptions.Key != null, "Relative path mapped to OpenApiPath key cannot be null.");
-            paths.Add(descriptions.Key, new OpenApiPathItem { Operations = GetOperations(descriptions, capturedTags) });
+            paths.Add(descriptions.Key, new OpenApiPathItem { Operations = await GetOperationsAsync(descriptions, capturedTags, cancellationToken) });
         }
         return paths;
     }
 
-    private Dictionary<OperationType, OpenApiOperation> GetOperations(IGrouping<string?, ApiDescription> descriptions, HashSet<OpenApiTag> capturedTags)
+    private async Task<Dictionary<OperationType, OpenApiOperation>> GetOperationsAsync(IGrouping<string?, ApiDescription> descriptions, HashSet<OpenApiTag> capturedTags, CancellationToken cancellationToken)
     {
         var operations = new Dictionary<OperationType, OpenApiOperation>();
         foreach (var description in descriptions)
         {
-            var operation = GetOperation(description, capturedTags);
+            var operation = await GetOperationAsync(description, capturedTags, cancellationToken);
             operation.Extensions.Add(OpenApiConstants.DescriptionId, new OpenApiString(description.ActionDescriptor.Id));
             _operationTransformerContextCache.TryAdd(description.ActionDescriptor.Id, new OpenApiOperationTransformerContext
             {
@@ -132,7 +132,7 @@ private Dictionary<OperationType, OpenApiOperation> GetOperations(IGrouping<stri
         return operations;
     }
 
-    private OpenApiOperation GetOperation(ApiDescription description, HashSet<OpenApiTag> capturedTags)
+    private async Task<OpenApiOperation> GetOperationAsync(ApiDescription description, HashSet<OpenApiTag> capturedTags, CancellationToken cancellationToken)
     {
         var tags = GetTags(description);
         if (tags != null)
@@ -147,9 +147,9 @@ private OpenApiOperation GetOperation(ApiDescription description, HashSet<OpenAp
             OperationId = GetOperationId(description),
             Summary = GetSummary(description),
             Description = GetDescription(description),
-            Responses = GetResponses(description),
-            Parameters = GetParameters(description),
-            RequestBody = GetRequestBody(description),
+            Responses = await GetResponsesAsync(description, cancellationToken),
+            Parameters = await GetParametersAsync(description, cancellationToken),
+            RequestBody = await GetRequestBodyAsync(description, cancellationToken),
             Tags = tags,
         };
         return operation;
@@ -177,7 +177,7 @@ private OpenApiOperation GetOperation(ApiDescription description, HashSet<OpenAp
         return [new OpenApiTag { Name = description.ActionDescriptor.RouteValues["controller"] }];
     }
 
-    private OpenApiResponses GetResponses(ApiDescription description)
+    private async Task<OpenApiResponses> GetResponsesAsync(ApiDescription description, CancellationToken cancellationToken)
     {
         // OpenAPI requires that each operation have a response, usually a successful one.
         // if there are no response types defined, we assume a successful 200 OK response
@@ -186,7 +186,7 @@ private OpenApiResponses GetResponses(ApiDescription description)
         {
             return new OpenApiResponses
             {
-                ["200"] = GetResponse(description, StatusCodes.Status200OK, _defaultApiResponseType)
+                ["200"] = await GetResponseAsync(description, StatusCodes.Status200OK, _defaultApiResponseType, cancellationToken)
             };
         }
 
@@ -200,12 +200,12 @@ private OpenApiResponses GetResponses(ApiDescription description)
             var responseKey = responseType.IsDefaultResponse
                 ? OpenApiConstants.DefaultOpenApiResponseKey
                 : responseType.StatusCode.ToString(CultureInfo.InvariantCulture);
-            responses.Add(responseKey, GetResponse(description, responseType.StatusCode, responseType));
+            responses.Add(responseKey, await GetResponseAsync(description, responseType.StatusCode, responseType, cancellationToken));
         }
         return responses;
     }
 
-    private OpenApiResponse GetResponse(ApiDescription apiDescription, int statusCode, ApiResponseType apiResponseType)
+    private async Task<OpenApiResponse> GetResponseAsync(ApiDescription apiDescription, int statusCode, ApiResponseType apiResponseType, CancellationToken cancellationToken)
     {
         var description = ReasonPhrases.GetReasonPhrase(statusCode);
         var response = new OpenApiResponse
@@ -222,7 +222,7 @@ private OpenApiResponse GetResponse(ApiDescription apiDescription, int statusCod
             .Select(responseFormat => responseFormat.MediaType);
         foreach (var contentType in apiResponseFormatContentTypes)
         {
-            var schema = apiResponseType.Type is { } type ? _componentService.GetOrCreateSchema(type) : new OpenApiSchema();
+            var schema = apiResponseType.Type is { } type ? await _componentService.GetOrCreateSchemaAsync(type, null, cancellationToken) : new OpenApiSchema();
             response.Content[contentType] = new OpenApiMediaType { Schema = schema };
         }
 
@@ -240,7 +240,7 @@ private OpenApiResponse GetResponse(ApiDescription apiDescription, int statusCod
         return response;
     }
 
-    private List<OpenApiParameter>? GetParameters(ApiDescription description)
+    private async Task<List<OpenApiParameter>?> GetParametersAsync(ApiDescription description, CancellationToken cancellationToken)
     {
         List<OpenApiParameter>? parameters = null;
         foreach (var parameter in description.ParameterDescriptions)
@@ -265,32 +265,32 @@ private OpenApiResponse GetResponse(ApiDescription apiDescription, int statusCod
                 // Per the OpenAPI specification, parameters that are sourced from the path
                 // are always required, regardless of the requiredness status of the parameter.
                 Required = parameter.Source == BindingSource.Path || parameter.IsRequired,
-                Schema = _componentService.GetOrCreateSchema(parameter.Type, parameter),
+                Schema = await _componentService.GetOrCreateSchemaAsync(parameter.Type, parameter, cancellationToken),
             };
             parameters ??= [];
             parameters.Add(openApiParameter);
         }
         return parameters;
     }
 
-    private OpenApiRequestBody? GetRequestBody(ApiDescription description)
+    private async Task<OpenApiRequestBody?> GetRequestBodyAsync(ApiDescription description, CancellationToken cancellationToken)
     {
         // Only one parameter can be bound from the body in each request.
         if (description.TryGetBodyParameter(out var bodyParameter))
         {
-            return GetJsonRequestBody(description.SupportedRequestFormats, bodyParameter);
+            return await GetJsonRequestBody(description.SupportedRequestFormats, bodyParameter, cancellationToken);
         }
         // If there are no body parameters, check for form parameters.
         // Note: Form parameters and body parameters cannot exist simultaneously
         // in the same endpoint.
         if (description.TryGetFormParameters(out var formParameters))
         {
-            return GetFormRequestBody(description.SupportedRequestFormats, formParameters);
+            return await GetFormRequestBody(description.SupportedRequestFormats, formParameters, cancellationToken);
         }
         return null;
     }
 
-    private OpenApiRequestBody GetFormRequestBody(IList<ApiRequestFormat> supportedRequestFormats, IEnumerable<ApiParameterDescription> formParameters)
+    private async Task<OpenApiRequestBody> GetFormRequestBody(IList<ApiRequestFormat> supportedRequestFormats, IEnumerable<ApiParameterDescription> formParameters, CancellationToken cancellationToken)
     {
         if (supportedRequestFormats.Count == 0)
         {
@@ -325,7 +325,7 @@ private OpenApiRequestBody GetFormRequestBody(IList<ApiRequestFormat> supportedR
             if (parameter.All(parameter => parameter.ModelMetadata.ContainerType is null))
             {
                 var description = parameter.Single();
-                var parameterSchema = _componentService.GetOrCreateSchema(description.Type);
+                var parameterSchema = await _componentService.GetOrCreateSchemaAsync(description.Type, null, cancellationToken);
                 // Form files are keyed by their parameter name so we must capture the parameter name
                 // as a property in the schema.
                 if (description.Type == typeof(IFormFile) || description.Type == typeof(IFormFileCollection))
@@ -388,15 +388,15 @@ private OpenApiRequestBody GetFormRequestBody(IList<ApiRequestFormat> supportedR
                     var propertySchema = new OpenApiSchema { Type = "object", Properties = new Dictionary<string, OpenApiSchema>() };
                     foreach (var description in parameter)
                     {
-                        propertySchema.Properties[description.Name] = _componentService.GetOrCreateSchema(description.Type);
+                        propertySchema.Properties[description.Name] = await _componentService.GetOrCreateSchemaAsync(description.Type, null, cancellationToken);
                     }
                     schema.AllOf.Add(propertySchema);
                 }
                 else
                 {
                     foreach (var description in parameter)
                     {
-                        schema.Properties[description.Name] = _componentService.GetOrCreateSchema(description.Type);
+                        schema.Properties[description.Name] = await _componentService.GetOrCreateSchemaAsync(description.Type, null, cancellationToken);
                     }
                 }
             }
@@ -415,7 +415,7 @@ private OpenApiRequestBody GetFormRequestBody(IList<ApiRequestFormat> supportedR
         return requestBody;
     }
 
-    private OpenApiRequestBody GetJsonRequestBody(IList<ApiRequestFormat> supportedRequestFormats, ApiParameterDescription bodyParameter)
+    private async Task<OpenApiRequestBody> GetJsonRequestBody(IList<ApiRequestFormat> supportedRequestFormats, ApiParameterDescription bodyParameter, CancellationToken cancellationToken)
     {
         if (supportedRequestFormats.Count == 0)
         {
@@ -442,7 +442,7 @@ private OpenApiRequestBody GetJsonRequestBody(IList<ApiRequestFormat> supportedR
         foreach (var requestForm in supportedRequestFormats)
         {
             var contentType = requestForm.MediaType;
-            requestBody.Content[contentType] = new OpenApiMediaType { Schema = _componentService.GetOrCreateSchema(bodyParameter.Type) };
+            requestBody.Content[contentType] = new OpenApiMediaType { Schema = await _componentService.GetOrCreateSchemaAsync(bodyParameter.Type, bodyParameter, cancellationToken) };
         }
 
         return requestBody;

src/OpenApi/src/Services/OpenApiOptions.cs

@@ -14,6 +14,7 @@ namespace Microsoft.AspNetCore.OpenApi;
 public sealed class OpenApiOptions
 {
     internal readonly List<IOpenApiDocumentTransformer> DocumentTransformers = [];
+    internal readonly List<Func<OpenApiSchema, OpenApiSchemaTransformerContext, CancellationToken, Task>> SchemaTransformers = [];
 
     /// <summary>
     /// Initializes a new instance of the <see cref="OpenApiOptions"/> class
@@ -89,4 +90,17 @@ public OpenApiOptions UseOperationTransformer(Func<OpenApiOperation, OpenApiOper
         DocumentTransformers.Add(new DelegateOpenApiDocumentTransformer(transformer));
         return this;
     }
+
+    /// <summary>
+    /// Registers a given delegate as a schema transformer on the current <see cref="OpenApiOptions"/> instance.
+    /// </summary>
+    /// <param name="transformer">The delegate representing the schema transformer.</param>
+    /// <returns>The <see cref="OpenApiOptions"/> instance for further customization.</returns>
+    public OpenApiOptions UseSchemaTransformer(Func<OpenApiSchema, OpenApiSchemaTransformerContext, CancellationToken, Task> transformer)
+    {
+        ArgumentNullException.ThrowIfNull(transformer, nameof(transformer));
+
+        SchemaTransformers.Add(transformer);
+        return this;
+    }
 }