Keyspace import

fdb-extensions/src/test/java/com/apple/foundationdb/tuple/TupleHelpersTest.java

@@ -0,0 +1,145 @@
+/*
+ * TupleHelpersTest.java
+ *
+ * This source file is part of the FoundationDB open source project
+ *
+ * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors
+ *
+ * Licensed 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 com.apple.foundationdb.tuple;
+
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import java.util.UUID;
+import java.util.stream.Stream;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class TupleHelpersTest {
+
+    static Stream<Arguments> isPrefixTrueCases() {
+        UUID uuid1 = UUID.fromString("12345678-1234-1234-1234-123456789abc");
+        byte[] data1 = {0x01, 0x02, 0x03};
+        byte[] data2 = {0x01, 0x02, 0x03};
+
+        return Stream.of(
+                Arguments.of("empty tuples", 
+                        Tuple.from(), 
+                        Tuple.from()),
+                Arguments.of("empty prefix and non-empty target", 
+                        Tuple.from(), 
+                        Tuple.from("test", 42)),
+                Arguments.of("single element match", 
+                        Tuple.from("test"), 
+                        Tuple.from("test", 42, true)),
+                Arguments.of("multiple elements match", 
+                        Tuple.from("test", 42, true), 
+                        Tuple.from("test", 42, true, "more", "data")),
+                Arguments.of("equal tuples", 
+                        Tuple.from("test", 42, true), 
+                        Tuple.from("test", 42, true)),
+                Arguments.of("different data types", 
+                        Tuple.from("string", 123L, 4.5f, 6.7d, true), 
+                        Tuple.from("string", 123L, 4.5f, 6.7d, true, "more", false)),
+                Arguments.of("mixed type match", 
+                        Tuple.from("string", 123), 
+                        Tuple.from("string", 123L)),
+                Arguments.of("binary data", 
+                        Tuple.from("test", data1), 
+                        Tuple.from("test", data2, "more")),
+                Arguments.of("uuids", 
+                        Tuple.from("test", uuid1), 
+                        Tuple.from("test", uuid1, UUID.randomUUID())),
+                Arguments.of("null values", 
+                        Tuple.from("test", null), 
+                        Tuple.from("test", null, "more")),
+                Arguments.of("number variations", 
+                        Tuple.from(42, 3.14, -7L), 
+                        Tuple.from(42, 3.14, -7L, "suffix")),
+                Arguments.of("complex data",
+                        Tuple.from("company", 1L, "engineering", 100L),
+                        Tuple.from("company", 1L, "engineering", 100L, "employee_data", true, 4.5f)),
+                Arguments.of("nested tuple match", 
+                        Tuple.from("outer", Tuple.from("inner", 42)), 
+                        Tuple.from("outer", Tuple.from("inner", 42), "more")),
+                Arguments.of("nested tuple with different depths", 
+                        Tuple.from("root", Tuple.from("level1", Tuple.from("level2", "value"))), 
+                        Tuple.from("root", Tuple.from("level1", Tuple.from("level2", "value")), "extra", "data")),
+                Arguments.of("empty nested tuple", 
+                        Tuple.from("outer", Tuple.from()), 
+                        Tuple.from("outer", Tuple.from(), "suffix"))
+        );
+    }
+
+    @ParameterizedTest(name = "{0}")
+    @MethodSource("isPrefixTrueCases")
+    void isPrefixShouldReturnTrue(String description, Tuple prefix, Tuple target) {
+        assertTrue(TupleHelpers.isPrefix(prefix, target));
+    }
+
+    static Stream<Arguments> isPrefixFalseCases() {
+        byte[] data1 = {0x01, 0x02, 0x03};
+        byte[] data2 = {0x01, 0x02, 0x04};
+        byte[] data3 = {0x01, 0x02};
+
+        return Stream.of(
+                Arguments.of("non-empty prefix and empty target", 
+                        Tuple.from("test"), 
+                        Tuple.from()),
+                Arguments.of("single element mismatch", 
+                        Tuple.from("test"), 
+                        Tuple.from("different", 42)),
+                Arguments.of("multiple elements partial match", 
+                        Tuple.from("test", 42, false), 
+                        Tuple.from("test", 42, true, "more")),
+                Arguments.of("prefix longer than target", 
+                        Tuple.from("test", 42, true, "extra"), 
+                        Tuple.from("test", 42)),
+                Arguments.of("different binary data", 
+                        Tuple.from("test", data1), 
+                        Tuple.from("test", data2, "more")),
+                Arguments.of("null mismatch", 
+                        Tuple.from("test", null), 
+                        Tuple.from("test", "not-null")),
+                Arguments.of("number mismatch", 
+                        Tuple.from(42, 3.14), 
+                        Tuple.from(42, 3.15, "suffix")),
+                Arguments.of("byte[] prefix of other", 
+                        Tuple.from().add(data3), 
+                        Tuple.from().add(data2)),
+                Arguments.of("string prefix of other", 
+                        Tuple.from().add("ap"), 
+                        Tuple.from().add("apple")),
+                Arguments.of("nested tuple mismatch", 
+                        Tuple.from("outer", Tuple.from("inner", 42)), 
+                        Tuple.from("outer", Tuple.from("inner", 43), "more")),
+                Arguments.of("nested tuple different structure", 
+                        Tuple.from("root", Tuple.from("level1", "value")), 
+                        Tuple.from("root", Tuple.from("level2", "value"), "extra")),
+                Arguments.of("nested vs non-nested", 
+                        Tuple.from("outer", Tuple.from("inner")), 
+                        Tuple.from("outer", "inner", "more"))
+        );
+    }
+
+    @ParameterizedTest(name = "{0}")
+    @MethodSource("isPrefixFalseCases")
+    void isPrefixShouldReturnFalse(String description, Tuple prefix, Tuple target) {
+        assertFalse(TupleHelpers.isPrefix(prefix, target));
+    }
+}

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/keyspace/DataInKeySpacePath.java

@@ -0,0 +1,76 @@
+/*
+ * DataInKeySpacePath.java
+ *
+ * This source file is part of the FoundationDB open source project
+ *
+ * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors
+ *
+ * Licensed 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 com.apple.foundationdb.record.provider.foundationdb.keyspace;
+
+import com.apple.foundationdb.KeyValue;
+import com.apple.foundationdb.record.provider.foundationdb.FDBRecordContext;
+import com.apple.foundationdb.tuple.Tuple;
+import com.apple.foundationdb.tuple.TupleHelpers;
+
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * Class representing a {@link KeyValue} pair within in {@link KeySpacePath}.
+ */
+public class DataInKeySpacePath {
+
+    final CompletableFuture<ResolvedKeySpacePath> resolvedPath;
+    final KeyValue rawKeyValue;
+
+    public DataInKeySpacePath(KeySpacePath path, KeyValue rawKeyValue, FDBRecordContext context) {
+        this.rawKeyValue = rawKeyValue;
+
+        // Convert the raw key to a Tuple and resolve it starting from the provided path
+        Tuple keyTuple = Tuple.fromBytes(rawKeyValue.getKey());
+
+        // First resolve the provided path to get its resolved form
+        this.resolvedPath = path.toResolvedPathAsync(context).thenCompose(resolvedPath -> {
+            // Now use the resolved path to find the child for the key
+            // We need to figure out how much of the key corresponds to the resolved path
+            Tuple pathTuple = resolvedPath.toTuple();
+            int pathLength = pathTuple.size();
+
+            // The remaining part of the key should be resolved from the resolved path's directory
+            if (keyTuple.size() > pathLength) {
+                // There's more in the key than just the path, so resolve the rest
+                if (resolvedPath.getDirectory().getSubdirectories().isEmpty()) {
+                    return CompletableFuture.completedFuture(
+                            new ResolvedKeySpacePath(resolvedPath.getParent(), resolvedPath.toPath(),
+                                    resolvedPath.getResolvedPathValue(),
+                                    TupleHelpers.subTuple(keyTuple, pathTuple.size(), keyTuple.size())));
+                } else {
+                    return resolvedPath.getDirectory().findChildForKey(context, resolvedPath, keyTuple, keyTuple.size(), pathLength);
+                }
+            } else {
+                // The key exactly matches the path
+                return CompletableFuture.completedFuture(resolvedPath);
+            }
+        });
+    }
+
+    public CompletableFuture<ResolvedKeySpacePath> getResolvedPath() {
+        return resolvedPath;
+    }
+
+    public KeyValue getRawKeyValue() {
+        return rawKeyValue;
+    }
+}

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/keyspace/KeySpacePath.java

@@ -566,4 +566,34 @@ default List<ResolvedKeySpacePath> listSubdirectory(@Nonnull FDBRecordContext co
      */
     @API(API.Status.UNSTABLE)
     String toString(@Nonnull Tuple tuple);
+
+    /**
+     * Exports all data stored under this {@code KeySpacePath} and returns it in a {@link RecordCursor}.
+     *
+     * @param context the transaction context in which to perform the data export
+     * @param continuation continuation from a previous export operation, or null to start from the beginning
+     * @param scanProperties properties controlling how the scan should be performed
+     * @return a {link RecordCursor} that provides all the data under this path
+     */
+    @API(API.Status.EXPERIMENTAL)
+    @Nonnull
+    RecordCursor<DataInKeySpacePath> exportAllData(@Nonnull FDBRecordContext context,
+                                                   @Nullable byte[] continuation,
+                                                   @Nonnull ScanProperties scanProperties);
+
+    /**
+     * Imports the provided data exported via {@link #exportAllData} into this {@code KeySpacePath}.
+     * This will validate that any data provided in {@code dataToImport} has a path that should be in this path,
+     * or one of the sub-directories, if not the future will complete exceptionally with
+     * {@link RecordCoreIllegalImportDataException}.
+     * If there is any data already existing under this path, the new data will overwrite if the keys are the same.
+     * This will use the logical value in the {@link DataInKeySpacePath#getResolvedPath()} to determine the key, rather
+     * than the raw key, meaning that this will work even if the data was exported from a different cluster.
+     * @param context the transaction context in which to save the data
+     * @param dataToImport the data to be saved to the database
+     * @return a future to be completed once all data has been important.
+     */
+    @API(API.Status.EXPERIMENTAL)
+    CompletableFuture<Void> importData(@Nonnull FDBRecordContext context,
+                                       @Nonnull Iterable<DataInKeySpacePath> dataToImport);
 }

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/keyspace/KeySpacePathImpl.java

@@ -26,10 +26,15 @@
 import com.apple.foundationdb.record.RecordCursor;
 import com.apple.foundationdb.record.ScanProperties;
 import com.apple.foundationdb.record.ValueRange;
+import com.apple.foundationdb.record.cursors.LazyCursor;
 import com.apple.foundationdb.record.provider.foundationdb.FDBRecordContext;
+import com.apple.foundationdb.record.provider.foundationdb.KeyValueCursor;
+import com.apple.foundationdb.subspace.Subspace;
 import com.apple.foundationdb.tuple.ByteArrayUtil;
 import com.apple.foundationdb.tuple.Tuple;
+import com.apple.foundationdb.tuple.TupleHelpers;
 import com.google.common.collect.Lists;
+
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import java.util.ArrayList;
@@ -337,6 +342,58 @@ public String toString() {
         return toString(null);
     }
 
+    @Nonnull
+    @Override
+    public RecordCursor<DataInKeySpacePath> exportAllData(@Nonnull FDBRecordContext context,
+                                                          @Nullable byte[] continuation,
+                                                          @Nonnull ScanProperties scanProperties) {
+        return new LazyCursor<>(toTupleAsync(context)
+                .thenApply(tuple -> KeyValueCursor.Builder.withSubspace(new Subspace(tuple))
+                        .setContext(context)
+                        .setContinuation(continuation)
+                        .setScanProperties(scanProperties)
+                        .build()),
+                context.getExecutor())
+                .map(keyValue -> new DataInKeySpacePath(this, keyValue, context));
+    }
+
+    @Nonnull
+    @Override
+    public CompletableFuture<Void> importData(@Nonnull FDBRecordContext context,
+                                              @Nonnull Iterable<DataInKeySpacePath> dataToImport) {
+        return toTupleAsync(context).thenCompose(targetTuple -> {
+            List<CompletableFuture<Void>> importFutures = new ArrayList<>();
+
+            for (DataInKeySpacePath dataItem : dataToImport) {
+                CompletableFuture<Void> importFuture = dataItem.getResolvedPath().thenCompose(resolvedPath -> {
+                    // Validate that this data belongs under this path
+                    Tuple itemTuple = resolvedPath.toTuple();
+                    if (!TupleHelpers.isPrefix(targetTuple, itemTuple)) {
+                        throw new RecordCoreIllegalImportDataException(
+                                "Data item path does not belong under target path",
+                                "target", targetTuple, "item", itemTuple);
+                    }
+
+                    // Reconstruct the key using logical values from the resolved path
+                    Tuple keyTuple = itemTuple;
+                    if (resolvedPath.getRemainder() != null) {
+                        keyTuple = keyTuple.addAll(resolvedPath.getRemainder());
+                    }
+
+                    // Store the data
+                    byte[] keyBytes = keyTuple.pack();
+                    byte[] valueBytes = dataItem.getRawKeyValue().getValue();
+                    context.ensureActive().set(keyBytes, valueBytes);
+
+                    return AsyncUtil.DONE;
+                });
+                importFutures.add(importFuture);
+            }
+
+            return AsyncUtil.whenAll(importFutures);
+        });
+    }
+
     /**
      * Returns this path properly wrapped in whatever implementation the directory the path is contained in dictates.
      */

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/keyspace/KeySpacePathWrapper.java

@@ -226,4 +226,19 @@ public String toString() {
     public String toString(@Nonnull Tuple t) {
         return inner.toString(t);
     }
+
+    @Nonnull
+    @Override
+    public RecordCursor<DataInKeySpacePath> exportAllData(@Nonnull FDBRecordContext context,
+                                                          @Nullable byte[] continuation,
+                                                          @Nonnull ScanProperties scanProperties) {
+        return inner.exportAllData(context, continuation, scanProperties);
+    }
+
+    @Nonnull
+    @Override
+    public CompletableFuture<Void> importData(@Nonnull FDBRecordContext context,
+                                              @Nonnull Iterable<DataInKeySpacePath> dataToImport) {
+        return inner.importData(context, dataToImport);
+    }
 }

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/keyspace/RecordCoreIllegalImportDataException.java

@@ -0,0 +1,33 @@
+/*
+ * RecordCoreIllegalImportDataException.java
+ *
+ * This source file is part of the FoundationDB open source project
+ *
+ * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors
+ *
+ * Licensed 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 com.apple.foundationdb.record.provider.foundationdb.keyspace;
+
+import com.apple.foundationdb.record.RecordCoreArgumentException;
+
+import javax.annotation.Nonnull;
+
+public class RecordCoreIllegalImportDataException extends RecordCoreArgumentException {
+    private static final long serialVersionUID = 1L;
+
+    public RecordCoreIllegalImportDataException(@Nonnull final String msg, @Nonnull final Object... keyValue) {
+        super(msg, keyValue);
+    }
+}