diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TreeMapMetaStore.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TreeMapMetaStore.java
index 2110283e6b..be5fbbe4fc 100644
--- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TreeMapMetaStore.java
+++ b/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
*
- *
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 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 copyValues(Collection values) {
+ List copied = new ArrayList<>(values.size());
+ for (T value : values) {
+ copied.add(this.copyRecord.apply(value));
+ }
+ return copied;
+ }
+
+ private List 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 elements = this.readRange(prefix);
- for (T element : elements) {
- this.delete(element);
+ List keys = this.copyKeysInRange(prefix);
+ for (String key : keys) {
+ this.delete(key);
}
}
diff --git a/polaris-core/src/test/java/org/apache/polaris/core/persistence/transactional/TreeMapMetaStoreTest.java b/polaris-core/src/test/java/org/apache/polaris/core/persistence/transactional/TreeMapMetaStoreTest.java
new file mode 100644
index 0000000000..bc2f6983aa
--- /dev/null
+++ b/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 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 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);
+ }
+}