fix: return new API keys from create/rotateApiKey methods

Author: bevzzz

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

@@ -21,25 +21,15 @@ public class Weaviate extends WeaviateContainer {
 
   private volatile SharedClient clientInstance;
 
-  public WeaviateClient getClient() {
-    return getClient(ObjectBuilder.identity());
-  }
-
   /**
    * Create a new instance of WeaviateClient connected to this container if none
    * exist. Get an existing client otherwise.
    *
    * The lifetime of this client is tied to that of its container, which means
    * that you do not need to {@code close} it manually. It will only truly close
    * after the parent Testcontainer is stopped.
-   *
-   * FIXME: we cannot return the same client for 2 different sets of
-   * configurations.
-   * What we should do is: {@link #getClient()} returns the shared client, while
-   * this one always constructs a new instance.
-   * Otherwise we'll get a race condition once the tests are parallelized.
    */
-  public WeaviateClient getClient(Function<Config.Custom, ObjectBuilder<Config>> fn) {
+  public WeaviateClient getClient() {
     if (!isRunning()) {
       start();
     }
@@ -49,19 +39,8 @@ public WeaviateClient getClient(Function<Config.Custom, ObjectBuilder<Config>> f
 
     synchronized (this) {
       if (clientInstance == null) {
-        var host = getHost();
-        var customFn = ObjectBuilder.partial(fn,
-            conn -> conn
-                .scheme("http")
-                .httpHost(host)
-                .grpcHost(host)
-                .httpPort(getMappedPort(8080))
-                .grpcPort(getMappedPort(50051)));
-        var config = customFn.apply(new Config.Custom()).build();
-        if (config.authentication() != null) {
-        }
         try {
-          clientInstance = new SharedClient(config, this);
+          clientInstance = new SharedClient(Config.of(defaultConfigFn()), this);
         } catch (Exception e) {
           throw new RuntimeException("create WeaviateClient for Weaviate container", e);
         }
@@ -75,19 +54,26 @@ public WeaviateClient getClient(Function<Config.Custom, ObjectBuilder<Config>> f
    * Prefer using {@link #getClient} unless your test needs the initialization
    * steps to run, e.g. OIDC authorization grant exchange.
    */
-  public WeaviateClient getNewClient(Function<Config.Custom, ObjectBuilder<Config>> fn) {
+  public WeaviateClient getClient(Function<Config.Custom, ObjectBuilder<Config>> fn) {
     if (!isRunning()) {
       start();
     }
+
+    var customFn = ObjectBuilder.partial(fn, defaultConfigFn());
+    var config = customFn.apply(new Config.Custom()).build();
+    try {
+      return new WeaviateClient(config);
+    } catch (Exception e) {
+      throw new RuntimeException("create WeaviateClient for Weaviate container", e);
+    }
+  }
+
+  private Function<Config.Custom, ObjectBuilder<Config>> defaultConfigFn() {
     var host = getHost();
-    var customFn = ObjectBuilder.partial(fn,
-        conn -> conn
-            .scheme("http")
-            .httpHost(host)
-            .grpcHost(host)
-            .httpPort(getMappedPort(8080))
-            .grpcPort(getMappedPort(50051)));
-    return WeaviateClient.connectToCustom(customFn);
+    return conn -> conn
+        .scheme("http")
+        .httpHost(host).httpPort(getMappedPort(8080))
+        .grpcHost(host).grpcPort(getMappedPort(50051));
   }
 
   public static Weaviate createDefault() {

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

@@ -128,7 +128,7 @@ public void test_clientCredentials() throws Exception {
 
   /** Send an HTTP and gRPC requests using a "sync" client. */
   private static void pingWeaviate(final Weaviate container, Authentication auth) throws Exception {
-    try (final var client = container.getNewClient(conn -> conn.authentication(auth))) {
+    try (final var client = container.getClient(conn -> conn.authentication(auth))) {
       // Make an authenticated HTTP call
       Assertions.assertThat(client.isLive()).isTrue();
 
@@ -143,7 +143,7 @@ private static void pingWeaviate(final Weaviate container, Authentication auth)
 
   /** Send an HTTP and gRPC requests using an "async" client. */
   private static void pingWeaviateAsync(final Weaviate container, Authentication auth) throws Exception {
-    try (final var client = container.getNewClient(conn -> conn.authentication(auth))) {
+    try (final var client = container.getClient(conn -> conn.authentication(auth))) {
       try (final var async = client.async()) {
         // Make an authenticated HTTP call
         Assertions.assertThat(async.isLive().join()).isTrue();

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

@@ -44,7 +44,7 @@ public class RbacITest extends ConcurrentTest {
   /** Name of the viewer role, which exists by default. */
   private static final String VIEWER_ROLE = "viewer";
 
-  private static final WeaviateClient client = Weaviate.custom()
+  private static final Weaviate container = Weaviate.custom()
       .withAdminUsers(ADMIN_USER)
       .withApiKeys(API_KEY)
       .withRbac()
@@ -53,7 +53,9 @@ public class RbacITest extends ConcurrentTest {
           "https://auth.wcs.api.weaviate.io/auth/realms/SeMI",
           "email",
           "groups")
-      .build()
+      .build();
+
+  private static final WeaviateClient client = container
       .getClient(fn -> fn.authentication(Authentication.apiKey(API_KEY)));
 
   @Test
@@ -187,7 +189,9 @@ public void test_users_db() throws IOException {
     var userId = ns("user");
     var roleName = ns("rock-n-role");
 
-    client.users.db.create(userId);
+    var apiKey = client.users.db.create(userId);
+    assertValidApiKey(apiKey);
+
     client.roles.create(roleName);
 
     client.users.db.assignRoles(userId, roleName);
@@ -207,6 +211,9 @@ public void test_users_db() throws IOException {
         .as("user is activated")
         .returns(true, DbUser::active);
 
+    apiKey = client.users.db.rotateKey(userId);
+    assertValidApiKey(apiKey);
+
     client.users.db.deactivate(userId);
     Assertions.assertThat(client.users.db.byName(userId)).get()
         .as("user is deactivated")
@@ -244,4 +251,16 @@ public void test_users_oidc() throws IOException {
         .extracting(Role::name)
         .doesNotContain(roleName);
   }
+
+  /**
+   * Create a new client with API-key authentication
+   * and check that it can make authenticated requests.
+   */
+  private void assertValidApiKey(String apiKey) {
+    try (final var c = container.getClient(cfg -> cfg.authentication(Authentication.apiKey(apiKey)))) {
+      Assertions.assertThatCode(() -> c.isLive()).as("check API key is valid").doesNotThrowAnyException();
+    } catch (Exception e) {
+      throw new AssertionError(e);
+    }
+  }
 }

src/main/java/io/weaviate/client6/v1/api/rbac/users/CreateDbUserRequest.java

@@ -2,14 +2,16 @@
 
 import java.util.Collections;
 
+import io.weaviate.client6.v1.internal.json.JSON;
 import io.weaviate.client6.v1.internal.rest.Endpoint;
 import io.weaviate.client6.v1.internal.rest.SimpleEndpoint;
 import io.weaviate.client6.v1.internal.rest.UrlEncoder;
 
 public record CreateDbUserRequest(String userId) {
 
-  public static final Endpoint<CreateDbUserRequest, Void> _ENDPOINT = SimpleEndpoint.sideEffect(
+  public static final Endpoint<CreateDbUserRequest, String> _ENDPOINT = SimpleEndpoint.noBody(
       __ -> "POST",
       request -> "/users/db/" + UrlEncoder.encodeValue(request.userId),
-      request -> Collections.emptyMap());
+      request -> Collections.emptyMap(),
+      (statusCode, response) -> JSON.deserialize(response, CreateDbUserResponse.class).apiKey());
 }

src/main/java/io/weaviate/client6/v1/api/rbac/users/CreateDbUserResponse.java

@@ -0,0 +1,6 @@
+package io.weaviate.client6.v1.api.rbac.users;
+
+import com.google.gson.annotations.SerializedName;
+
+public record CreateDbUserResponse(@SerializedName("apikey") String apiKey) {
+}

src/main/java/io/weaviate/client6/v1/api/rbac/users/DbUsersClient.java

@@ -19,14 +19,15 @@ public DbUsersClient(RestTransport restTransport) {
    * Create a new "db" user.
    *
    * @param userId User ID.
+   * @return API key for the created user.
    * @throws WeaviateApiException in case the server returned with an
    *                              error status code.
    * @throws IOException          in case the request was not sent successfully
    *                              due to a malformed request, a networking error
    *                              or the server being unavailable.
    */
-  public void create(String userId) throws IOException {
-    this.restTransport.performRequest(new CreateDbUserRequest(userId), CreateDbUserRequest._ENDPOINT);
+  public String create(String userId) throws IOException {
+    return this.restTransport.performRequest(new CreateDbUserRequest(userId), CreateDbUserRequest._ENDPOINT);
   }
 
   /**

src/main/java/io/weaviate/client6/v1/api/rbac/users/DbUsersClientAsync.java

@@ -19,8 +19,9 @@ public DbUsersClientAsync(RestTransport restTransport) {
    * Create a new "db" user.
    *
    * @param userId User ID.
+   * @return API key for the created user.
    */
-  public CompletableFuture<Void> create(String userId) throws IOException {
+  public CompletableFuture<String> create(String userId) throws IOException {
     return this.restTransport.performRequestAsync(new CreateDbUserRequest(userId), CreateDbUserRequest._ENDPOINT);
   }
 

src/main/java/io/weaviate/client6/v1/api/rbac/users/RotateDbUserKeyResponse.java

@@ -2,5 +2,5 @@
 
 import com.google.gson.annotations.SerializedName;
 
-public record RotateDbUserKeyResponse(@SerializedName("apiKey") String apiKey) {
+public record RotateDbUserKeyResponse(@SerializedName("apikey") String apiKey) {
 }