feat: add tenants API

Author: bevzzz

pom.xml

@@ -244,6 +244,7 @@
           <artifactId>maven-surefire-plugin</artifactId>
           <version>2.22.2</version>
           <configuration>
+            <trimStackTrace>false</trimStackTrace>
             <argLine>
               <!--
                 Gson (used for JSON serialization) utilizes reflection and needs to be able to access private fields of

src/it/java/io/weaviate/ConcurrentTest.java

@@ -2,8 +2,14 @@
 
 import java.util.Random;
 import java.util.UUID;
+import java.util.concurrent.Callable;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 import org.apache.commons.lang3.RandomStringUtils;
+import org.assertj.core.api.Assertions;
 import org.junit.Rule;
 import org.junit.rules.TestName;
 
@@ -62,4 +68,46 @@ protected static float[] randomVector(int length, float origin, float bound) {
     }
     return vector;
   }
+
+  /**
+   * Check that a condition is eventually met.
+   *
+   * @param cond           Arbitrary code that evaluates the test condition..
+   * @param intervalMillis Polling interval.
+   * @param timeoutSeconds Maximum waiting time.
+   * @param message        Optional failure message.
+   *
+   * @throws AssertionError   if the condition does not evaluate to true
+   *                          within {@code timeoutSeconds} or a thread
+   *                          was interrupted in the meantime.
+   * @throws RuntimeException if an exception occurred when envalating condition.
+   */
+  public static void eventually(Callable<Boolean> cond, int intervalMillis, int timeoutSeconds, String... message) {
+    var check = CompletableFuture.runAsync(() -> {
+      try {
+        while (!Thread.currentThread().isInterrupted() && !cond.call()) {
+          try {
+            Thread.sleep(intervalMillis);
+          } catch (InterruptedException ex) {
+            Thread.currentThread().interrupt();
+          }
+        }
+      } catch (Exception e) {
+        // Propagate to callee
+        throw new RuntimeException(e);
+      }
+    });
+
+    try {
+      check.get(timeoutSeconds, TimeUnit.SECONDS);
+    } catch (TimeoutException ex) {
+      check.cancel(true);
+      Assertions.fail(message.length >= 0 ? message[0] : null, ex);
+    } catch (InterruptedException ex) {
+      Thread.currentThread().interrupt();
+      Assertions.fail(ex);
+    } catch (ExecutionException ex) {
+      throw new RuntimeException(ex);
+    }
+  }
 }

src/it/java/io/weaviate/containers/Container.java

@@ -16,6 +16,7 @@ public class Container {
   public static final Weaviate WEAVIATE = Weaviate.createDefault();
   public static final Contextionary CONTEXTIONARY = Contextionary.createDefault();
   public static final Img2VecNeural IMG2VEC_NEURAL = Img2VecNeural.createDefault();
+  public static final MinIo MINIO = MinIo.createDefault();
 
   static {
     startAll();

src/it/java/io/weaviate/containers/MinIo.java

@@ -0,0 +1,20 @@
+package io.weaviate.containers;
+
+import org.testcontainers.containers.MinIOContainer;
+
+public class MinIo extends MinIOContainer {
+  private static final String DOCKER_IMAGE = "minio/minio";
+  public static final String ACCESS_KEY = "minioadmin";
+  public static final String SECRET_KEY = "minioadmin";
+
+  static MinIo createDefault() {
+    return new MinIo();
+  }
+
+  private MinIo() {
+    super(DOCKER_IMAGE);
+    withUserName(ACCESS_KEY);
+    withPassword(SECRET_KEY);
+    withCreateContainerCmdModifier(cmd -> cmd.withHostName("minio"));
+  }
+}

src/it/java/io/weaviate/containers/Weaviate.java

@@ -89,6 +89,15 @@ public Builder withImageInference(String url, String module) {
       return this;
     }
 
+    public Builder withOffloadS3(String accessKey, String secretKey) {
+      addModules("offload-s3");
+      environment.put("OFFLOAD_S3_ENDPOINT", "http://minio:9000");
+      environment.put("OFFLOAD_S3_BUCKET_AUTO_CREATE", "true");
+      environment.put("AWS_ACCESS_KEY_ID", accessKey);
+      environment.put("AWS_SECRET_KEY", secretKey);
+      return this;
+    }
+
     public Builder enableTelemetry(boolean enable) {
       telemetry = enable;
       return this;

src/it/java/io/weaviate/integration/TenantsITest.java

@@ -0,0 +1,74 @@
+package io.weaviate.integration;
+
+import org.assertj.core.api.Assertions;
+import org.junit.Test;
+
+import io.weaviate.ConcurrentTest;
+import io.weaviate.client6.v1.api.WeaviateClient;
+import io.weaviate.client6.v1.api.collections.tenants.Tenant;
+import io.weaviate.containers.Container;
+import io.weaviate.containers.Container.ContainerGroup;
+import io.weaviate.containers.MinIo;
+import io.weaviate.containers.Weaviate;
+
+public class TenantsITest extends ConcurrentTest {
+  private static final ContainerGroup compose = Container.compose(
+      Weaviate.custom()
+          .withOffloadS3(MinIo.ACCESS_KEY, MinIo.SECRET_KEY)
+          .build(),
+      Container.MINIO);
+
+  private static WeaviateClient client = compose.getClient();
+
+  @Test
+  public void test_tenantLifecycle() throws Exception {
+    var nsThings = ns("Things");
+
+    client.collections.create(
+        nsThings, c -> c
+            .multiTenancy(mt -> mt
+                .autoTenantCreation(false)
+                .autoTenantActivation(false)));
+
+    var things = client.collections.use(nsThings);
+
+    // No tenants at first
+    Assertions.assertThat(things.tenants.list()).as("no tenants initially").isEmpty();
+
+    var allison = Tenant.active("active-allison");
+    var isaac = Tenant.inactive("inactive-isaac");
+    var owen = Tenant.inactive("offloaded-owen");
+
+    things.tenants.create(allison, isaac, owen);
+
+    // Collection has 2 tenants creted just now.
+    Assertions.assertThat(things.tenants.list()).as("list created tenants").hasSize(3);
+    Assertions.assertThat(things.tenants.exists(allison.name()))
+        .describedAs("%s exists", allison.name()).isTrue();
+    Assertions.assertThat(things.tenants.exists(isaac.name()))
+        .describedAs("%s exists", isaac.name()).isTrue();
+    Assertions.assertThat(things.tenants.exists(owen.name()))
+        .describedAs("%s exists", owen.name()).isTrue();
+
+    things.tenants.activate(isaac.name());
+    eventually(() -> things.tenants.get(isaac.name()).get().isActive(),
+        200, 2, isaac.name() + " not activated");
+
+    things.tenants.deactivate(allison.name());
+    eventually(() -> things.tenants.get(allison.name()).get().isInactive(),
+        200, 2, allison.name() + " not deactivated");
+
+    things.tenants.offload(owen.name());
+    eventually(() -> things.tenants.get(owen.name()).get().isOffloaded(),
+        200, 2, owen.name() + " not offloaded");
+
+    things.tenants.delete(allison.name(), isaac.name(), owen.name());
+    Assertions.assertThat(things.tenants.list()).as("no tenants after deletion").isEmpty();
+    Assertions.assertThat(things.tenants.exists(allison.name()))
+        .describedAs("%s not exists", allison.name()).isFalse();
+    Assertions.assertThat(things.tenants.exists(isaac.name()))
+        .describedAs("%s not exists", isaac.name()).isFalse();
+    Assertions.assertThat(things.tenants.exists(owen.name()))
+        .describedAs("%s not exists", owen.name()).isFalse();
+  }
+}

src/main/java/io/weaviate/client6/v1/api/collections/CollectionHandle.java

@@ -9,6 +9,7 @@
 import io.weaviate.client6.v1.api.collections.pagination.Paginator;
 import io.weaviate.client6.v1.api.collections.query.ConsistencyLevel;
 import io.weaviate.client6.v1.api.collections.query.WeaviateQueryClient;
+import io.weaviate.client6.v1.api.collections.tenants.WeaviateTenantsClient;
 import io.weaviate.client6.v1.internal.ObjectBuilder;
 import io.weaviate.client6.v1.internal.grpc.GrpcTransport;
 import io.weaviate.client6.v1.internal.orm.CollectionDescriptor;
@@ -19,6 +20,7 @@ public class CollectionHandle<PropertiesT> {
   public final WeaviateDataClient<PropertiesT> data;
   public final WeaviateQueryClient<PropertiesT> query;
   public final WeaviateAggregateClient aggregate;
+  public final WeaviateTenantsClient tenants;
 
   private final CollectionHandleDefaults defaults;
 
@@ -31,6 +33,7 @@ public CollectionHandle(
     this.aggregate = new WeaviateAggregateClient(collection, grpcTransport, defaults);
     this.query = new WeaviateQueryClient<>(collection, grpcTransport, defaults);
     this.data = new WeaviateDataClient<>(collection, restTransport, grpcTransport, defaults);
+    this.tenants = new WeaviateTenantsClient(collection, restTransport, grpcTransport);
 
     this.defaults = defaults;
   }
@@ -41,6 +44,7 @@ private CollectionHandle(CollectionHandle<PropertiesT> c, CollectionHandleDefaul
     this.aggregate = c.aggregate;
     this.query = new WeaviateQueryClient<>(c.query, defaults);
     this.data = new WeaviateDataClient<>(c.data, defaults);
+    this.tenants = c.tenants;
 
     this.defaults = defaults;
   }

src/main/java/io/weaviate/client6/v1/api/collections/MultiTenancy.java

@@ -7,30 +7,38 @@
 import io.weaviate.client6.v1.internal.ObjectBuilder;
 
 public record MultiTenancy(
+    @SerializedName("enabled") Boolean enabled,
     @SerializedName("autoTenantCreation") Boolean createAutomatically,
-    @SerializedName("autoTenantActivate") Boolean activateAutomatically) {
+    @SerializedName("autoTenantActivation") Boolean activateAutomatically) {
 
   public static MultiTenancy of(Function<Builder, ObjectBuilder<MultiTenancy>> fn) {
     return fn.apply(new Builder()).build();
   }
 
   public MultiTenancy(Builder builder) {
     this(
+        builder.enabled,
         builder.createAutomatically,
         builder.activateAutomatically);
   }
 
   public static class Builder implements ObjectBuilder<MultiTenancy> {
+    private Boolean enabled = true;
     private Boolean createAutomatically;
     private Boolean activateAutomatically;
 
-    public Builder createAutomatically(boolean createAutomatically) {
-      this.createAutomatically = createAutomatically;
+    public Builder enabled(boolean enabled) {
+      this.enabled = enabled;
       return this;
     }
 
-    public Builder activateAutomatically(boolean activateAutomatically) {
-      this.activateAutomatically = activateAutomatically;
+    public Builder autoTenantCreation(boolean enabled) {
+      this.createAutomatically = enabled;
+      return this;
+    }
+
+    public Builder autoTenantActivation(boolean enabled) {
+      this.activateAutomatically = enabled;
       return this;
     }
 

src/main/java/io/weaviate/client6/v1/api/collections/tenants/CreateTenantsRequest.java

@@ -0,0 +1,19 @@
+package io.weaviate.client6.v1.api.collections.tenants;
+
+import java.util.Collections;
+import java.util.List;
+
+import io.weaviate.client6.v1.internal.json.JSON;
+import io.weaviate.client6.v1.internal.orm.CollectionDescriptor;
+import io.weaviate.client6.v1.internal.rest.Endpoint;
+import io.weaviate.client6.v1.internal.rest.SimpleEndpoint;
+
+public record CreateTenantsRequest(List<Tenant> tenants) {
+  static Endpoint<CreateTenantsRequest, Void> endpoint(CollectionDescriptor<?> collection) {
+    return SimpleEndpoint.sideEffect(
+        __ -> "POST",
+        __ -> "/schema/" + collection.name() + "/tenants",
+        __ -> Collections.emptyMap(),
+        request -> JSON.serialize(request.tenants));
+  }
+}

src/main/java/io/weaviate/client6/v1/api/collections/tenants/DeleteTenantsRequest.java

@@ -0,0 +1,19 @@
+package io.weaviate.client6.v1.api.collections.tenants;
+
+import java.util.Collections;
+import java.util.List;
+
+import io.weaviate.client6.v1.internal.json.JSON;
+import io.weaviate.client6.v1.internal.orm.CollectionDescriptor;
+import io.weaviate.client6.v1.internal.rest.Endpoint;
+import io.weaviate.client6.v1.internal.rest.SimpleEndpoint;
+
+public record DeleteTenantsRequest(List<String> tenants) {
+  static Endpoint<DeleteTenantsRequest, Void> endpoint(CollectionDescriptor<?> collection) {
+    return SimpleEndpoint.sideEffect(
+        __ -> "DELETE",
+        __ -> "/schema/" + collection.name() + "/tenants",
+        __ -> Collections.emptyMap(),
+        request -> JSON.serialize(request.tenants));
+  }
+}