Merge pull request #445 from weaviate/v6-tenant

v6: Per-collection-handle tenant

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

@@ -7,34 +7,54 @@
 import io.weaviate.client6.v1.api.collections.config.WeaviateConfigClient;
 import io.weaviate.client6.v1.api.collections.data.WeaviateDataClient;
 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;
 import io.weaviate.client6.v1.internal.rest.RestTransport;
 
-public class CollectionHandle<T> {
+public class CollectionHandle<PropertiesT> {
   public final WeaviateConfigClient config;
-  public final WeaviateDataClient<T> data;
-  public final WeaviateQueryClient<T> query;
+  public final WeaviateDataClient<PropertiesT> data;
+  public final WeaviateQueryClient<PropertiesT> query;
   public final WeaviateAggregateClient aggregate;
+  public final WeaviateTenantsClient tenants;
+
+  private final CollectionHandleDefaults defaults;
 
   public CollectionHandle(
       RestTransport restTransport,
       GrpcTransport grpcTransport,
-      CollectionDescriptor<T> collectionDescriptor) {
+      CollectionDescriptor<PropertiesT> collection,
+      CollectionHandleDefaults defaults) {
+    this.config = new WeaviateConfigClient(collection, restTransport, grpcTransport, defaults);
+    this.aggregate = new WeaviateAggregateClient(collection, grpcTransport, defaults);
+    this.query = new WeaviateQueryClient<>(collection, grpcTransport, defaults);
+    this.data = new WeaviateDataClient<>(collection, restTransport, grpcTransport, defaults);
+    this.defaults = defaults;
+
+    this.tenants = new WeaviateTenantsClient(collection, restTransport, grpcTransport);
+  }
+
+  /** Copy constructor that sets new defaults. */
+  private CollectionHandle(CollectionHandle<PropertiesT> c, CollectionHandleDefaults defaults) {
+    this.config = new WeaviateConfigClient(c.config, defaults);
+    this.aggregate = new WeaviateAggregateClient(c.aggregate, defaults);
+    this.query = new WeaviateQueryClient<>(c.query, defaults);
+    this.data = new WeaviateDataClient<>(c.data, defaults);
+    this.defaults = defaults;
 
-    this.config = new WeaviateConfigClient(collectionDescriptor, restTransport, grpcTransport);
-    this.data = new WeaviateDataClient<>(collectionDescriptor, restTransport, grpcTransport);
-    this.query = new WeaviateQueryClient<>(collectionDescriptor, grpcTransport);
-    this.aggregate = new WeaviateAggregateClient(collectionDescriptor, grpcTransport);
+    this.tenants = c.tenants;
   }
 
-  public Paginator<T> paginate() {
+  public Paginator<PropertiesT> paginate() {
     return Paginator.of(this.query);
   }
 
-  public Paginator<T> paginate(Function<Paginator.Builder<T>, ObjectBuilder<Paginator<T>>> fn) {
+  public Paginator<PropertiesT> paginate(
+      Function<Paginator.Builder<PropertiesT>, ObjectBuilder<Paginator<PropertiesT>>> fn) {
     return Paginator.of(this.query, fn);
   }
 
@@ -57,4 +77,25 @@ public Paginator<T> paginate(Function<Paginator.Builder<T>, ObjectBuilder<Pagina
   public long size() {
     return this.aggregate.overAll(all -> all.includeTotalCount(true)).totalCount();
   }
+
+  public ConsistencyLevel consistencyLevel() {
+    return defaults.consistencyLevel();
+  }
+
+  public CollectionHandle<PropertiesT> withConsistencyLevel(ConsistencyLevel consistencyLevel) {
+    return new CollectionHandle<>(this, CollectionHandleDefaults.of(with -> with.consistencyLevel(consistencyLevel)));
+  }
+
+  public String tenant() {
+    return defaults.tenant();
+  }
+
+  public CollectionHandle<PropertiesT> withTenant(String tenant) {
+    return new CollectionHandle<>(this, CollectionHandleDefaults.of(with -> with.tenant(tenant)));
+  }
+
+  public CollectionHandle<PropertiesT> withDefaults(
+      Function<CollectionHandleDefaults.Builder, ObjectBuilder<CollectionHandleDefaults>> fn) {
+    return new CollectionHandle<>(this, CollectionHandleDefaults.of(fn));
+  }
 }

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

@@ -9,7 +9,9 @@
 import io.weaviate.client6.v1.api.collections.config.WeaviateConfigClientAsync;
 import io.weaviate.client6.v1.api.collections.data.WeaviateDataClientAsync;
 import io.weaviate.client6.v1.api.collections.pagination.AsyncPaginator;
+import io.weaviate.client6.v1.api.collections.query.ConsistencyLevel;
 import io.weaviate.client6.v1.api.collections.query.WeaviateQueryClientAsync;
+import io.weaviate.client6.v1.api.collections.tenants.WeaviateTenantsClientAsync;
 import io.weaviate.client6.v1.internal.ObjectBuilder;
 import io.weaviate.client6.v1.internal.grpc.GrpcTransport;
 import io.weaviate.client6.v1.internal.orm.CollectionDescriptor;
@@ -20,16 +22,34 @@ public class CollectionHandleAsync<PropertiesT> {
   public final WeaviateDataClientAsync<PropertiesT> data;
   public final WeaviateQueryClientAsync<PropertiesT> query;
   public final WeaviateAggregateClientAsync aggregate;
+  public final WeaviateTenantsClientAsync tenants;
+
+  private final CollectionHandleDefaults defaults;
 
   public CollectionHandleAsync(
       RestTransport restTransport,
       GrpcTransport grpcTransport,
-      CollectionDescriptor<PropertiesT> collectionDescriptor) {
+      CollectionDescriptor<PropertiesT> collection,
+      CollectionHandleDefaults defaults) {
+
+    this.config = new WeaviateConfigClientAsync(collection, restTransport, grpcTransport, defaults);
+    this.aggregate = new WeaviateAggregateClientAsync(collection, grpcTransport, defaults);
+    this.query = new WeaviateQueryClientAsync<>(collection, grpcTransport, defaults);
+    this.data = new WeaviateDataClientAsync<>(collection, restTransport, grpcTransport, defaults);
+    this.defaults = defaults;
+
+    this.tenants = new WeaviateTenantsClientAsync(collection, restTransport, grpcTransport);
+  }
 
-    this.config = new WeaviateConfigClientAsync(collectionDescriptor, restTransport, grpcTransport);
-    this.data = new WeaviateDataClientAsync<>(collectionDescriptor, restTransport, grpcTransport);
-    this.query = new WeaviateQueryClientAsync<>(collectionDescriptor, grpcTransport);
-    this.aggregate = new WeaviateAggregateClientAsync(collectionDescriptor, grpcTransport);
+  /** Copy constructor that sets new defaults. */
+  private CollectionHandleAsync(CollectionHandleAsync<PropertiesT> c, CollectionHandleDefaults defaults) {
+    this.config = new WeaviateConfigClientAsync(c.config, defaults);
+    this.aggregate = new WeaviateAggregateClientAsync(c.aggregate, defaults);
+    this.query = new WeaviateQueryClientAsync<>(c.query, defaults);
+    this.data = new WeaviateDataClientAsync<>(c.data, defaults);
+    this.defaults = defaults;
+
+    this.tenants = c.tenants;
   }
 
   public AsyncPaginator<PropertiesT> paginate() {
@@ -64,4 +84,26 @@ public CompletableFuture<Long> size() {
     return this.aggregate.overAll(all -> all.includeTotalCount(true))
         .thenApply(AggregateResponse::totalCount);
   }
+
+  public ConsistencyLevel consistencyLevel() {
+    return defaults.consistencyLevel();
+  }
+
+  public CollectionHandleAsync<PropertiesT> withConsistencyLevel(ConsistencyLevel consistencyLevel) {
+    return new CollectionHandleAsync<>(this, CollectionHandleDefaults.of(
+        def -> def.consistencyLevel(consistencyLevel)));
+  }
+
+  public String tenant() {
+    return defaults.tenant();
+  }
+
+  public CollectionHandleAsync<PropertiesT> withTenant(String tenant) {
+    return new CollectionHandleAsync<>(this, CollectionHandleDefaults.of(with -> with.tenant(tenant)));
+  }
+
+  public CollectionHandleAsync<PropertiesT> withDefaults(
+      Function<CollectionHandleDefaults.Builder, ObjectBuilder<CollectionHandleDefaults>> fn) {
+    return new CollectionHandleAsync<>(this, CollectionHandleDefaults.of(fn));
+  }
 }