diff --git a/helix-core/src/main/java/org/apache/helix/manager/zk/ZKHelixAdmin.java b/helix-core/src/main/java/org/apache/helix/manager/zk/ZKHelixAdmin.java index 8a42e7ba7e..7bc7ab617e 100644 --- a/helix-core/src/main/java/org/apache/helix/manager/zk/ZKHelixAdmin.java +++ b/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); } /** diff --git a/helix-core/src/main/java/org/apache/helix/manager/zk/ZKUtil.java b/helix-core/src/main/java/org/apache/helix/manager/zk/ZKUtil.java index 1015834760..4606d91846 100644 --- a/helix-core/src/main/java/org/apache/helix/manager/zk/ZKUtil.java +++ b/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); } diff --git a/helix-core/src/test/java/org/apache/helix/manager/zk/TestZKUtil.java b/helix-core/src/test/java/org/apache/helix/manager/zk/TestZKUtil.java index baf4311b3b..c00b32ce08 100644 --- a/helix-core/src/test/java/org/apache/helix/manager/zk/TestZKUtil.java +++ b/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")); + } }