WFCORE-6948 Add ability for wildfly-subsystem to register an operatio…

…n transformation for any specific operation (not just add).

subsystem/src/main/java/org/wildfly/subsystem/resource/ResourceDescriptor.java

@@ -81,8 +81,19 @@ default Set<ResourceCapabilityReferenceRecorder<?>> getResourceCapabilityReferen
      * Returns a transformer to be applied to all operations that operate on an existing resource.
      * This is typically used to adapt legacy operations to conform to the current version of the model.
      * @return an operation handler transformer.
+     * @deprecated Superseded by {@link #getOperationTransformation(String)}.
      */
+    @Deprecated(forRemoval = true)
     default UnaryOperator<OperationStepHandler> getResourceOperationTransformation() {
+        return this.getOperationTransformation(""); // Returns default operation transformer
+    }
+
+    /**
+     * Returns a transformer to be applied the specified operation.
+     * This is typically used to adapt legacy operations to conform to the current version of the model.
+     * @return an operation handler transformer.
+     */
+    default UnaryOperator<OperationStepHandler> getOperationTransformation(String operationName) {
         return UnaryOperator.identity();
     }
 
@@ -114,9 +125,11 @@ default Set<RuntimeCapability<?>> getCapabilities() {
      * Returns a transformer for the add operation handler.
      * This is typically used to adapt legacy operations to conform to the current version of the model.
      * @return an operation handler transformer.
+     * @deprecated Superseded by {@link #getOperationTransformation(String)}
      */
+    @Deprecated(forRemoval = true)
     default UnaryOperator<OperationStepHandler> getAddOperationTransformation() {
-        return UnaryOperator.identity();
+        return this.getOperationTransformation(ModelDescriptionConstants.ADD);
     }
 
     /**
@@ -150,8 +163,8 @@ class DefaultResourceDescriptor implements ResourceDescriptor {
         private final Set<PathElement> requiredSingletonChildren;
         private final Map<AttributeDefinition, AttributeTranslation> attributeTranslations;
         private final Set<ResourceCapabilityReferenceRecorder<?>> resourceCapabilityReferences;
-        private final UnaryOperator<OperationStepHandler> addOperationTransformer;
-        private final UnaryOperator<OperationStepHandler> operationTransformer;
+        private final Map<String, UnaryOperator<OperationStepHandler>> operationTransformers;
+        private final UnaryOperator<OperationStepHandler> defaultOperationTransformer;
         private final UnaryOperator<Resource> resourceTransformer;
         private final Optional<Consumer<DeploymentProcessorTarget>> deploymentChainContributor;
         private final OperationEntry.Flag addOperationRestartFlag;
@@ -187,8 +200,8 @@ public BiPredicate<OperationContext, Resource> getCapabilityFilter(RuntimeCapabi
             this.requiredChildren = builder.requiredChildren;
             this.requiredSingletonChildren = builder.requiredSingletonChildren;
             this.resourceCapabilityReferences = builder.resourceCapabilityReferences;
-            this.addOperationTransformer = builder.addOperationTransformer;
-            this.operationTransformer = builder.operationTransformer;
+            this.operationTransformers = builder.operationTransformers;
+            this.defaultOperationTransformer = builder.defaultOperationTransformer;
             this.resourceTransformer = builder.resourceTransformer;
             this.deploymentChainContributor = builder.deploymentChainContributor;
             this.addOperationRestartFlag = builder.addOperationRestartFlag;
@@ -221,8 +234,8 @@ public Set<ResourceCapabilityReferenceRecorder<?>> getResourceCapabilityReferenc
         }
 
         @Override
-        public UnaryOperator<OperationStepHandler> getResourceOperationTransformation() {
-            return this.operationTransformer;
+        public UnaryOperator<OperationStepHandler> getOperationTransformation(String operationName) {
+            return this.operationTransformers.getOrDefault(operationName, this.defaultOperationTransformer);
         }
 
         @Override
@@ -250,11 +263,6 @@ public Set<PathElement> getRequiredSingletonChildren() {
             return this.requiredSingletonChildren;
         }
 
-        @Override
-        public UnaryOperator<OperationStepHandler> getAddOperationTransformation() {
-            return this.addOperationTransformer;
-        }
-
         @Override
         public UnaryOperator<Resource> getResourceTransformation() {
             return this.resourceTransformer;
@@ -435,16 +443,36 @@ default C addResourceCapabilityReference(ResourceCapabilityReferenceRecorder<?>
          * Applies the specified transformation to the {@value ModelDescriptionConstants#ADD} operation of this resource.
          * @param transformation an operation handler transformation
          * @return a reference to this configurator
+         * @deprecated Superseded by {@link #withOperationTransformation(String, UnaryOperator)}.
          */
-        C withAddResourceOperationTransformation(UnaryOperator<OperationStepHandler> transformation);
+        @Deprecated(forRemoval = true)
+        default C withAddResourceOperationTransformation(UnaryOperator<OperationStepHandler> transformation) {
+            return this.withOperationTransformation(ModelDescriptionConstants.ADD, transformation);
+        }
 
         /**
-         * Applies the specified transformation to the {@value ModelDescriptionConstants#REMOVE} and all global operations of this resource.
+         * Applies the specified transformation to all operations for this resource.
          * @param transformation an operation handler transformation
          * @return a reference to this configurator
          */
         C withOperationTransformation(UnaryOperator<OperationStepHandler> transformation);
 
+        /**
+         * Applies the specified transformation to the specified operation for this resource.
+         * @param transformation an operation handler transformation
+         * @return a reference to this configurator
+         */
+        default C withOperationTransformation(String operationName, UnaryOperator<OperationStepHandler> transformation) {
+            return this.withOperationTransformation(Set.of(operationName), transformation);
+        }
+
+        /**
+         * Applies the specified transformation to the specified operations for this resource.
+         * @param transformation an operation handler transformation
+         * @return a reference to this configurator
+         */
+        C withOperationTransformation(Set<String> operationNames, UnaryOperator<OperationStepHandler> transformation);
+
         /**
          * Applies the specified transformation to the {@link Resource} created by this resource's {@value ModelDescriptionConstants#ADD} operation.
          * @param transformation an operation handler transformation
@@ -545,8 +573,8 @@ abstract static class AbstractConfigurator<C extends Configurator<C>> implements
         private Set<PathElement> requiredSingletonChildren = Set.of();
         private Map<AttributeDefinition, AttributeTranslation> attributeTranslations = Map.of();
         private Set<ResourceCapabilityReferenceRecorder<?>> resourceCapabilityReferences = Set.of();
-        private UnaryOperator<OperationStepHandler> addOperationTransformer = UnaryOperator.identity();
-        private UnaryOperator<OperationStepHandler> operationTransformer = UnaryOperator.identity();
+        private Map<String, UnaryOperator<OperationStepHandler>> operationTransformers = Map.of();
+        private UnaryOperator<OperationStepHandler> defaultOperationTransformer = UnaryOperator.identity();
         private UnaryOperator<Resource> resourceTransformer = UnaryOperator.identity();
         private Optional<Consumer<DeploymentProcessorTarget>> deploymentChainContributor = Optional.empty();
 
@@ -645,14 +673,14 @@ public C addResourceCapabilityReferences(Collection<ResourceCapabilityReferenceR
         }
 
         @Override
-        public C withAddResourceOperationTransformation(UnaryOperator<OperationStepHandler> transformation) {
-            this.addOperationTransformer = transformation;
+        public C withOperationTransformation(UnaryOperator<OperationStepHandler> transformation) {
+            this.defaultOperationTransformer = transformation;
             return this.self();
         }
 
         @Override
-        public C withOperationTransformation(UnaryOperator<OperationStepHandler> transformation) {
-            this.operationTransformer = transformation;
+        public C withOperationTransformation(Set<String> operationNames, UnaryOperator<OperationStepHandler> transformation) {
+            this.operationTransformers = concat(this.operationTransformers, operationNames.stream(), transformation);
             return this.self();
         }
 

subsystem/src/main/java/org/wildfly/subsystem/resource/ResourceDescriptorRegistrar.java

@@ -135,23 +135,24 @@ public void register(ManagementResourceRegistration registration) {
                 .setStability(registration.getStability())
                 .withFlag(this.descriptor.getAddOperationRestartFlag())
                 .build();
-        registration.registerOperationHandler(addDefinition, this.descriptor.getAddOperationTransformation().apply(new AddResourceOperationStepHandler(this.descriptor)));
+        registration.registerOperationHandler(addDefinition, this.descriptor.getOperationTransformation(ModelDescriptionConstants.ADD).apply(new AddResourceOperationStepHandler(this.descriptor)));
 
         // Register remove resource operation handler
         OperationDefinition removeDefinition = new SimpleOperationDefinitionBuilder(ModelDescriptionConstants.REMOVE, this.descriptor.getResourceDescriptionResolver())
                 .setDescriptionProvider(new DefaultResourceRemoveDescriptionProvider(this.descriptor.getResourceDescriptionResolver()))
                 .setStability(registration.getStability())
                 .withFlag(this.descriptor.getRemoveOperationRestartFlag())
                 .build();
-        registration.registerOperationHandler(removeDefinition, this.descriptor.getResourceOperationTransformation().apply(new RemoveResourceOperationStepHandler(this.descriptor)));
+        registration.registerOperationHandler(removeDefinition, this.descriptor.getOperationTransformation(ModelDescriptionConstants.REMOVE).apply(new RemoveResourceOperationStepHandler(this.descriptor)));
 
         // Override global operations with transformed operations, if necessary
         for (Map.Entry<OperationDefinition, OperationStepHandler> entry : GLOBAL_OPERATIONS.entrySet()) {
+            OperationDefinition definition = entry.getKey();
             OperationStepHandler handler = entry.getValue();
             // Only override global operation handlers for non-identity transformations
-            OperationStepHandler transformedHandler = this.descriptor.getResourceOperationTransformation().apply(handler);
+            OperationStepHandler transformedHandler = this.descriptor.getOperationTransformation(definition.getName()).apply(handler);
             if (handler != transformedHandler) {
-                registration.registerOperationHandler(entry.getKey(), transformedHandler);
+                registration.registerOperationHandler(definition, transformedHandler);
             }
         }
     }