Support generating OpenApiSchemas from transformers (#61050)

Author: captainsafia

src/OpenApi/src/PublicAPI.Unshipped.txt

@@ -1,2 +1,9 @@
 #nullable enable
 static Microsoft.AspNetCore.Builder.OpenApiEndpointConventionBuilderExtensions.AddOpenApiOperationTransformer<TBuilder>(this TBuilder builder, System.Func<Microsoft.OpenApi.Models.OpenApiOperation!, Microsoft.AspNetCore.OpenApi.OpenApiOperationTransformerContext!, System.Threading.CancellationToken, System.Threading.Tasks.Task!>! transformer) -> TBuilder
+Microsoft.AspNetCore.OpenApi.OpenApiDocumentTransformerContext.GetOrCreateSchemaAsync(System.Type! type, Microsoft.AspNetCore.Mvc.ApiExplorer.ApiParameterDescription? parameterDescription = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task<Microsoft.OpenApi.Models.OpenApiSchema!>!
+Microsoft.AspNetCore.OpenApi.OpenApiOperationTransformerContext.Document.get -> Microsoft.OpenApi.Models.OpenApiDocument?
+Microsoft.AspNetCore.OpenApi.OpenApiOperationTransformerContext.Document.init -> void
+Microsoft.AspNetCore.OpenApi.OpenApiOperationTransformerContext.GetOrCreateSchemaAsync(System.Type! type, Microsoft.AspNetCore.Mvc.ApiExplorer.ApiParameterDescription? parameterDescription = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task<Microsoft.OpenApi.Models.OpenApiSchema!>!
+Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext.Document.get -> Microsoft.OpenApi.Models.OpenApiDocument?
+Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext.Document.init -> void
+Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext.GetOrCreateSchemaAsync(System.Type! type, Microsoft.AspNetCore.Mvc.ApiExplorer.ApiParameterDescription? parameterDescription = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task<Microsoft.OpenApi.Models.OpenApiSchema!>!

src/OpenApi/src/Services/OpenApiDocumentService.cs

@@ -77,7 +77,7 @@ public async Task<OpenApiDocument> GetOpenApiDocumentAsync(IServiceProvider scop
         document.Paths = await GetOpenApiPathsAsync(document, scopedServiceProvider, operationTransformers, schemaTransformers, cancellationToken);
         try
         {
-            await ApplyTransformersAsync(document, scopedServiceProvider, cancellationToken);
+            await ApplyTransformersAsync(document, scopedServiceProvider, schemaTransformers, cancellationToken);
         }
 
         finally
@@ -95,13 +95,15 @@ public async Task<OpenApiDocument> GetOpenApiDocumentAsync(IServiceProvider scop
         return document;
     }
 
-    private async Task ApplyTransformersAsync(OpenApiDocument document, IServiceProvider scopedServiceProvider, CancellationToken cancellationToken)
+    private async Task ApplyTransformersAsync(OpenApiDocument document, IServiceProvider scopedServiceProvider, IOpenApiSchemaTransformer[] schemaTransformers, CancellationToken cancellationToken)
     {
         var documentTransformerContext = new OpenApiDocumentTransformerContext
         {
             DocumentName = documentName,
             ApplicationServices = scopedServiceProvider,
             DescriptionGroups = apiDescriptionGroupCollectionProvider.ApiDescriptionGroups.Items,
+            Document = document,
+            SchemaTransformers = schemaTransformers
         };
         // Use index-based for loop to avoid allocating an enumerator with a foreach.
         for (var i = 0; i < _options.DocumentTransformers.Count; i++)
@@ -271,6 +273,8 @@ private async Task<Dictionary<OperationType, OpenApiOperation>> GetOperationsAsy
                 DocumentName = documentName,
                 Description = description,
                 ApplicationServices = scopedServiceProvider,
+                Document = document,
+                SchemaTransformers = schemaTransformers
             };
 
             _operationTransformerContextCache.TryAdd(description.ActionDescriptor.Id, operationContext);

src/OpenApi/src/Services/Schemas/OpenApiSchemaService.cs

@@ -118,7 +118,7 @@ internal sealed class OpenApiSchemaService(
         }
     };
 
-    internal async Task<IOpenApiSchema> GetOrCreateSchemaAsync(OpenApiDocument document, Type type, IServiceProvider scopedServiceProvider, IOpenApiSchemaTransformer[] schemaTransformers, ApiParameterDescription? parameterDescription = null, CancellationToken cancellationToken = default)
+    internal async Task<OpenApiSchema> GetOrCreateUnresolvedSchemaAsync(OpenApiDocument? document, Type type, IServiceProvider scopedServiceProvider, IOpenApiSchemaTransformer[] schemaTransformers, ApiParameterDescription? parameterDescription = null, CancellationToken cancellationToken = default)
     {
         var key = parameterDescription?.ParameterDescriptor is IParameterInfoParameterDescriptor parameterInfoDescription
             && parameterDescription.ModelMetadata.PropertyName is null
@@ -133,7 +133,13 @@ internal async Task<IOpenApiSchema> GetOrCreateSchemaAsync(OpenApiDocument docum
         var deserializedSchema = JsonSerializer.Deserialize(schemaAsJsonObject, _jsonSchemaContext.OpenApiJsonSchema);
         Debug.Assert(deserializedSchema != null, "The schema should have been deserialized successfully and materialize a non-null value.");
         var schema = deserializedSchema.Schema;
-        await ApplySchemaTransformersAsync(schema, type, scopedServiceProvider, schemaTransformers, parameterDescription, cancellationToken);
+        await ApplySchemaTransformersAsync(document, schema, type, scopedServiceProvider, schemaTransformers, parameterDescription, cancellationToken);
+        return schema;
+    }
+
+    internal async Task<IOpenApiSchema> GetOrCreateSchemaAsync(OpenApiDocument document, Type type, IServiceProvider scopedServiceProvider, IOpenApiSchemaTransformer[] schemaTransformers, ApiParameterDescription? parameterDescription = null, CancellationToken cancellationToken = default)
+    {
+        var schema = await GetOrCreateUnresolvedSchemaAsync(document, type, scopedServiceProvider, schemaTransformers, parameterDescription, cancellationToken);
         return ResolveReferenceForSchema(document, schema);
     }
 
@@ -229,7 +235,7 @@ internal static IOpenApiSchema ResolveReferenceForSchema(OpenApiDocument documen
         return schema;
     }
 
-    internal async Task ApplySchemaTransformersAsync(IOpenApiSchema schema, Type type, IServiceProvider scopedServiceProvider, IOpenApiSchemaTransformer[] schemaTransformers, ApiParameterDescription? parameterDescription = null, CancellationToken cancellationToken = default)
+    internal async Task ApplySchemaTransformersAsync(OpenApiDocument? document, IOpenApiSchema schema, Type type, IServiceProvider scopedServiceProvider, IOpenApiSchemaTransformer[] schemaTransformers, ApiParameterDescription? parameterDescription = null, CancellationToken cancellationToken = default)
     {
         if (schemaTransformers.Length == 0)
         {
@@ -242,7 +248,9 @@ internal async Task ApplySchemaTransformersAsync(IOpenApiSchema schema, Type typ
             JsonTypeInfo = jsonTypeInfo,
             JsonPropertyInfo = null,
             ParameterDescription = parameterDescription,
-            ApplicationServices = scopedServiceProvider
+            ApplicationServices = scopedServiceProvider,
+            Document = document,
+            SchemaTransformers = schemaTransformers
         };
         for (var i = 0; i < schemaTransformers.Length; i++)
         {

src/OpenApi/src/Transformers/OpenApiDocumentTransformerContext.cs

@@ -1,7 +1,10 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Diagnostics;
 using Microsoft.AspNetCore.Mvc.ApiExplorer;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.OpenApi.Models;
 
 namespace Microsoft.AspNetCore.OpenApi;
 
@@ -24,4 +27,33 @@ public sealed class OpenApiDocumentTransformerContext
     /// Gets the application services associated with current document.
     /// </summary>
     public required IServiceProvider ApplicationServices { get; init; }
+
+    internal IOpenApiSchemaTransformer[] SchemaTransformers { get; init; } = [];
+
+    // Internal because we expect users to interact with the `Document` provided in
+    // the `IOpenApiDocumentTransformer` itself instead of the context object.
+    internal OpenApiDocument? Document { get; init; }
+
+    /// <summary>
+    /// Gets or creates an <see cref="OpenApiSchema"/> for the specified type. Augments
+    /// the schema with any <see cref="IOpenApiSchemaTransformer"/>s that are registered
+    /// on the document. If <paramref name="parameterDescription"/> is not null, the schema will be
+    /// augmented with the <see cref="ApiParameterDescription"/> information.
+    /// </summary>
+    /// <param name="type">The type for which the schema is being created.</param>
+    /// <param name="parameterDescription">An optional parameter description to augment the schema.</param>
+    /// <param name="cancellationToken">A cancellation token to cancel the operation.</param>
+    /// <returns>A task that represents the asynchronous operation, with a value of type <see cref="OpenApiSchema"/>.</returns>
+    public Task<OpenApiSchema> GetOrCreateSchemaAsync(Type type, ApiParameterDescription? parameterDescription = null, CancellationToken cancellationToken = default)
+    {
+        Debug.Assert(Document is not null, "Document should have been initialized by framework.");
+        var schemaService = ApplicationServices.GetRequiredKeyedService<OpenApiSchemaService>(DocumentName);
+        return schemaService.GetOrCreateUnresolvedSchemaAsync(
+            document: Document,
+            type: type,
+            parameterDescription: parameterDescription,
+            scopedServiceProvider: ApplicationServices,
+            schemaTransformers: SchemaTransformers,
+            cancellationToken: cancellationToken);
+    }
 }

src/OpenApi/src/Transformers/OpenApiOperationTransformerContext.cs

@@ -2,6 +2,8 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using Microsoft.AspNetCore.Mvc.ApiExplorer;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.OpenApi.Models;
 
 namespace Microsoft.AspNetCore.OpenApi;
 
@@ -24,4 +26,33 @@ public sealed class OpenApiOperationTransformerContext
     /// Gets the application services associated with the current document the target operation is in.
     /// </summary>
     public required IServiceProvider ApplicationServices { get; init; }
+
+    /// <summary>
+    /// Gets the OpenAPI document the current endpoint belongs to.
+    /// </summary>
+    public OpenApiDocument? Document { get; init; }
+
+    internal IOpenApiSchemaTransformer[] SchemaTransformers { get; init; } = [];
+
+    /// <summary>
+    /// Gets or creates an <see cref="OpenApiSchema"/> for the specified type. Augments
+    /// the schema with any <see cref="IOpenApiSchemaTransformer"/>s that are registered
+    /// on the document. If <paramref name="parameterDescription"/> is not null, the schema will be
+    /// augmented with the <see cref="ApiParameterDescription"/> information.
+    /// </summary>
+    /// <param name="type">The type for which the schema is being created.</param>
+    /// <param name="parameterDescription">An optional parameter description to augment the schema.</param>
+    /// <param name="cancellationToken">A cancellation token to cancel the operation.</param>
+    /// <returns>A task that represents the asynchronous operation, with a value of type <see cref="OpenApiSchema"/>.</returns>
+    public Task<OpenApiSchema> GetOrCreateSchemaAsync(Type type, ApiParameterDescription? parameterDescription = null, CancellationToken cancellationToken = default)
+    {
+        var schemaService = ApplicationServices.GetRequiredKeyedService<OpenApiSchemaService>(DocumentName);
+        return schemaService.GetOrCreateUnresolvedSchemaAsync(
+            document: Document,
+            type: type,
+            parameterDescription: parameterDescription,
+            scopedServiceProvider: ApplicationServices,
+            schemaTransformers: SchemaTransformers,
+            cancellationToken: cancellationToken);
+    }
 }

src/OpenApi/src/Transformers/OpenApiSchemaTransformerContext.cs

@@ -3,6 +3,8 @@
 
 using System.Text.Json.Serialization.Metadata;
 using Microsoft.AspNetCore.Mvc.ApiExplorer;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.OpenApi.Models;
 
 namespace Microsoft.AspNetCore.OpenApi;
 
@@ -41,11 +43,40 @@ public sealed class OpenApiSchemaTransformerContext
     /// </summary>
     public required IServiceProvider ApplicationServices { get; init; }
 
+    /// <summary>
+    /// Gets the OpenAPI document the current schema belongs to.
+    /// </summary>
+    public OpenApiDocument? Document { get; init; }
+
     // Expose internal setters for the properties that only allow initializations to avoid allocating
     // new instances of the context for each sub-schema transformation.
     internal void UpdateJsonTypeInfo(JsonTypeInfo jsonTypeInfo, JsonPropertyInfo? jsonPropertyInfo)
     {
         _jsonTypeInfo = jsonTypeInfo;
         _jsonPropertyInfo = jsonPropertyInfo;
     }
+
+    internal IOpenApiSchemaTransformer[] SchemaTransformers { get; init; } = [];
+
+    /// <summary>
+    /// Gets or creates an <see cref="OpenApiSchema"/> for the specified type. Augments
+    /// the schema with any <see cref="IOpenApiSchemaTransformer"/>s that are registered
+    /// on the document. If <paramref name="parameterDescription"/> is not null, the schema will be
+    /// augmented with the <see cref="ApiParameterDescription"/> information.
+    /// </summary>
+    /// <param name="type">The type for which the schema is being created.</param>
+    /// <param name="parameterDescription">An optional parameter description to augment the schema.</param>
+    /// <param name="cancellationToken">A cancellation token to cancel the operation.</param>
+    /// <returns>A task that represents the asynchronous operation, with a value of type <see cref="OpenApiSchema"/>.</returns>
+    public Task<OpenApiSchema> GetOrCreateSchemaAsync(Type type, ApiParameterDescription? parameterDescription = null, CancellationToken cancellationToken = default)
+    {
+        var schemaService = ApplicationServices.GetRequiredKeyedService<OpenApiSchemaService>(DocumentName);
+        return schemaService.GetOrCreateUnresolvedSchemaAsync(
+            document: Document,
+            type: type,
+            parameterDescription: parameterDescription,
+            scopedServiceProvider: ApplicationServices,
+            schemaTransformers: SchemaTransformers,
+            cancellationToken: cancellationToken);
+    }
 }