diff --git a/src/main/java/dev/openfga/sdk/api/client/ClientWriteRequest.java b/src/main/java/dev/openfga/sdk/api/client/ClientWriteRequest.java index ecbebbc..8dc1f35 100644 --- a/src/main/java/dev/openfga/sdk/api/client/ClientWriteRequest.java +++ b/src/main/java/dev/openfga/sdk/api/client/ClientWriteRequest.java @@ -18,6 +18,10 @@ public class ClientWriteRequest { private List writes; private List deletes; + public static ClientWriteRequest ofWrites(List writes) { + return new ClientWriteRequest().writes(writes); + } + public ClientWriteRequest writes(List writes) { this.writes = writes; return this; @@ -27,6 +31,10 @@ public List getWrites() { return writes; } + public static ClientWriteRequest ofDeletes(List deletes) { + return new ClientWriteRequest().deletes(deletes); + } + public ClientWriteRequest deletes(List deletes) { this.deletes = deletes; return this; diff --git a/src/main/java/dev/openfga/sdk/api/client/OpenFgaClient.java b/src/main/java/dev/openfga/sdk/api/client/OpenFgaClient.java index d253127..1c1ce95 100644 --- a/src/main/java/dev/openfga/sdk/api/client/OpenFgaClient.java +++ b/src/main/java/dev/openfga/sdk/api/client/OpenFgaClient.java @@ -23,6 +23,7 @@ import java.util.concurrent.*; import java.util.function.Consumer; import java.util.stream.Collectors; +import java.util.stream.Stream; public class OpenFgaClient { private final ApiClient apiClient; @@ -276,6 +277,15 @@ public CompletableFuture write(ClientWriteRequest request, configuration.assertValid(); String storeId = configuration.getStoreIdChecked(); + if (options != null && options.enableTransactions()) { + return writeTransactions(storeId, request, options); + } else { + return writeOneTransaction(storeId, request, options); + } + } + + private CompletableFuture writeOneTransaction( + String storeId, ClientWriteRequest request, ClientWriteOptions options) { WriteRequest body = new WriteRequest(); if (request != null) { @@ -300,6 +310,44 @@ public CompletableFuture write(ClientWriteRequest request, return call(() -> api.write(storeId, body)).thenApply(ClientWriteResponse::new); } + private CompletableFuture writeTransactions( + String storeId, ClientWriteRequest request, ClientWriteOptions options) { + + int chunkSize = options.getTransactionChunkSize(); + + var writeTransactions = + chunksOf(chunkSize, request.getWrites()).stream().map(ClientWriteRequest::ofWrites); + var deleteTransactions = + chunksOf(chunkSize, request.getDeletes()).stream().map(ClientWriteRequest::ofDeletes); + + var transactions = Stream.concat(writeTransactions, deleteTransactions).collect(Collectors.toList()); + var futureResponse = this.writeOneTransaction(storeId, transactions.get(0), options); + + for (int i = 1; i < transactions.size(); i++) { + final int index = i; // Must be final in this scope for closure. + + // The resulting completable future of this chain will result in either: + // 1. The first exception thrown in a failed completion. Other thenApply() will not be evaluated. + // 2. The final successful ClientWriteResponse. + futureResponse.thenApply(_response -> this.writeOneTransaction(storeId, transactions.get(index), options)); + } + + return futureResponse; + } + + private List> chunksOf(int chunkSize, List list) { + int nChunks = list.size() / chunkSize; + + int finalEndExclusive = list.size(); + List> chunks = new ArrayList<>(); + + for (int i = 0; i < nChunks; i++) { + List chunk = list.subList(i * chunkSize, Math.min((i + 1) * chunkSize, finalEndExclusive)); + chunks.add(chunk); + } + return chunks; + } + /** * WriteTuples - Utility method to write tuples, wraps Write * diff --git a/src/main/java/dev/openfga/sdk/api/configuration/ClientWriteOptions.java b/src/main/java/dev/openfga/sdk/api/configuration/ClientWriteOptions.java index cd60f92..ac2c08c 100644 --- a/src/main/java/dev/openfga/sdk/api/configuration/ClientWriteOptions.java +++ b/src/main/java/dev/openfga/sdk/api/configuration/ClientWriteOptions.java @@ -14,6 +14,8 @@ public class ClientWriteOptions { private String authorizationModelId; + private Boolean enableTransactions; + private int transactionChunkSize; public ClientWriteOptions authorizationModelId(String authorizationModelId) { this.authorizationModelId = authorizationModelId; @@ -23,4 +25,22 @@ public ClientWriteOptions authorizationModelId(String authorizationModelId) { public String getAuthorizationModelId() { return authorizationModelId; } + + public ClientWriteOptions enableTransactions(boolean enableTransactions) { + this.enableTransactions = enableTransactions; + return this; + } + + public boolean enableTransactions() { + return enableTransactions != null && enableTransactions; + } + + public ClientWriteOptions transactionChunkSize(int transactionChunkSize) { + this.transactionChunkSize = transactionChunkSize; + return this; + } + + public int getTransactionChunkSize() { + return transactionChunkSize >= 0 ? transactionChunkSize : 1; + } }