Ensure TreeMapMetaStore range reads return copies

polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TreeMapMetaStore.java

@@ -20,6 +20,7 @@
 
 import jakarta.annotation.Nonnull;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.TreeMap;
 import java.util.concurrent.atomic.AtomicLong;
@@ -63,8 +64,6 @@ public String buildKey(T value) {
     /**
      * read a value in the slice, will return null if not found
      *
-     * <p>TODO: return a copy of each object to avoid mutating the records
-     *
      * @param key key for that value
      */
     public T read(String key) {
@@ -81,16 +80,32 @@ public T read(String key) {
     public List<T> readRange(String prefix) {
       ensureReadTr();
       if (prefix.isEmpty()) {
-        return new ArrayList<>(this.slice.values());
+        return copyValues(this.slice.values());
+      }
+
+      // Get the sub-map with keys in the range [prefix, rangeEndKey(prefix))
+      return copyValues(slice.subMap(prefix, true, rangeEndKey(prefix), false).values());
+    }
+
+    private List<T> copyValues(Collection<T> values) {
+      List<T> copied = new ArrayList<>(values.size());
+      for (T value : values) {
+        copied.add(this.copyRecord.apply(value));
+      }
+      return copied;
+    }
+
+    private List<String> copyKeysInRange(String prefix) {
+      if (prefix.isEmpty()) {
+        return new ArrayList<>(this.slice.keySet());
       }
 
-      // end of the key
-      String endKey =
-          prefix.substring(0, prefix.length() - 1)
-              + (char) (prefix.charAt(prefix.length() - 1) + 1);
+      return new ArrayList<>(slice.subMap(prefix, true, rangeEndKey(prefix), false).keySet());
+    }
 
-      // Get the sub-map with keys in the range [prefix, endKey)
-      return new ArrayList<>(slice.subMap(prefix, true, endKey, false).values());
+    private String rangeEndKey(String prefix) {
+      return prefix.substring(0, prefix.length() - 1)
+          + (char) (prefix.charAt(prefix.length() - 1) + 1);
     }
 
     /**
@@ -132,9 +147,9 @@ public void delete(String key) {
      */
     public void deleteRange(String prefix) {
       ensureReadWriteTr();
-      List<T> elements = this.readRange(prefix);
-      for (T element : elements) {
-        this.delete(element);
+      List<String> keys = this.copyKeysInRange(prefix);
+      for (String key : keys) {
+        this.delete(key);
       }
     }
 

polaris-core/src/test/java/org/apache/polaris/core/persistence/transactional/TreeMapMetaStoreTest.java

@@ -0,0 +1,72 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.polaris.core.persistence.transactional;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.List;
+import org.apache.polaris.core.PolarisDefaultDiagServiceImpl;
+import org.apache.polaris.core.PolarisDiagnostics;
+import org.apache.polaris.core.entity.PolarisGrantRecord;
+import org.junit.jupiter.api.Test;
+
+class TreeMapMetaStoreTest {
+
+  private final PolarisDiagnostics diagnostics = new PolarisDefaultDiagServiceImpl();
+
+  @Test
+  void readRangeReturnsCopiesForEmptyPrefix() {
+    TreeMapMetaStore store = new TreeMapMetaStore(diagnostics);
+    PolarisGrantRecord grantRecord = new PolarisGrantRecord(1L, 2L, 3L, 4L, 5);
+
+    store.runActionInTransaction(
+        diagnostics, () -> store.getSliceGrantRecords().write(grantRecord));
+
+    List<PolarisGrantRecord> range =
+        store.runInReadTransaction(diagnostics, () -> store.getSliceGrantRecords().readRange(""));
+    range.get(0).setPrivilegeCode(99);
+
+    PolarisGrantRecord stored =
+        store.runInReadTransaction(
+            diagnostics, () -> store.getSliceGrantRecords().readRange("").get(0));
+
+    assertThat(stored.getPrivilegeCode()).isEqualTo(5);
+  }
+
+  @Test
+  void readRangeReturnsCopiesForNonEmptyPrefix() {
+    TreeMapMetaStore store = new TreeMapMetaStore(diagnostics);
+    PolarisGrantRecord grantRecord = new PolarisGrantRecord(1L, 2L, 3L, 4L, 5);
+    String prefix = store.buildPrefixKeyComposite(1L, 2L);
+
+    store.runActionInTransaction(
+        diagnostics, () -> store.getSliceGrantRecords().write(grantRecord));
+
+    List<PolarisGrantRecord> range =
+        store.runInReadTransaction(
+            diagnostics, () -> store.getSliceGrantRecords().readRange(prefix));
+    range.get(0).setPrivilegeCode(99);
+
+    PolarisGrantRecord stored =
+        store.runInReadTransaction(
+            diagnostics, () -> store.getSliceGrantRecords().readRange(prefix).get(0));
+
+    assertThat(stored.getPrivilegeCode()).isEqualTo(5);
+  }
+}