NIFI-16235 - Support rebasing locally added Controller Services in versioned Process Groups

nifi-registry/nifi-registry-core/nifi-registry-flow-diff/src/main/java/org/apache/nifi/registry/flow/diff/ComponentAddedRebaseHandler.java

@@ -0,0 +1,109 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.nifi.registry.flow.diff;
+
+import org.apache.nifi.flow.VersionedComponent;
+import org.apache.nifi.flow.VersionedControllerService;
+import org.apache.nifi.flow.VersionedProcessGroup;
+
+import java.util.HashSet;
+import java.util.Set;
+
+public class ComponentAddedRebaseHandler implements RebaseHandler {
+
+    private static final String NULL_COMPONENT_TYPE = "null";
+
+    @Override
+    public DifferenceType getSupportedType() {
+        return DifferenceType.COMPONENT_ADDED;
+    }
+
+    @Override
+    public RebaseAnalysis.ClassifiedDifference classify(final FlowDifference localDifference, final Set<FlowDifference> upstreamDifferences,
+                                                        final VersionedProcessGroup targetSnapshot) {
+        final VersionedComponent addedComponent = localDifference.getComponentB();
+        if (!(addedComponent instanceof VersionedControllerService controllerService)) {
+            final String componentType = addedComponent == null ? NULL_COMPONENT_TYPE : addedComponent.getClass().getSimpleName();
+            return RebaseAnalysis.ClassifiedDifference.unsupported(localDifference, RebaseConflictCode.UNSUPPORTED_COMPONENT_TYPE,
+                    "Local component addition type %s is not supported for rebase".formatted(componentType));
+        }
+
+        final String parentGroupIdentifier = controllerService.getGroupIdentifier();
+        if (parentGroupIdentifier == null) {
+            return RebaseAnalysis.ClassifiedDifference.unsupported(localDifference, RebaseConflictCode.COMPONENT_NOT_FOUND,
+                    "Controller Service %s does not specify a parent Process Group".formatted(controllerService.getIdentifier()));
+        }
+
+        final VersionedProcessGroup parentGroup = resolveParentGroup(targetSnapshot, parentGroupIdentifier, upstreamDifferences);
+        if (parentGroup == null) {
+            return RebaseAnalysis.ClassifiedDifference.unsupported(localDifference, RebaseConflictCode.COMPONENT_NOT_FOUND,
+                    "Parent Process Group %s for Controller Service %s not found in target snapshot"
+                            .formatted(parentGroupIdentifier, controllerService.getIdentifier()));
+        }
+
+        final VersionedComponent collidingComponent = RebaseHandlerUtils.findComponentById(targetSnapshot, controllerService.getIdentifier());
+        if (collidingComponent != null) {
+            return RebaseAnalysis.ClassifiedDifference.conflicting(localDifference, RebaseConflictCode.COMPONENT_IDENTIFIER_COLLISION,
+                    "Target snapshot already contains component %s with identifier %s"
+                            .formatted(collidingComponent.getClass().getSimpleName(), controllerService.getIdentifier()));
+        }
+
+        controllerService.setGroupIdentifier(parentGroup.getIdentifier());
+        return RebaseAnalysis.ClassifiedDifference.compatible(localDifference);
+    }
+
+    @Override
+    public void apply(final FlowDifference localDifference, final VersionedProcessGroup mergedFlow) {
+        final VersionedControllerService controllerService = (VersionedControllerService) localDifference.getComponentB();
+        final VersionedProcessGroup parentGroup = RebaseHandlerUtils.findProcessGroupById(mergedFlow, controllerService.getGroupIdentifier());
+        if (parentGroup == null) {
+            throw new IllegalStateException("Parent Process Group %s for Controller Service %s was verified during classification but is absent during apply"
+                    .formatted(controllerService.getGroupIdentifier(), controllerService.getIdentifier()));
+        }
+
+        final VersionedComponent existingComponent = RebaseHandlerUtils.findComponentById(mergedFlow, controllerService.getIdentifier());
+        if (existingComponent != null) {
+            throw new IllegalStateException("Merged flow already contains component %s with identifier %s"
+                    .formatted(existingComponent.getClass().getSimpleName(), controllerService.getIdentifier()));
+        }
+
+        final Set<VersionedControllerService> controllerServices = parentGroup.getControllerServices();
+        if (controllerServices == null) {
+            parentGroup.setControllerServices(new HashSet<>());
+        }
+        parentGroup.getControllerServices().add(controllerService);
+    }
+
+    private VersionedProcessGroup resolveParentGroup(final VersionedProcessGroup targetSnapshot, final String parentGroupIdentifier,
+                                                     final Set<FlowDifference> upstreamDifferences) {
+        final VersionedProcessGroup parentGroup = RebaseHandlerUtils.findProcessGroupById(targetSnapshot, parentGroupIdentifier);
+        if (parentGroup != null) {
+            return parentGroup;
+        }
+
+        final boolean parentRemoved = upstreamDifferences.stream()
+                .filter(difference -> difference.getDifferenceType() == DifferenceType.COMPONENT_REMOVED)
+                .map(FlowDifference::getComponentA)
+                .filter(VersionedProcessGroup.class::isInstance)
+                .map(VersionedComponent::getIdentifier)
+                .anyMatch(parentGroupIdentifier::equals);
+
+        return parentRemoved ? null : targetSnapshot;
+    }
+
+}

nifi-registry/nifi-registry-core/nifi-registry-flow-diff/src/main/java/org/apache/nifi/registry/flow/diff/RebaseConflictCode.java

@@ -32,6 +32,11 @@ public enum RebaseConflictCode {
      */
     MISSING_FIELD_NAME,
 
+    /**
+     * The registered handler does not support the local change's component type.
+     */
+    UNSUPPORTED_COMPONENT_TYPE,
+
     /**
      * Both the local and upstream flows modified the same property on the same component.
      */
@@ -47,6 +52,11 @@ public enum RebaseConflictCode {
      */
     COMPONENT_NOT_FOUND,
 
+    /**
+     * The target version already contains a component with the same identifier as the local addition.
+     */
+    COMPONENT_IDENTIFIER_COLLISION,
+
     /**
      * The property descriptor targeted by the local change changed in an incompatible way in the target version.
      */

nifi-registry/nifi-registry-core/nifi-registry-flow-diff/src/main/java/org/apache/nifi/registry/flow/diff/RebaseHandlerUtils.java

@@ -89,6 +89,21 @@ static VersionedComponent findComponentById(final VersionedProcessGroup group, f
         return null;
     }
 
+    static VersionedProcessGroup findProcessGroupById(final VersionedProcessGroup group, final String identifier) {
+        if (identifier.equals(group.getIdentifier()) || identifier.equals(group.getInstanceIdentifier())) {
+            return group;
+        }
+
+        for (final VersionedProcessGroup childGroup : group.getProcessGroups()) {
+            final VersionedProcessGroup result = findProcessGroupById(childGroup, identifier);
+            if (result != null) {
+                return result;
+            }
+        }
+
+        return null;
+    }
+
     static VersionedConnection findConnectionById(final VersionedProcessGroup group, final String identifier) {
         for (final VersionedConnection connection : group.getConnections()) {
             if (identifier.equals(connection.getIdentifier())) {

nifi-registry/nifi-registry-core/nifi-registry-flow-diff/src/main/java/org/apache/nifi/registry/flow/diff/StandardRebaseEngine.java

@@ -46,6 +46,7 @@ public StandardRebaseEngine() {
         registerHandler(new PositionChangedRebaseHandler());
         registerHandler(new SizeChangedRebaseHandler());
         registerHandler(new BendpointsChangedRebaseHandler());
+        registerHandler(new ComponentAddedRebaseHandler());
         registerHandler(new PropertyChangedRebaseHandler());
         registerHandler(new PropertyAddedRebaseHandler());
         registerHandler(new CommentsChangedRebaseHandler());