Core: Support removing keys from EnvironmentContext (#12103)

apache · Jan 30, 2025 · 02c8b2d · 02c8b2d
1 parent 2a0d5e8
commit 02c8b2d
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/core/src/main/java/org/apache/iceberg/EnvironmentContext.java b/core/src/main/java/org/apache/iceberg/EnvironmentContext.java
@@ -52,4 +52,14 @@ public static Map<String, String> get() {
   public static void put(String key, String value) {
     PROPERTIES.put(key, value);
   }
+
+  /**
+   * Remove the key from the global properties map.
+   *
+   * @param key The key whose value to remove
+   * @return The previous value associated with the key or null
+   */
+  public static String remove(String key) {
+    return PROPERTIES.remove(key);
+  }
 }
diff --git a/core/src/test/java/org/apache/iceberg/TestEnvironmentContext.java b/core/src/test/java/org/apache/iceberg/TestEnvironmentContext.java
@@ -29,4 +29,15 @@ public void testDefaultValue() {
     assertThat(EnvironmentContext.get().get("iceberg-version"))
         .isEqualTo(IcebergBuild.fullVersion());
   }
+
+  @Test
+  public void testPutAndRemove() {
+    EnvironmentContext.put("test-key", "test-value");
+    assertThat(EnvironmentContext.get()).containsEntry("test-key", "test-value");
+
+    assertThat(EnvironmentContext.remove("test-key")).isEqualTo("test-value");
+    assertThat(EnvironmentContext.get()).doesNotContainKey("test-key");
+
+    assertThat(EnvironmentContext.remove("test-key")).isNull();
+  }
 }