From d4fa20e92faec8f37dc4aedba6c0180d0aeb6e60 Mon Sep 17 00:00:00 2001
From: Pierre Villard <pierre.villard@snowflake.com>
Date: Fri, 4 Apr 2025 16:05:09 +0200
Subject: [PATCH] Add PG name where CS is added

---
 .../java/com/snowflake/openflow/FlowDiff.java   | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/flow-diff/src/main/java/com/snowflake/openflow/FlowDiff.java b/flow-diff/src/main/java/com/snowflake/openflow/FlowDiff.java
index f44b7f0..d3d7689 100644
--- a/flow-diff/src/main/java/com/snowflake/openflow/FlowDiff.java
+++ b/flow-diff/src/main/java/com/snowflake/openflow/FlowDiff.java
@@ -58,6 +58,7 @@ public class FlowDiff {
 
     private static String flowName;
     private static Map<String, VersionedParameterContext> parameterContexts;
+    private static Map<String, VersionedProcessGroup> processGroups;
 
     public static void main(String[] args) throws IOException {
 
@@ -98,8 +99,9 @@ public static void main(String[] args) throws IOException {
                     printConfigurableExtensionProperties(proc);
                 } else if (diff.getComponentB().getComponentType().equals(ComponentType.CONTROLLER_SERVICE)) {
                     final VersionedControllerService cs = (VersionedControllerService) diff.getComponentB();
+                    final String pgName = processGroups.get(cs.getGroupIdentifier()).getName();
                     System.out.println("- A " + printComponent(diff.getComponentB())
-                            + " has been added with the below properties:");
+                            + " has been added in Process Group `" + pgName + "` with the below properties:");
                     printConfigurableExtensionProperties(cs);
                 } else if (diff.getComponentB().getComponentType().equals(ComponentType.LABEL)) {
                     final VersionedLabel label = (VersionedLabel) diff.getComponentB();
@@ -415,6 +417,11 @@ public static Set<FlowDifference> getDiff(final String pathA, final String pathB
         final FlowSnapshotContainer snapshotA = getFlowContainer(pathA, factory);
         final FlowSnapshotContainer snapshotB = getFlowContainer(pathB, factory);
 
+        processGroups = new HashMap<>();
+        VersionedProcessGroup rootPG = snapshotB.getFlowSnapshot().getFlowContents();
+        processGroups.put(rootPG.getIdentifier(), rootPG);
+        registerProcessGroups(rootPG);
+
         // identifier is null for parameter contexts, and we know that names are unique so setting name as id
         snapshotA.getFlowSnapshot().getParameterContexts().values().forEach(pc -> pc.setIdentifier(pc.getName()));
         snapshotB.getFlowSnapshot().getParameterContexts().values().forEach(pc -> pc.setIdentifier(pc.getName()));
@@ -469,6 +476,14 @@ public int compare(FlowDifference o1, FlowDifference o2) {
         return sortedDiffs;
     }
 
+    private static void registerProcessGroups(VersionedProcessGroup rootPG) {
+        Set<VersionedProcessGroup> childPGs = rootPG.getProcessGroups();
+        for (VersionedProcessGroup pg : childPGs) {
+            processGroups.put(pg.getIdentifier(), pg);
+            registerProcessGroups(pg);
+        }
+    }
+
     static FlowSnapshotContainer getFlowContainer(final String path, final JsonFactory factory) throws IOException {
         final File snapshotFile = new File(path);
         try (final JsonParser parser = factory.createParser(snapshotFile)) {