com.aerospike.vector.client.adminclient.AdminClient Maven / Gradle / Ivy
/* (C)2024 */
package com.aerospike.vector.client.adminclient;
import com.aerospike.vector.client.*;
import com.aerospike.vector.client.internal.ClusterTenderer;
import com.aerospike.vector.client.proto.*;
import com.google.common.base.Preconditions;
import com.google.protobuf.Empty;
import io.grpc.StatusRuntimeException;
import java.util.*;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;
import javax.annotation.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/** Implementation of Vector DB admin client */
public class AdminClient implements IAdminClient {
private static final Logger log = LoggerFactory.getLogger(AdminClient.class);
private final ExecutorService adminExecutorService;
private final ClusterTenderer clusterTenderer;
/**
* Constructor for creating a new AdminClient with specified parameters for the cluster.
*
* @param connectionConfig the configuration settings for connecting to the cluster. This
* includes parameters such as the cluster URL, credentials, TLS and other required
* information.
*/
public AdminClient(ConnectionConfig connectionConfig) {
this.clusterTenderer = new ClusterTenderer(connectionConfig, "adminclient");
this.adminExecutorService =
Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
}
/**
* {@inheritDoc}
*
* This implementation checks if the index already exists before attempting to create it. If
* the index exists, no action is taken. This check helps prevent unnecessary creation attempts
* that might result in exceptions or errors from the database.
*
*
Note: This method uses a high-concurrency ExecutorService to manage gRPC calls, improving
* performance under load but requiring proper shutdown management.
*
* @param indexId {@inheritDoc} Unique identifier for the index; must not be null.
* @param vectorBinName {@inheritDoc} Name of the bin storing vector data; must not be null.
* @param dimensions {@inheritDoc} Must be a positive integer, throws IllegalArgumentException
* otherwise.
* @param vectorDistanceMetric {@inheritDoc}
* @param setFilter {@inheritDoc} Can be null, which means no set filter is applied.
* @param indexParams {@inheritDoc} Optional parameters for fine-tuning the index.
* @param storage {@inheritDoc} Storage settings; can be null if default storage options are
* used.
* @param labels {@inheritDoc} Can be empty.
* @throws RuntimeException if an error occurs during the index creation process, including
* timeout exceptions.
*/
@Override
public void indexCreate(
IndexId indexId,
String vectorBinName,
int dimensions,
VectorDistanceMetric vectorDistanceMetric,
@Nullable String setFilter,
@Nullable HnswParams indexParams,
@Nullable IndexStorage storage,
@Nullable Map labels,
long timeoutInMillis,
long waitTimeInMillis) {
Objects.requireNonNull(indexId, "Index ID cannot be null.");
Objects.requireNonNull(vectorBinName, "Vector bin name cannot be null.");
Objects.requireNonNull(vectorDistanceMetric, "Vector distance metric cannot be null.");
if (dimensions <= 0) {
throw new IllegalArgumentException("Dimensions must be a positive integer.");
}
IndexServiceGrpc.IndexServiceBlockingStub indexService =
clusterTenderer.getIndexServiceBlockingStub();
// Check if the index already exists
boolean indexExists =
indexService.list(IndexListRequest.getDefaultInstance()).getIndicesList().stream()
.anyMatch(i -> i.getId().equals(indexId));
if (indexExists) {
return;
}
IndexDefinition.Builder builder =
IndexDefinition.newBuilder()
.setField(vectorBinName)
.setId(indexId)
.setVectorDistanceMetric(vectorDistanceMetric)
.setDimensions(dimensions);
if (setFilter != null) {
builder.setSetFilter(setFilter);
}
if (indexParams != null) {
builder.setHnswParams(indexParams);
}
if (storage != null) {
builder.setStorage(storage);
}
if (null == labels) {
builder.putAllLabels(new java.util.HashMap<>());
} else {
builder.putAllLabels(labels);
}
indexService.create(IndexCreateRequest.newBuilder().setDefinition(builder.build()).build());
try {
waitForIndexCreation(indexId, timeoutInMillis, waitTimeInMillis);
} catch (InterruptedException e) {
throw new RuntimeException(e);
}
}
/**
* todo complete the documentation Update an existing index.
*
* @param indexId unique identifier for the index
* @param maxMemQueueSize
* @param batchingParams
* @param cachingParams
* @param healerParams
* @param mergeParams
* @param labels
*/
@Override
public void indexUpdate(
IndexId indexId,
@Nullable Integer maxMemQueueSize,
@Nullable HnswBatchingParams batchingParams,
@Nullable HnswCachingParams cachingParams,
@Nullable HnswHealerParams healerParams,
@Nullable HnswIndexMergeParams mergeParams,
@Nullable Map labels)
throws RuntimeException {
IndexUpdateRequest.Builder updateRequest = IndexUpdateRequest.newBuilder();
updateRequest.setIndexId(indexId);
if (labels != null && !labels.isEmpty()) {
updateRequest.putAllLabels(labels);
}
HnswIndexUpdate.Builder hnswIndexUpdateBuilder = HnswIndexUpdate.newBuilder();
if (maxMemQueueSize != null) {
Preconditions.checkArgument(
maxMemQueueSize > 0,
String.format("maxMemQueueSize must be >0, found %d", maxMemQueueSize));
hnswIndexUpdateBuilder.setMaxMemQueueSize(maxMemQueueSize);
}
if (batchingParams != null) {
hnswIndexUpdateBuilder.setBatchingParams(batchingParams);
}
if (cachingParams != null) {
hnswIndexUpdateBuilder.setCachingParams(cachingParams);
}
if (mergeParams != null) {
hnswIndexUpdateBuilder.setMergeParams(mergeParams);
}
if (healerParams != null) {
hnswIndexUpdateBuilder.setHealerParams(healerParams);
}
updateRequest.setHnswIndexUpdate(hnswIndexUpdateBuilder.build());
IndexServiceGrpc.IndexServiceBlockingStub indexService =
clusterTenderer.getIndexServiceBlockingStub();
indexService.update(updateRequest.build());
}
@Override
public void indexDrop(IndexId indexId, long timeoutInMillis, long waitTimeInMillis) {
IndexServiceGrpc.IndexServiceBlockingStub indexService =
clusterTenderer.getIndexServiceBlockingStub();
indexService.drop(IndexDropRequest.newBuilder().setIndexId(indexId).build());
try {
waitForIndexDeletion(indexService, indexId, timeoutInMillis, waitTimeInMillis);
} catch (InterruptedException e) {
throw new RuntimeException(e);
}
}
/**
* Garbage collect vertices identified as invalid before cutoff timestamp.
*
* @param indexId the index from which to garbage collect the invalid vertices
* @param cutoffTimestamp the cutoff timestamp (Unix timestamp) for garbage collecting invalid
* vertices
*/
@Override
public void gcInvalidVertices(IndexId indexId, long cutoffTimestamp) {
IndexServiceGrpc.IndexServiceBlockingStub indexService =
clusterTenderer.getIndexServiceBlockingStub();
GcInvalidVerticesRequest request =
GcInvalidVerticesRequest.newBuilder()
.setIndexId(indexId)
.setCutoffTimestamp(cutoffTimestamp)
.build();
indexService.gcInvalidVertices(request);
}
/**
* List all available vector index
*
* @return list all index related information
*/
@Override
public List indexList(boolean applyDefaults) {
IndexServiceGrpc.IndexServiceBlockingStub indexService =
clusterTenderer.getIndexServiceBlockingStub();
return indexService
.list(IndexListRequest.newBuilder().setApplyDefaults(applyDefaults).build())
.getIndicesList();
}
/**
* Number of unmerged index records in the given index
*
* @param indexId Unique identifier for the index.
* @return number of unmerged record in the index
*/
@Override
public IndexStatusResponse indexStatus(IndexId indexId) {
IndexServiceGrpc.IndexServiceBlockingStub indexService =
clusterTenderer.getIndexServiceBlockingStub();
return indexService.getStatus(IndexStatusRequest.newBuilder().setIndexId(indexId).build());
}
@Override
public IndexDefinition getIndex(IndexId indexId, boolean applyDefaults) {
IndexServiceGrpc.IndexServiceBlockingStub indexService =
clusterTenderer.getIndexServiceBlockingStub();
return indexService.get(
IndexGetRequest.newBuilder()
.setIndexId(indexId)
.setApplyDefaults(applyDefaults)
.build());
}
/**
* Get information about a user.
*
* @param username the username to fetch the user information for
* @return the user information
*/
@Override
public User getUser(String username) {
UserAdminServiceGrpc.UserAdminServiceBlockingStub userAdminService =
clusterTenderer.getUserAdminServiceBlockingStub();
return userAdminService.getUser(GetUserRequest.newBuilder().setUsername(username).build());
}
/**
* Add a user and grant roles.
*
* @param credentials new user credentials.
* @param roles the roles to grant to the new application user.
*/
@Override
public void addUser(
com.aerospike.vector.client.auth.Credentials credentials, Set roles) {
if (!(credentials instanceof com.aerospike.vector.client.auth.PasswordCredentials)) {
throw new IllegalArgumentException("only password credentials supported");
}
UserAdminServiceGrpc.UserAdminServiceBlockingStub userAdminService =
clusterTenderer.getUserAdminServiceBlockingStub();
try {
userAdminService.addUser(
AddUserRequest.newBuilder()
.setCredentials(
toGrpcCredentials(
(com.aerospike.vector.client.auth.PasswordCredentials)
credentials))
.addAllRoles(roles)
.build());
} catch (StatusRuntimeException e) {
log.error("RPC failed: {}", e.getStatus());
throw e;
}
}
/**
* Update user with new credentials.
*
* @param credentials new user credentials.
*/
@Override
public void updateCredentials(com.aerospike.vector.client.auth.Credentials credentials) {
if (!(credentials instanceof com.aerospike.vector.client.auth.Credentials)) {
throw new IllegalArgumentException("only password credentials supported");
}
UserAdminServiceGrpc.UserAdminServiceBlockingStub userAdminService =
clusterTenderer.getUserAdminServiceBlockingStub();
userAdminService.updateCredentials(
UpdateCredentialsRequest.newBuilder()
.setCredentials(
toGrpcCredentials(
(com.aerospike.vector.client.auth.PasswordCredentials)
credentials))
.build());
}
/**
* Drop a user.
*
* @param username the username to drop.
*/
@Override
public void dropUser(String username) {
UserAdminServiceGrpc.UserAdminServiceBlockingStub userAdminService =
clusterTenderer.getUserAdminServiceBlockingStub();
userAdminService.dropUser(DropUserRequest.newBuilder().setUsername(username).build());
}
/**
* List all users.
*
* @return a list of all users
*/
@Override
public List userList() {
UserAdminServiceGrpc.UserAdminServiceBlockingStub userAdminService =
clusterTenderer.getUserAdminServiceBlockingStub();
return userAdminService.listUsers(Empty.getDefaultInstance()).getUsersList();
}
/**
* Grant roles to a user.
*
* @param username the username to grant roles to
* @param roles the roles to grant
*/
@Override
public void grantRoles(String username, Set roles) {
UserAdminServiceGrpc.UserAdminServiceBlockingStub userAdminService =
clusterTenderer.getUserAdminServiceBlockingStub();
userAdminService.grantRoles(
GrantRolesRequest.newBuilder().setUsername(username).addAllRoles(roles).build());
}
/**
* Revoke roles from a user.
*
* @param username the username to revoke roles from
* @param roles the roles to revoke
*/
@Override
public void revokeRoles(String username, Set roles) {
UserAdminServiceGrpc.UserAdminServiceBlockingStub userAdminService =
clusterTenderer.getUserAdminServiceBlockingStub();
userAdminService.revokeRoles(
RevokeRolesRequest.newBuilder().setUsername(username).addAllRoles(roles).build());
}
/**
* List all roles.
*
* @return a list of all roles
*/
@Override
public List roleList() {
UserAdminServiceGrpc.UserAdminServiceBlockingStub userAdminService =
clusterTenderer.getUserAdminServiceBlockingStub();
return userAdminService.listRoles(Empty.getDefaultInstance()).getRolesList();
}
private Credentials toGrpcCredentials(
com.aerospike.vector.client.auth.PasswordCredentials credentials) {
if (credentials == null) {
return null; // or handle it based on your nullability requirements
}
// Create a new builder for PasswordCredentials
PasswordCredentials passwordCredentials =
PasswordCredentials.newBuilder().setPassword(credentials.password()).build();
// Create a new builder for Credentials and set the PasswordCredentials
Credentials grpcCredentials =
Credentials.newBuilder()
.setUsername(credentials.username())
.setPasswordCredentials(passwordCredentials)
.build();
return grpcCredentials;
}
/** {@inheritDoc} */
@Override
public void close() {
try {
if (clusterTenderer != null) {
clusterTenderer.close();
}
if (adminExecutorService != null && !adminExecutorService.isShutdown()) {
adminExecutorService.shutdown();
}
} catch (Exception e) {
throw new RuntimeException("Failed to close resources properly", e);
}
}
private void waitForIndexCreation(IndexId indexId, long timeoutInMillis, long waitTimeInMillis)
throws InterruptedException {
Callable task =
() -> {
long startTime = System.currentTimeMillis();
while (System.currentTimeMillis() - startTime < timeoutInMillis) {
try {
if (getIndex(indexId, true) != null) {
return (java.lang.Boolean) true; // Index is found
}
} catch (StatusRuntimeException e) {
// Handle the case where the index is not found yet
Thread.sleep(waitTimeInMillis);
}
}
return (java.lang.Boolean) false;
};
Future future = adminExecutorService.submit(task);
try {
if (!future.get()) {
throw new TimeoutException(
"Failed to verify index creation within the timeout period.");
}
} catch (InterruptedException | ExecutionException | TimeoutException e) {
throw new RuntimeException("Error waiting for index to be created", e);
}
// Not shutting down executor because other tasks may be using it.
}
private void waitForIndexDeletion(
IndexServiceGrpc.IndexServiceBlockingStub indexService,
IndexId indexId,
long timeoutInMillis,
long waitTime)
throws InterruptedException {
long startTime = System.currentTimeMillis();
while (System.currentTimeMillis() - startTime < timeoutInMillis) {
// if index is not in the list then it was deleted
boolean indexExists =
indexService
.list(IndexListRequest.getDefaultInstance())
.getIndicesList()
.stream()
.anyMatch(i -> i.getId().equals(indexId));
if (!indexExists) {
return;
} else {
Thread.sleep(waitTime);
}
}
throw new RuntimeException(String.format("Timed out in %s index deletion", indexId));
}
}