Allow document transformers to access operation descriptions (#61264)

src/OpenApi/src/Services/OpenApiDocumentService.cs

@@ -72,10 +72,11 @@ public async Task<OpenApiDocument> GetOpenApiDocumentAsync(IServiceProvider scop
             Info = GetOpenApiInfo(),
             Servers = GetOpenApiServers(httpRequest)
         };
-        document.Paths = await GetOpenApiPathsAsync(document, scopedServiceProvider, operationTransformers, schemaTransformers, cancellationToken);
+        var operationDescriptions = new Dictionary<OpenApiOperation, ApiDescription>();
+        document.Paths = await GetOpenApiPathsAsync(document, scopedServiceProvider, operationTransformers, schemaTransformers, operationDescriptions, cancellationToken);
         try
         {
-            await ApplyTransformersAsync(document, scopedServiceProvider, schemaTransformers, cancellationToken);
+            await ApplyTransformersAsync(document, scopedServiceProvider, schemaTransformers, operationDescriptions, cancellationToken);
         }
 
         finally
@@ -97,15 +98,16 @@ public async Task<OpenApiDocument> GetOpenApiDocumentAsync(IServiceProvider scop
         return document;
     }
 
-    private async Task ApplyTransformersAsync(OpenApiDocument document, IServiceProvider scopedServiceProvider, IOpenApiSchemaTransformer[] schemaTransformers, CancellationToken cancellationToken)
+    private async Task ApplyTransformersAsync(OpenApiDocument document, IServiceProvider scopedServiceProvider, IOpenApiSchemaTransformer[] schemaTransformers, IReadOnlyDictionary<OpenApiOperation, ApiDescription> operationDescriptions, CancellationToken cancellationToken)
     {
         var documentTransformerContext = new OpenApiDocumentTransformerContext
         {
             DocumentName = documentName,
             ApplicationServices = scopedServiceProvider,
             DescriptionGroups = apiDescriptionGroupCollectionProvider.ApiDescriptionGroups.Items,
             Document = document,
-            SchemaTransformers = schemaTransformers
+            SchemaTransformers = schemaTransformers,
+            OperationDescriptions = operationDescriptions
         };
         // Use index-based for loop to avoid allocating an enumerator with a foreach.
         for (var i = 0; i < _options.DocumentTransformers.Count; i++)
@@ -246,6 +248,7 @@ private async Task<OpenApiPaths> GetOpenApiPathsAsync(
         IServiceProvider scopedServiceProvider,
         IOpenApiOperationTransformer[] operationTransformers,
         IOpenApiSchemaTransformer[] schemaTransformers,
+        Dictionary<OpenApiOperation, ApiDescription> operationDescriptions,
         CancellationToken cancellationToken)
     {
         var descriptionsByPath = apiDescriptionGroupCollectionProvider.ApiDescriptionGroups.Items
@@ -256,7 +259,7 @@ private async Task<OpenApiPaths> GetOpenApiPathsAsync(
         foreach (var descriptions in descriptionsByPath)
         {
             Debug.Assert(descriptions.Key != null, "Relative path mapped to OpenApiPath key cannot be null.");
-            var operations = await GetOperationsAsync(descriptions, document, scopedServiceProvider, operationTransformers, schemaTransformers, cancellationToken);
+            var operations = await GetOperationsAsync(descriptions, document, scopedServiceProvider, operationTransformers, schemaTransformers, operationDescriptions, cancellationToken);
             if (operations.Count > 0)
             {
                 paths.Add(descriptions.Key, new OpenApiPathItem { Operations = operations });
@@ -272,6 +275,7 @@ private async Task<Dictionary<HttpMethod, OpenApiOperation>> GetOperationsAsync(
         IServiceProvider scopedServiceProvider,
         IOpenApiOperationTransformer[] operationTransformers,
         IOpenApiSchemaTransformer[] schemaTransformers,
+        Dictionary<OpenApiOperation, ApiDescription> operationDescriptions,
         CancellationToken cancellationToken)
     {
         var operations = new Dictionary<HttpMethod, OpenApiOperation>();
@@ -299,6 +303,7 @@ private async Task<Dictionary<HttpMethod, OpenApiOperation>> GetOperationsAsync(
             }
 
             operations[method] = operation;
+            operationDescriptions.Add(operation, description);
 
             // Use index-based for loop to avoid allocating an enumerator with a foreach.
             for (var i = 0; i < operationTransformers.Length; i++)

src/OpenApi/src/Transformers/OpenApiDocumentTransformerContext.cs

@@ -27,6 +27,11 @@
     /// </summary>
     public required IServiceProvider ApplicationServices { get; init; }
 
+    /// <summary>
+    /// Map of <see cref="OpenApiOperation"/> instances to their associated <see cref="ApiDescription"/>.
+    /// </summary>
+    public required IReadOnlyDictionary<OpenApiOperation, ApiDescription> OperationDescriptions { get; init; }
+
     internal IOpenApiSchemaTransformer[] SchemaTransformers { get; init; } = [];
 
     // Internal because we expect users to interact with the `Document` provided in

src/OpenApi/test/Microsoft.AspNetCore.OpenApi.Tests/Transformers/DocumentTransformerTests.cs

@@ -389,4 +389,89 @@ internal void TestMethod() { }
 
         internal static int InstantiationCount = 0;
     }
+
+    [Fact]
+    public async Task DocumentTransformer_CanAccessOperationDescriptions()
+    {
+        var builder = CreateBuilder();
+
+        builder.MapGet("/todo", () => "todo");
+        builder.MapGet("/user", () => "user");
+        builder.MapPost("/todo", () => { });
+
+        var expectedPaths = new HashSet<string>(StringComparer.Ordinal)
+        {
+            "/todo",
+            "/user"
+        };
+
+        var mappedOperationsCount = 0;
+        var options = new OpenApiOptions();
+        options.AddDocumentTransformer((document, context, cancellationToken) =>
+        {
+            Assert.Equal(3, context.OperationDescriptions.Count);
+
+            foreach (var (operation, apiDescription) in context.OperationDescriptions)
+            {
+                var documentPath = "/" + apiDescription.RelativePath;
+                Assert.Contains(documentPath, expectedPaths);
+
+                Assert.True(document.Paths.TryGetValue(documentPath, out var pathItem));
+                Assert.Contains(operation, pathItem.Operations.Values);
+                mappedOperationsCount++;
+            }
+
+            return Task.CompletedTask;
+        });
+
+        await VerifyOpenApiDocument(builder, options, document =>
+        {
+            Assert.Equal(2, document.Paths.Count);
+        });
+
+        Assert.Equal(3, mappedOperationsCount);
+    }
+
+    [Fact]
+    public async Task DocumentTransformer_CanAccessEndpointMetadataViaOperationDescriptions()
+    {
+        var builder = CreateBuilder();
+
+        builder.MapGet("/public", () => "public")
+            .WithMetadata(new OperationMarker());
+        builder.MapGet("/secret", () => "secret");
+
+        var options = new OpenApiOptions();
+        var publicEndpointHasMarker = false;
+        var secretEndpointHasMarker = false;
+
+        options.AddDocumentTransformer((document, context, cancellationToken) =>
+        {
+            Assert.Equal(2, context.OperationDescriptions.Count);
+
+            foreach (var (_, apiDescription) in context.OperationDescriptions)
+            {
+                var hasMarker = apiDescription.ActionDescriptor.EndpointMetadata
+                    .Any(m => m is OperationMarker);
+
+                if (apiDescription.RelativePath == "public")
+                {
+                    publicEndpointHasMarker = hasMarker;
+                }
+                else if (apiDescription.RelativePath == "secret")
+                {
+                    secretEndpointHasMarker = hasMarker;
+                }
+            }
+
+            return Task.CompletedTask;
+        });
+
+        await VerifyOpenApiDocument(builder, options, document => { });
+
+        Assert.True(publicEndpointHasMarker);
+        Assert.False(secretEndpointHasMarker);
+    }
+
+    private sealed class OperationMarker;
 }