diff --git a/src/it/java/io/weaviate/integration/TenantsITest.java b/src/it/java/io/weaviate/integration/TenantsITest.java index c765ae080..8c9e9bc13 100644 --- a/src/it/java/io/weaviate/integration/TenantsITest.java +++ b/src/it/java/io/weaviate/integration/TenantsITest.java @@ -1,5 +1,8 @@ package io.weaviate.integration; +import java.util.List; +import java.util.stream.IntStream; + import org.assertj.core.api.Assertions; import org.junit.Test; @@ -71,4 +74,24 @@ public void test_tenantLifecycle() throws Exception { Assertions.assertThat(things.tenants.exists(owen.name())) .describedAs("%s not exists", owen.name()).isFalse(); } + + @Test + public void test_updateMoreThanOneHundredTenants() throws Exception { + var nsMany = ns("ManyTenants"); + client.collections.create(nsMany, c -> c.multiTenancy(mt -> mt.autoTenantCreation(false))); + var many = client.collections.use(nsMany); + + List names = IntStream.rangeClosed(1, 250).mapToObj(i -> "tenant-" + i).toList(); + many.tenants.create(names.stream().map(Tenant::active).toList()); + Assertions.assertThat(many.tenants.list()).as("created").hasSize(250); + + many.tenants.deactivate(names); + + eventually(() -> many.tenants.list().stream().allMatch(Tenant::isInactive), + 200, 5, "not all tenants were deactivated"); + + many.tenants.activate(names); + eventually(() -> many.tenants.list().stream().allMatch(Tenant::isActive), + 200, 5, "not all tenants were activated"); + } } diff --git a/src/main/java/io/weaviate/client6/v1/api/collections/tenants/UpdateTenantsRequest.java b/src/main/java/io/weaviate/client6/v1/api/collections/tenants/UpdateTenantsRequest.java index c17334be0..6771929b2 100644 --- a/src/main/java/io/weaviate/client6/v1/api/collections/tenants/UpdateTenantsRequest.java +++ b/src/main/java/io/weaviate/client6/v1/api/collections/tenants/UpdateTenantsRequest.java @@ -1,5 +1,6 @@ package io.weaviate.client6.v1.api.collections.tenants; +import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -9,6 +10,20 @@ import io.weaviate.client6.v1.internal.rest.SimpleEndpoint; public record UpdateTenantsRequest(List tenants) { + static final int MAX_TENANTS_PER_REQUEST = 100; + + static List> batches(List tenants) { + if (tenants.size() <= MAX_TENANTS_PER_REQUEST) { + return List.of(List.copyOf(tenants)); + } + + var batches = new ArrayList>(); + for (int from = 0; from < tenants.size(); from += MAX_TENANTS_PER_REQUEST) { + batches.add(List.copyOf(tenants.subList(from, Math.min(from + MAX_TENANTS_PER_REQUEST, tenants.size())))); + } + return batches; + } + static Endpoint endpoint(CollectionDescriptor collection) { return SimpleEndpoint.sideEffect( __ -> "PUT", diff --git a/src/main/java/io/weaviate/client6/v1/api/collections/tenants/WeaviateTenantsClient.java b/src/main/java/io/weaviate/client6/v1/api/collections/tenants/WeaviateTenantsClient.java index fadb596db..94d16e305 100644 --- a/src/main/java/io/weaviate/client6/v1/api/collections/tenants/WeaviateTenantsClient.java +++ b/src/main/java/io/weaviate/client6/v1/api/collections/tenants/WeaviateTenantsClient.java @@ -105,7 +105,7 @@ public List get(List tenants) { /** * Update tenant configuration. * - * @param tenants Tenant names. + * @param tenants Tenant configurations. * @throws WeaviateApiException in case the server returned with an * error status code. * @throws IOException in case the request was not sent successfully @@ -119,7 +119,7 @@ public void update(Tenant... tenants) throws IOException { /** * Update tenant configuration. * - * @param tenants Tenant names. + * @param tenants Tenant configurations. * @throws WeaviateApiException in case the server returned with an * error status code. * @throws IOException in case the request was not sent successfully @@ -127,7 +127,9 @@ public void update(Tenant... tenants) throws IOException { * or the server being unavailable. */ public void update(List tenants) throws IOException { - this.restTransport.performRequest(new UpdateTenantsRequest(tenants), UpdateTenantsRequest.endpoint(collection)); + for (var batch : UpdateTenantsRequest.batches(tenants)) { + this.restTransport.performRequest(new UpdateTenantsRequest(batch), UpdateTenantsRequest.endpoint(collection)); + } } /** diff --git a/src/main/java/io/weaviate/client6/v1/api/collections/tenants/WeaviateTenantsClientAsync.java b/src/main/java/io/weaviate/client6/v1/api/collections/tenants/WeaviateTenantsClientAsync.java index 5725f55ac..c895c4731 100644 --- a/src/main/java/io/weaviate/client6/v1/api/collections/tenants/WeaviateTenantsClientAsync.java +++ b/src/main/java/io/weaviate/client6/v1/api/collections/tenants/WeaviateTenantsClientAsync.java @@ -55,8 +55,12 @@ public CompletableFuture update(Tenant... tenants) throws IOException { } public CompletableFuture update(List tenants) throws IOException { - return this.restTransport.performRequestAsync(new UpdateTenantsRequest(tenants), - UpdateTenantsRequest.endpoint(collection)); + CompletableFuture updated = CompletableFuture.completedFuture(null); + for (var batch : UpdateTenantsRequest.batches(tenants)) { + updated = updated.thenCompose(__ -> this.restTransport.performRequestAsync( + new UpdateTenantsRequest(batch), UpdateTenantsRequest.endpoint(collection))); + } + return updated; } public CompletableFuture delete(String... tenants) throws IOException { diff --git a/src/test/java/io/weaviate/client6/v1/api/collections/tenants/UpdateTenantsBatchingTest.java b/src/test/java/io/weaviate/client6/v1/api/collections/tenants/UpdateTenantsBatchingTest.java new file mode 100644 index 000000000..713a3560a --- /dev/null +++ b/src/test/java/io/weaviate/client6/v1/api/collections/tenants/UpdateTenantsBatchingTest.java @@ -0,0 +1,114 @@ +package io.weaviate.client6.v1.api.collections.tenants; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.stream.IntStream; + +import org.assertj.core.api.Assertions; +import org.junit.Test; + +import io.weaviate.client6.v1.internal.orm.CollectionDescriptor; +import io.weaviate.testutil.transport.MockRestTransport; + +public class UpdateTenantsBatchingTest { + private static final CollectionDescriptor COLLECTION = CollectionDescriptor.ofMap("Things"); + + private static List tenants(int n) { + return IntStream.rangeClosed(1, n).mapToObj(i -> Tenant.active("tenant-" + i)).toList(); + } + + private static List names(int n) { + return IntStream.rangeClosed(1, n).mapToObj(i -> "tenant-" + i).toList(); + } + + private static List bodies(MockRestTransport transport, int expectAtMost) { + var seen = new ArrayList(); + var assertions = new MockRestTransport.AssertFunction[expectAtMost]; + for (var i = 0; i < expectAtMost; i++) { + assertions[i] = (method, url, body, query) -> seen.add(body); + } + transport.assertNext(assertions); + return seen; + } + + @Test + public void test_batchesOfAtMost100() { + Assertions.assertThat(UpdateTenantsRequest.batches(tenants(100))) + .as("exactly at the limit is one request").hasSize(1); + Assertions.assertThat(UpdateTenantsRequest.batches(tenants(101))) + .as("one over the limit splits") + .extracting(List::size).containsExactly(100, 1); + Assertions.assertThat(UpdateTenantsRequest.batches(tenants(250))) + .extracting(List::size).containsExactly(100, 100, 50); + } + + @Test + public void test_batchesCoverEveryTenantInOrder() { + var all = tenants(250); + + var flattened = new ArrayList(); + UpdateTenantsRequest.batches(all).forEach(flattened::addAll); + + Assertions.assertThat(flattened).isEqualTo(all); + } + + @Test + public void test_emptyListIsASingleEmptyBatch() { + Assertions.assertThat(UpdateTenantsRequest.batches(List.of())).containsExactly(List.of()); + } + + @Test + public void test_updateSplitsIntoSeveralRequests() throws IOException { + var transport = new MockRestTransport(); + + new WeaviateTenantsClient(COLLECTION, transport, null).update(tenants(101)); + + var bodies = bodies(transport, 3); + Assertions.assertThat(bodies).as("101 tenants -> 2 requests").hasSize(2); + Assertions.assertThat(bodies.get(0)).contains("tenant-100").doesNotContain("tenant-101"); + Assertions.assertThat(bodies.get(1)).contains("tenant-101").doesNotContain("tenant-100"); + } + + @Test + public void test_updateAtTheLimitIsASingleRequest() throws IOException { + var transport = new MockRestTransport(); + + new WeaviateTenantsClient(COLLECTION, transport, null).update(tenants(100)); + + Assertions.assertThat(bodies(transport, 2)).hasSize(1); + } + + @Test + public void test_deactivateSplitsIntoSeveralRequests() throws IOException { + var transport = new MockRestTransport(); + + new WeaviateTenantsClient(COLLECTION, transport, null).deactivate(names(201)); + + var bodies = bodies(transport, 4); + Assertions.assertThat(bodies).as("201 tenants -> 3 requests").hasSize(3); + Assertions.assertThat(bodies).allSatisfy(body -> Assertions.assertThat(body).contains("INACTIVE")); + } + + @Test + public void test_createIsNotSplit() throws IOException { + var transport = new MockRestTransport(); + + new WeaviateTenantsClient(COLLECTION, transport, null).create(tenants(250)); + + var bodies = bodies(transport, 2); + Assertions.assertThat(bodies).as("adding tenants is not capped by the server").hasSize(1); + Assertions.assertThat(bodies.get(0)).contains("tenant-1").contains("tenant-250"); + } + + @Test + public void test_asyncUpdateSplitsIntoSeveralRequests() throws Exception { + var transport = new MockRestTransport(); + + new WeaviateTenantsClientAsync(COLLECTION, transport, null).update(tenants(101)).get(); + + var bodies = bodies(transport, 3); + Assertions.assertThat(bodies).hasSize(2); + Assertions.assertThat(bodies.get(1)).contains("tenant-101"); + } +} diff --git a/src/test/java/io/weaviate/testutil/transport/MockRestTransport.java b/src/test/java/io/weaviate/testutil/transport/MockRestTransport.java index 7991078cd..19225964a 100644 --- a/src/test/java/io/weaviate/testutil/transport/MockRestTransport.java +++ b/src/test/java/io/weaviate/testutil/transport/MockRestTransport.java @@ -51,11 +51,15 @@ public ResponseT performRequest(RequestT reque return null; } + @SuppressWarnings("unchecked") @Override public CompletableFuture performRequestAsync(RequestT request, Endpoint endpoint) { requests.add(new Request<>(request, endpoint)); - return null; + if (endpoint instanceof BooleanEndpoint) { + return CompletableFuture.completedFuture((ResponseT) Boolean.TRUE); + } + return CompletableFuture.completedFuture(null); } @Override