Merge pull request #73 from linkedin/ngangwar/ngangwar_update_ideal_state

Update Ideal State API creating new Ideal State

Author: ngngwr

helix-core/src/main/java/org/apache/helix/manager/zk/ZKHelixAdmin.java

@@ -1734,7 +1734,7 @@ public void updateIdealState(String clusterName, String resourceName, IdealState
           resourceName));
     }
     // Update by way of merge
-    ZKUtil.createOrUpdate(_zkClient, zkPath, idealState.getRecord(), true, true);
+    ZKUtil.upsertWithOptionalCreate(_zkClient, zkPath, idealState.getRecord(), true, true, false);
   }
 
   /**

helix-core/src/main/java/org/apache/helix/manager/zk/ZKUtil.java

@@ -427,6 +427,11 @@ public static void createOrUpdate(String zkAddress, String path, final ZNRecord
 
   public static void createOrUpdate(RealmAwareZkClient client, String path, final ZNRecord record,
       final boolean persistent, final boolean mergeOnUpdate) {
+    upsertWithOptionalCreate(client, path, record, persistent, mergeOnUpdate, true);
+  }
+
+  public static void upsertWithOptionalCreate(RealmAwareZkClient client, String path, final ZNRecord record,
+      final boolean persistent, final boolean mergeOnUpdate, final boolean allowCreate) {
     int retryCount = 0;
     while (retryCount < RETRYLIMIT) {
       try {
@@ -443,6 +448,10 @@ public ZNRecord update(ZNRecord currentData) {
           };
           client.updateDataSerialized(path, updater);
         } else {
+          if (!allowCreate) {
+            throw new HelixException(("Path " + path + " does not exist"));
+          }
+
           CreateMode mode = (persistent) ? CreateMode.PERSISTENT : CreateMode.EPHEMERAL;
           client.create(path, record, mode);
         }

helix-core/src/test/java/org/apache/helix/manager/zk/TestZKUtil.java

@@ -218,4 +218,62 @@ record = _gZkClient.readData(path);
       put("k2", "v2");
     }}, record.getMapField("map"));
   }
+
+  @Test()
+  public void testUpsertWithOptionalCreate_CreateWhenNotExists() {
+    // Test: Create node when it doesn't exist and allowCreate is true
+    String path = PropertyPathBuilder.instanceConfig(clusterName, "id10");
+    ZNRecord record = new ZNRecord("id10");
+    record.setSimpleField("key1", "value1");
+
+    // Node doesn't exist yet
+    AssertJUnit.assertFalse(_gZkClient.exists(path));
+
+    // Create with allowCreate=true, persistent=true
+    ZKUtil.upsertWithOptionalCreate(_gZkClient, path, record, true, true, true);
+
+    // Verify node was created
+    AssertJUnit.assertTrue(_gZkClient.exists(path));
+    ZNRecord result = _gZkClient.readData(path);
+    AssertJUnit.assertEquals("id10", result.getId());
+    AssertJUnit.assertEquals("value1", result.getSimpleField("key1"));
+  }
+
+  @Test()
+  public void testUpsertWithOptionalCreate_NoCreateWhenNotExists() {
+    // Test: Don't create node when it doesn't exist and allowCreate is false
+    String path = PropertyPathBuilder.instanceConfig(clusterName, "id11");
+    ZNRecord record = new ZNRecord("id11");
+    record.setSimpleField("key1", "value1");
+
+    // Node doesn't exist yet
+    AssertJUnit.assertFalse(_gZkClient.exists(path));
+
+    // Try to update with allowCreate=false
+    ZKUtil.upsertWithOptionalCreate(_gZkClient, path, record, true, true, false);
+
+    // Verify node was NOT created
+    AssertJUnit.assertFalse(_gZkClient.exists(path));
+  }
+
+  @Test()
+  public void testUpsertWithOptionalCreate_UpdateExistingWithAllowCreateFalse() {
+    // Test: Update should work even when allowCreate=false if node exists
+    String path = PropertyPathBuilder.instanceConfig(clusterName, "id15");
+
+    // Create initial node
+    ZNRecord initial = new ZNRecord("id15");
+    initial.setSimpleField("key1", "value1");
+    _gZkClient.createPersistent(path, initial);
+
+    // Update with allowCreate=false (should still work since node exists)
+    ZNRecord update = new ZNRecord("id15");
+    update.setSimpleField("key2", "value2");
+    ZKUtil.upsertWithOptionalCreate(_gZkClient, path, update, true, true, false);
+
+    // Verify update happened
+    ZNRecord result = _gZkClient.readData(path);
+    AssertJUnit.assertEquals("value1", result.getSimpleField("key1"));
+    AssertJUnit.assertEquals("value2", result.getSimpleField("key2"));
+  }
 }