From df2147829742129c039b37c5d6f4f11aa54785a2 Mon Sep 17 00:00:00 2001
From: Bence Simon <bsimon@apache.org>
Date: Thu, 24 Nov 2022 21:41:50 +0100
Subject: [PATCH] NIFI-10874 Fixed issue with multiple levels of versioned
 flows

This closes #6719

Signed-off-by: David Handermann <exceptionfactory@apache.org>
---
 .../org/apache/nifi/web/api/ProcessGroupResource.java | 11 +++++++++--
 .../java/org/apache/nifi/web/api/dto/DtoFactory.java  |  2 ++
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/ProcessGroupResource.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/ProcessGroupResource.java
index 6704ed5678..dee9a7e9bd 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/ProcessGroupResource.java
+++ b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/ProcessGroupResource.java
@@ -51,6 +51,7 @@ import org.apache.nifi.connectable.ConnectableType;
 import org.apache.nifi.controller.ScheduledState;
 import org.apache.nifi.controller.serialization.FlowEncodingVersion;
 import org.apache.nifi.controller.service.ControllerServiceState;
+import org.apache.nifi.flow.VersionedFlowCoordinates;
 import org.apache.nifi.flow.VersionedParameterContext;
 import org.apache.nifi.flow.VersionedProcessGroup;
 import org.apache.nifi.parameter.ParameterContext;
@@ -2013,7 +2014,13 @@ public class ProcessGroupResource extends FlowUpdateResource<ProcessGroupImportE
             // Step 2: Retrieve flow from Flow Registry
             final RegisteredFlowSnapshot flowSnapshot = getFlowFromRegistry(versionControlInfo);
 
-            // Step 3: Resolve Bundle info
+            // Step 3: Enrich version control info came from UI
+            if (flowSnapshot.getFlowContents() != null && flowSnapshot.getFlowContents().getVersionedFlowCoordinates() != null) {
+                final VersionedFlowCoordinates versionedFlowCoordinates = flowSnapshot.getFlowContents().getVersionedFlowCoordinates();
+                flowSnapshot.getFlowContents().getVersionedFlowCoordinates().setStorageLocation(versionedFlowCoordinates.getStorageLocation());
+            }
+
+            // Step 4: Resolve Bundle info
             serviceFacade.discoverCompatibleBundles(flowSnapshot.getFlowContents());
 
             // If there are any Controller Services referenced that are inherited from the parent group, resolve those to point to the appropriate Controller Service, if we are able to.
@@ -2022,7 +2029,7 @@ public class ProcessGroupResource extends FlowUpdateResource<ProcessGroupImportE
             // If there are any Parameter Providers referenced by Parameter Contexts, resolve these to point to the appropriate Parameter Provider, if we are able to.
             serviceFacade.resolveParameterProviders(flowSnapshot, NiFiUserUtils.getNiFiUser());
 
-            // Step 4: Update contents of the ProcessGroupDTO passed in to include the components that need to be added.
+            // Step 5: Update contents of the ProcessGroupDTO passed in to include the components that need to be added.
             requestProcessGroupEntity.setVersionedFlowSnapshot(flowSnapshot);
         }
 
diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/dto/DtoFactory.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/dto/DtoFactory.java
index c2421b54f4..d28f9d93a2 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/dto/DtoFactory.java
+++ b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/dto/DtoFactory.java
@@ -2810,6 +2810,7 @@ public final class DtoFactory {
         dto.setFlowName(versionControlInfo.getFlowName());
         dto.setFlowDescription(versionControlInfo.getFlowDescription());
         dto.setVersion(versionControlInfo.getVersion());
+        dto.setStorageLocation(versionControlInfo.getStorageLocation());
 
         final VersionedFlowStatus status = versionControlInfo.getStatus();
         final VersionedFlowState state = status.getState();
@@ -4594,6 +4595,7 @@ public final class DtoFactory {
         copy.setVersion(original.getVersion());
         copy.setState(original.getState());
         copy.setStateExplanation(original.getStateExplanation());
+        copy.setStorageLocation(original.getStorageLocation());
         return copy;
     }