Apply schema transformers on properties and other subschemas

Author: captainsafia

src/OpenApi/perf/Microbenchmarks/TransformersBenchmark.cs

@@ -73,7 +73,7 @@ public void SchemaTransformer_Setup()
         {
             _options.UseSchemaTransformer((schema, context, token) =>
             {
-                if (context.Type == typeof(Todo) && context.ParameterDescription != null)
+                if (context.JsonTypeInfo.Type == typeof(Todo) && context.ParameterDescription != null)
                 {
                     schema.Extensions["x-my-extension"] = new OpenApiString(context.ParameterDescription.Name);
                 }

src/OpenApi/src/PublicAPI.Unshipped.txt

@@ -20,11 +20,13 @@ Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext.ApplicationServices
 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.JsonPropertyInfo.get -> System.Text.Json.Serialization.Metadata.JsonPropertyInfo?
+Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext.JsonPropertyInfo.init -> void
+Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext.JsonTypeInfo.get -> System.Text.Json.Serialization.Metadata.JsonTypeInfo!
+Microsoft.AspNetCore.OpenApi.OpenApiSchemaTransformerContext.JsonTypeInfo.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/Schemas/OpenApiSchemaService.cs

@@ -140,17 +140,58 @@ internal async Task<OpenApiSchema> GetOrCreateSchemaAsync(Type type, ApiParamete
 
     internal async Task ApplySchemaTransformersAsync(OpenApiSchema schema, Type type, ApiParameterDescription? parameterDescription = null, CancellationToken cancellationToken = default)
     {
+        var jsonTypeInfo = _jsonSerializerOptions.GetTypeInfo(type);
         var context = new OpenApiSchemaTransformerContext
         {
             DocumentName = documentName,
-            Type = type,
+            JsonTypeInfo = jsonTypeInfo,
+            JsonPropertyInfo = null,
             ParameterDescription = parameterDescription,
             ApplicationServices = serviceProvider
         };
         for (var i = 0; i < _openApiOptions.SchemaTransformers.Count; i++)
         {
             var transformer = _openApiOptions.SchemaTransformers[i];
-            await transformer(schema, context, cancellationToken);
+            await InnerApplySchemaTransformersAsync(schema, jsonTypeInfo, context, transformer, cancellationToken);
+        }
+    }
+
+    private async Task InnerApplySchemaTransformersAsync(OpenApiSchema schema,
+        JsonTypeInfo jsonTypeInfo,
+        OpenApiSchemaTransformerContext context,
+        Func<OpenApiSchema, OpenApiSchemaTransformerContext, CancellationToken, Task> transformer,
+        CancellationToken cancellationToken = default)
+    {
+        await transformer(schema, context, cancellationToken);
+
+        // Only apply transformers on polymorphic schemas where we can resolve the derived
+        // types associated with the base type.
+        if (schema.AnyOf is { Count: > 0 } && jsonTypeInfo.PolymorphismOptions is not null)
+        {
+            var anyOfIndex = 0;
+            foreach (var derivedType in jsonTypeInfo.PolymorphismOptions.DerivedTypes)
+            {
+                var derivedJsonTypeInfo = _jsonSerializerOptions.GetTypeInfo(derivedType.DerivedType);
+                context.UpdateJsonTypeInfo(derivedJsonTypeInfo, null);
+                await InnerApplySchemaTransformersAsync(schema.AnyOf[anyOfIndex], derivedJsonTypeInfo, context, transformer, cancellationToken);
+                anyOfIndex++;
+            }
+        }
+
+        if (schema.Items is not null)
+        {
+            var elementTypeInfo = _jsonSerializerOptions.GetTypeInfo(jsonTypeInfo.ElementType!);
+            context.UpdateJsonTypeInfo(elementTypeInfo, null);
+            await InnerApplySchemaTransformersAsync(schema.Items, elementTypeInfo, context, transformer, cancellationToken);
+        }
+
+        if (schema.Properties is { Count: > 0 })
+        {
+            foreach (var propertyInfo in jsonTypeInfo.Properties)
+            {
+                context.UpdateJsonTypeInfo(_jsonSerializerOptions.GetTypeInfo(propertyInfo.PropertyType), propertyInfo);
+                await InnerApplySchemaTransformersAsync(schema.Properties[propertyInfo.Name], jsonTypeInfo, context, transformer, cancellationToken);
+            }
         }
     }
 

src/OpenApi/src/Transformers/OpenApiSchemaTransformerContext.cs

@@ -1,8 +1,8 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Text.Json.Serialization.Metadata;
 using Microsoft.AspNetCore.Mvc.ApiExplorer;
-using Microsoft.OpenApi.Models;
 
 namespace Microsoft.AspNetCore.OpenApi;
 
@@ -11,24 +11,41 @@ namespace Microsoft.AspNetCore.OpenApi;
 /// </summary>
 public sealed class OpenApiSchemaTransformerContext
 {
+    private JsonTypeInfo? _jsonTypeInfo;
+    private JsonPropertyInfo? _jsonPropertyInfo;
+
     /// <summary>
     /// Gets the name of the associated OpenAPI document.
     /// </summary>
     public required string DocumentName { get; init; }
 
-    /// <summary>
-    /// Gets the <see cref="Type" /> associated with the current <see cref="OpenApiSchema"/>.
-    /// </summary>
-    public required Type Type { get; init; }
-
     /// <summary>
     /// Gets the <see cref="ApiParameterDescription"/> associated with the target schema.
     /// Null when processing an OpenAPI schema for a response type.
     /// </summary>
     public required ApiParameterDescription? ParameterDescription { get; init; }
 
+    /// <summary>
+    /// Gets the <see cref="JsonTypeInfo"/> associated with the target schema.
+    /// </summary>
+    public required JsonTypeInfo JsonTypeInfo { get => _jsonTypeInfo!; init => _jsonTypeInfo = value; }
+
+    /// <summary>
+    /// Gets the <see cref="JsonPropertyInfo"/> associated with the target schema if the
+    /// target schema is a property of a parent schema.
+    /// </summary>
+    public required JsonPropertyInfo? JsonPropertyInfo { get => _jsonPropertyInfo; init => _jsonPropertyInfo = value; }
+
     /// <summary>
     /// Gets the application services associated with the current document the target schema is in.
     /// </summary>
     public required IServiceProvider ApplicationServices { 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;
+    }
 }

src/OpenApi/test/Transformers/Implementations/OpenApiSchemaReferenceTransformerTests.cs

@@ -269,7 +269,7 @@ public async Task TypeModifiedWithSchemaTransformerMapsToDifferentReferenceId()
         var options = new OpenApiOptions();
         options.UseSchemaTransformer((schema, context, cancellationToken) =>
         {
-            if (context.Type == typeof(Todo) && context.ParameterDescription is not null)
+            if (context.JsonTypeInfo.Type == typeof(Todo) && context.ParameterDescription is not null)
             {
                 schema.Extensions["x-my-extension"] = new OpenApiString(context.ParameterDescription.Name);
             }