com.aerospike.vector.client.adminclient.AdminClient Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of avs-client-java Show documentation
Show all versions of avs-client-java Show documentation
This project includes the Java client for Aerospike Vector Search for high-performance data interactions.
The newest version!
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.protobuf.Empty;
import io.grpc.StatusRuntimeException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import javax.annotation.Nullable;
import java.lang.Boolean;
import java.util.*;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;
/**
* 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.newVirtualThreadPerTaskExecutor();
}
/**
* {@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,
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();
if (indexExists(indexService, indexId)) {
return; // Index already exists, so we do nothing
}
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(builder.build());
try {
waitForIndexCreation(indexId, timeoutInMillis, waitTimeInMillis);
} catch (InterruptedException e) {
throw new RuntimeException(e);
}
}
private boolean indexExists(IndexServiceGrpc.IndexServiceBlockingStub indexService, IndexId indexId) {
// Check if index already exists
List existingIndices = indexService.list(Empty.getDefaultInstance()).getIndicesList();
return existingIndices.stream().anyMatch(def -> def.getId().equals(indexId));
}
@Override
public void indexDrop(IndexId indexId, long timeoutInMillis, long waitTimeInMillis ) {
IndexServiceGrpc.IndexServiceBlockingStub indexService = clusterTenderer.getIndexServiceBlockingStub();
indexService.drop(indexId);
try {
waitForIndexDeletion(indexService, indexId, timeoutInMillis, waitTimeInMillis);
} catch (InterruptedException e) {
throw new RuntimeException(e);
}
}
/**
* List all available vector index
* @return list all index related information
*/
@Override
public List indexList() {
IndexServiceGrpc.IndexServiceBlockingStub indexService = clusterTenderer.getIndexServiceBlockingStub();
return indexService.list(Empty.getDefaultInstance()).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(indexId);
}
@Override
public IndexDefinition getIndex(IndexId indexId) {
IndexServiceGrpc.IndexServiceBlockingStub indexService = clusterTenderer.getIndexServiceBlockingStub();
return indexService.get(indexId);
}
/**
* 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) != null) {
return true; // Index is found
}
} catch (StatusRuntimeException e) {
// Handle the case where the index is not found yet
Thread.sleep(waitTimeInMillis);
}
}
return 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
if (!indexExists(indexService, indexId) ) {
return;
}else {
Thread.sleep(waitTime);
}
}
throw new RuntimeException(String.format("Timed out in %s index deletion", indexId));
}
}