Apply schema transformers on properties and other subschemas (dotnet#56709)

* Apply schema transformers on properties and other subschemas

* Add test for tranformers on unsupported subschemas

* Avoid aggressive disposable for type-based schema transformers

* Add some guards around recursion into properties and anyOf

* Manage activated IOpenApiSchemaTransformer in OpenApiSchemaService

Author: captainsafia

src/OpenApi/perf/Microbenchmarks/TransformersBenchmark.cs

@@ -95,7 +95,7 @@ public void SchemaTransformer_Setup()
         {
             _options.AddSchemaTransformer((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);
                 }
@@ -167,7 +167,7 @@ private class SchemaTransformer : IOpenApiSchemaTransformer
     {
         public Task TransformAsync(OpenApiSchema schema, OpenApiSchemaTransformerContext context, CancellationToken cancellationToken)
         {
-            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

@@ -30,11 +30,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.AspNetCore.OpenApi.OpenApiOptions.CreateDefaultSchemaReferenceId(System.Text.Json.Serialization.Metadata.JsonTypeInfo! jsonTypeInfo) -> string?

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

@@ -144,17 +144,83 @@ 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.TransformAsync(schema, context, cancellationToken);
+            // If the transformer is a type-based transformer, we need to initialize and finalize it
+            // once in the context of the top-level assembly and not the child properties we are invoking
+            // it on.
+            if (transformer is TypeBasedOpenApiSchemaTransformer typeBasedTransformer)
+            {
+                var initializedTransformer = typeBasedTransformer.InitializeTransformer(serviceProvider);
+                try
+                {
+                    await InnerApplySchemaTransformersAsync(schema, jsonTypeInfo, context, initializedTransformer, cancellationToken);
+                }
+                finally
+                {
+                    await TypeBasedOpenApiSchemaTransformer.FinalizeTransformer(initializedTransformer);
+                }
+            }
+            else
+            {
+                await InnerApplySchemaTransformersAsync(schema, jsonTypeInfo, context, transformer, cancellationToken);
+            }
+        }
+    }
+
+    private async Task InnerApplySchemaTransformersAsync(OpenApiSchema schema,
+        JsonTypeInfo jsonTypeInfo,
+        OpenApiSchemaTransformerContext context,
+        IOpenApiSchemaTransformer transformer,
+        CancellationToken cancellationToken = default)
+    {
+        await transformer.TransformAsync(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);
+                if (schema.AnyOf.Count <= anyOfIndex)
+                {
+                    break;
+                }
+                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);
+                if (schema.Properties.TryGetValue(propertyInfo.Name, out var propertySchema))
+                {
+                    await InnerApplySchemaTransformersAsync(propertySchema, _jsonSerializerOptions.GetTypeInfo(propertyInfo.PropertyType), 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/src/Transformers/TypeBasedOpenApiSchemaTransformer.cs

@@ -20,24 +20,26 @@ internal TypeBasedOpenApiSchemaTransformer([DynamicallyAccessedMembers(Dynamical
         _transformerFactory = ActivatorUtilities.CreateFactory(_transformerType, []);
     }
 
-    public async Task TransformAsync(OpenApiSchema schema, OpenApiSchemaTransformerContext context, CancellationToken cancellationToken)
+    internal IOpenApiSchemaTransformer InitializeTransformer(IServiceProvider serviceProvider)
     {
-        var transformer = _transformerFactory.Invoke(context.ApplicationServices, []) as IOpenApiSchemaTransformer;
+        var transformer = _transformerFactory.Invoke(serviceProvider, []) as IOpenApiSchemaTransformer;
         Debug.Assert(transformer != null, $"The type {_transformerType} does not implement {nameof(IOpenApiSchemaTransformer)}.");
-        try
+        return transformer;
+    }
+
+    internal static async Task FinalizeTransformer(IOpenApiSchemaTransformer transformer)
+    {
+        if (transformer is IAsyncDisposable asyncDisposable)
         {
-            await transformer.TransformAsync(schema, context, cancellationToken);
+            await asyncDisposable.DisposeAsync();
         }
-        finally
+        else if (transformer is IDisposable disposable)
         {
-            if (transformer is IAsyncDisposable asyncDisposable)
-            {
-                await asyncDisposable.DisposeAsync();
-            }
-            else if (transformer is IDisposable disposable)
-            {
-                disposable.Dispose();
-            }
+            disposable.Dispose();
         }
     }
+
+    // No-op because the activate instance is invoked by the OpenApiSchema service.
+    public Task TransformAsync(OpenApiSchema schema, OpenApiSchemaTransformerContext context, CancellationToken cancellationToken)
+        => Task.CompletedTask;
 }

src/OpenApi/test/Integration/snapshots/OpenApiDocumentIntegrationTests.VerifyOpenApiDocument_documentName=v2.verified.txt

@@ -52,7 +52,11 @@
       "ArrayOfstring": {
         "type": "array",
         "items": {
-          "type": "string"
+          "type": "string",
+          "externalDocs": {
+            "description": "Documentation for this OpenAPI schema",
+            "url": "https://example.com/api/docs/schemas/string"
+          }
         },
         "externalDocs": {
           "description": "Documentation for this OpenAPI schema",

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

@@ -269,7 +269,7 @@ public async Task TypeModifiedWithSchemaTransformerMapsToDifferentReferenceId()
         var options = new OpenApiOptions();
         options.AddSchemaTransformer((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);
             }