com.aerospike.vector.client.dbclient.Client Maven / Gradle / Ivy
Show all versions of avs-client-java Show documentation
/* (C)2024 */
package com.aerospike.vector.client.dbclient;
import com.aerospike.vector.client.*;
import com.aerospike.vector.client.Projection;
import com.aerospike.vector.client.VectorSearchQuery;
import com.aerospike.vector.client.internal.*;
import com.aerospike.vector.client.proto.*;
import com.aerospike.vector.client.proto.Vector;
import com.google.common.base.Preconditions;
import com.google.protobuf.Empty;
import io.grpc.Deadline;
import io.grpc.StatusRuntimeException;
import io.grpc.stub.StreamObserver;
import java.io.Closeable;
import java.util.*;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* This class implements the {@code IClient} and {@code IAdminClient} interfaces using gRPC for
* communication with a vector database.
*
* The client provides the following functionalities:
*
*
* - Index management
*
- User management
*
- Vector CRUD (Create, Read, Update, Delete) operations
*
- Search capabilities
*
*/
public class Client implements IClient, IAdminClient, Closeable {
public long WAIT_TIME_IN_MILLIS = 100L; // 0.1 second sleep between index status checks
private static final Logger log = LoggerFactory.getLogger(Client.class);
private final ExecutorService clientExecutor =
Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2);
private final ClusterTenderer clusterTenderer;
/**
* Constructs a Client to manage database operations.
*
* @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 Client(ConnectionConfig connectionConfig) {
Preconditions.checkArgument(connectionConfig != null, "connectionConfig must not be null.");
this.clusterTenderer = new ClusterTenderer(connectionConfig, "dbclient");
}
/**
* {@inheritDoc}
*
* @param namespace the database namespace where the data will be stored.
* @param set the set within the namespace; must not be null. Represents a subgroup within the
* namespace where the data item will be placed.
* @param key the key used to uniquely identify the data item within the set.
* @param fields a map of bin names to their respective values, representing the data to be
* stored.
* @param ignoreMemQueueFull if true, ignores the RESOURCE_EXHAUSTED error on vector record.
* write and silently ignores the indexing queue full error. If this option is used, the
* healer becomes responsible for indexing this record later.
* @param writeType the type of write operation to be performed (see {@link WriteType}), null
* defaults to UPSERT.
*/
@Override
public void put(
String namespace,
@Nullable String set,
Object key,
Map fields,
boolean ignoreMemQueueFull,
@Nullable WriteType writeType) {
Preconditions.checkArgument(
namespace != null && !namespace.isEmpty(), "namespace must not be empty or null.");
Preconditions.checkArgument(key != null, "key must not be null.");
Preconditions.checkArgument(
fields != null && !fields.isEmpty(), "fields must not be null.");
TransactServiceGrpc.TransactServiceBlockingStub transactService =
clusterTenderer.getTransactBlockingStub();
transactService.put(
buildPutRequest(namespace, set, key, fields, ignoreMemQueueFull, writeType));
}
/**
* {@inheritDoc}
*
* @param namespace the database namespace where the data will be stored.
* @param set the set within the namespace; must not be null. Represents a subgroup within the
* namespace where the data item will be placed.
* @param key the unique identifier for the data item within the set.
* @param fields a map of bin names to their respective values, representing the data to be
* stored.
* @param ignoreMemQueueFull if true, ignores the RESOURCE_EXHAUSTED error on vector record
* write and silently ignores the indexing queue full error. If this option is used, the
* healer becomes responsible for indexing this record later.
* @param writeType the type of write operation to be performed (see {@link WriteType}).
* @throws IllegalArgumentException if the set is null.
*/
@Override
public void putAsync(
String namespace,
@Nullable String set,
Object key,
Map fields,
boolean ignoreMemQueueFull,
WriteType writeType) {
Preconditions.checkArgument(
namespace != null && !namespace.isEmpty(), "namespace must not be empty or null.");
Preconditions.checkArgument(key != null, "key must not be null.");
Preconditions.checkArgument(fields != null && !fields.isEmpty(), "bins must not be null.");
Preconditions.checkArgument(writeType != null, "writeType must not be null.");
TransactServiceGrpc.TransactServiceStub transactService =
clusterTenderer.getTransactNonBlockingStub().withExecutor(clientExecutor);
transactService.put(
buildPutRequest(namespace, set, key, fields, ignoreMemQueueFull, writeType),
new StreamObserver<>() {
@Override
public void onNext(Empty empty) {}
@Override
public void onError(Throwable throwable) {
throw new RuntimeException("Error in putAsync.", throwable);
}
@Override
public void onCompleted() {}
});
}
private PutRequest buildPutRequest(
String namespace,
@Nullable String set,
Object key,
Map fields,
boolean ignoreMemQueueFull,
@Nullable WriteType writeType) {
Key grpcKey = Conversions.buildKey(namespace, set, key);
List grpcFields =
fields.entrySet().stream()
.map(entry -> Conversions.buildField(entry.getKey(), entry.getValue()))
.collect(Collectors.toList());
PutRequest.Builder builder =
PutRequest.newBuilder()
.setKey(grpcKey)
.setIgnoreMemQueueFull(ignoreMemQueueFull)
.addAllFields(grpcFields);
if (writeType != null) {
builder.setWriteTypeValue(writeType.getNumber());
}
return builder.build();
}
/**
* {@inheritDoc}
*
* @param namespace the database namespace.
* @param set the set within the namespace, can be null.
* @param key the key for the data item.
* @param projection specifies projections to be included in the result.
* @return a map of bin names to their respective values, null if record was not found.
*/
@Override
public Map get(
String namespace, @Nullable String set, Object key, @Nullable Projection projection) {
Preconditions.checkArgument(
namespace != null && !namespace.isEmpty(), "namespace must not be empty or null.");
Preconditions.checkArgument(key != null, "key must not be null.");
if (projection == null) {
projection = Projection.DEFAULT;
}
TransactServiceGrpc.TransactServiceBlockingStub transactService =
clusterTenderer.getTransactBlockingStub();
GetRequest getRequest =
GetRequest.newBuilder()
.setKey(Conversions.buildKey(namespace, set, key))
.setProjection(projection.toProjectionSpec())
.build();
try {
com.aerospike.vector.client.proto.Record response = transactService.get(getRequest);
return response.getFieldsList().stream()
.collect(Collectors.toMap(Field::getName, Field::getValue));
} catch (io.grpc.StatusRuntimeException e) {
if (e.getStatus().getCode() == io.grpc.Status.Code.NOT_FOUND) {
// Key not found; return null
return null;
} else {
// Re-throw other exceptions
throw e;
}
}
}
/** {@inheritDoc} */
@Override
public boolean exists(String namespace, @Nullable String set, Object key) {
Preconditions.checkArgument(
namespace != null && !namespace.isEmpty(), "namespace must not be empty or null.");
Preconditions.checkArgument(key != null, "key must not be null.");
TransactServiceGrpc.TransactServiceBlockingStub transactService =
clusterTenderer.getTransactBlockingStub();
Key grpcKey = Conversions.buildKey(namespace, set, key);
return transactService
.exists(ExistsRequest.newBuilder().setKey(grpcKey).build())
.getValue();
}
/** {@inheritDoc} */
@Override
public boolean isIndexed(String namespace, @Nullable String set, Object key, String indexName) {
Preconditions.checkArgument(
namespace != null && !namespace.isEmpty(), "namespace must not be empty or null.");
Preconditions.checkArgument(key != null, "key must not be null.");
Preconditions.checkArgument(
indexName != null && !indexName.isEmpty(), "indexName must not be empty or null.");
TransactServiceGrpc.TransactServiceBlockingStub transactService =
clusterTenderer.getTransactBlockingStub();
IsIndexedRequest request =
IsIndexedRequest.newBuilder()
.setKey(Conversions.buildKey(namespace, set, key))
.setIndexId(
IndexId.newBuilder()
.setNamespace(namespace)
.setName(indexName)
.build())
.build();
return transactService.isIndexed(request).getValue();
}
/** {@inheritDoc} */
@Override
public void waitForIndexCompletion(IndexId indexId, long timeoutMillis) {
Preconditions.checkArgument(indexId != null, "indexId must not be null.");
Preconditions.checkArgument(timeoutMillis > 0, "timeoutMillis must be >0.");
Callable task =
() -> {
long waitInterval = 20_000L; // 20 seconds
long endTime = System.currentTimeMillis() + timeoutMillis;
while (System.currentTimeMillis() < endTime) {
Thread.sleep(waitInterval);
IndexStatusResponse indexStatus = indexStatus(indexId);
if (indexStatus.getUnmergedRecordCount() == 0) {
return null; // Indexing completed
}
}
throw new TimeoutException(
"Indexing did not complete within the allotted time.");
};
Future future = clientExecutor.submit(task);
try {
future.get(timeoutMillis, TimeUnit.MILLISECONDS);
} catch (InterruptedException | ExecutionException | TimeoutException e) {
throw new RuntimeException("Error waiting for index completion.", e);
}
}
/** {@inheritDoc} */
@Override
public List vectorSearch(VectorSearchQuery query) {
Preconditions.checkArgument(query != null, "query must not be null.");
TransactServiceGrpc.TransactServiceBlockingStub transactService =
clusterTenderer.getTransactBlockingStub();
if (query.getTimeout() != Integer.MAX_VALUE) {
transactService =
transactService.withDeadline(
Deadline.after(query.getTimeout(), TimeUnit.MILLISECONDS));
}
IndexId indexId =
IndexId.newBuilder()
.setNamespace(query.getNamespace())
.setName(query.getIndexName())
.build();
VectorSearchRequest.Builder requestBuilder =
VectorSearchRequest.newBuilder()
.setIndex(indexId)
.setQueryVector(query.getVector())
.setLimit(query.getLimit())
.setProjection(
query.toVectorSearchRequest().getProjection()); // set timeout
if (query.getSearchParams() != null) {
requestBuilder.setHnswSearchParams(query.getSearchParams());
}
VectorSearchRequest request = requestBuilder.build();
Iterator response = transactService.vectorSearch(request);
List neighbors = new ArrayList<>();
while (response.hasNext()) {
neighbors.add(response.next());
}
return neighbors;
}
/*
{@inheritDoc}
*/
@Override
public void delete(String namespace, @Nullable String set, Object key) {
Preconditions.checkArgument(
namespace != null && !namespace.isEmpty(), "namespace must not be empty or null.");
Preconditions.checkArgument(key != null, "key must not be null.");
TransactServiceGrpc.TransactServiceBlockingStub transactService =
clusterTenderer.getTransactBlockingStub();
Key grpcKey = Conversions.buildKey(namespace, set, key);
DeleteRequest deleteRequest = DeleteRequest.newBuilder().setKey(grpcKey).build();
transactService.delete(deleteRequest);
}
/**
* Performs an asynchronous vector search, passing results back through the provided listener.
*
* @param listener Listener to handle asynchronous results.
* @param query Query Vector {@code } for the search.
*/
@Override
public void vectorSearchAsync(VectorSearchQuery query, VectorSearchListener listener) {
Preconditions.checkArgument(query != null, "query must not be null.");
Preconditions.checkArgument(listener != null, "listener must not be null.");
TransactServiceGrpc.TransactServiceStub transactService =
clusterTenderer.getTransactNonBlockingStub().withExecutor(clientExecutor);
if (query.getTimeout() != Integer.MAX_VALUE) {
transactService =
transactService.withDeadline(
Deadline.after(query.getTimeout(), TimeUnit.MILLISECONDS));
}
transactService.vectorSearch(
query.toVectorSearchRequest(),
new StreamObserver<>() {
@Override
public void onNext(Neighbor result) {
listener.onNext(result);
}
@Override
public void onError(Throwable t) {
listener.onError(t);
}
@Override
public void onCompleted() {
listener.onComplete();
}
});
}
/**
* {@inheritDoc}
*
* @param query KeySearchQuery object.
* @param vectorField field name which contains vector information.
* @return list of neighbors based retrieved based on the query search criteria.
*/
@Override
public List searchByKey(KeySearchQuery query, String vectorField) {
Preconditions.checkArgument(query != null, "query must not be null.");
Preconditions.checkArgument(
vectorField != null && !vectorField.isBlank(), "vectorField must not be empty.");
Map keyResult =
get(
query.getNamespace(),
query.getSet(),
query.getUserKey(),
query.getProjection());
log.debug("keyResult:{}", keyResult);
// check if the key exists or not.
if (keyResult == null) {
return null;
}
Preconditions.checkArgument(
keyResult.containsKey(vectorField), "requested record is missing vectorField.");
Vector vector =
((com.aerospike.vector.client.proto.Value) keyResult.get(vectorField))
.getVectorValue();
Vector vec;
if (vector.hasFloatData()) {
List f = vector.getFloatData().getValueList();
float[] floatArray = new float[f.size()];
for (int i = 0; i < f.size(); i++) {
floatArray[i] = f.get(i);
}
vec = Conversions.buildVectorValue(floatArray);
} else {
List f = vector.getBoolData().getValueList();
boolean[] booleanArray = new boolean[f.size()];
for (int i = 0; i < f.size(); i++) {
booleanArray[i] = f.get(i);
}
vec = Conversions.buildVectorValue(booleanArray);
}
VectorSearchRequest vectorSearchRequest = query.toVectorSearchRequest(vec);
TransactServiceGrpc.TransactServiceBlockingStub transactService =
clusterTenderer.getTransactBlockingStub();
Iterator response = transactService.vectorSearch(vectorSearchRequest);
List neighbors = new ArrayList<>();
while (response.hasNext()) {
neighbors.add(response.next());
}
return neighbors;
}
@Override
public void close() {
try {
clusterTenderer.close();
clientExecutor.shutdown();
if (!clientExecutor.awaitTermination(30, TimeUnit.SECONDS)) {
clientExecutor.shutdownNow();
}
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
log.error("Interrupted while closing.", e);
}
}
/**
* {@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.
*
* @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) {
Preconditions.checkArgument(indexId != null, "Index ID must not be null.");
Preconditions.checkArgument(
vectorBinName != null && !vectorBinName.isBlank(),
"Vector bin name must not be null.");
Preconditions.checkArgument(
vectorDistanceMetric != null, "Vector distance metric must not be null.");
Preconditions.checkArgument(dimensions > 0, "Dimensions must be a positive integer.");
Preconditions.checkArgument(
timeoutInMillis > 0, "timeoutInMillis 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);
} catch (InterruptedException e) {
throw new RuntimeException(e);
}
}
/**
* Update an existing index.
*
* @param indexId unique identifier for the index.
* @param maxMemQueueSize Maximum size of in-memory queue for inserted/updated vector records.
* @param batchingParams Updated batching behaviour for batch based index update.
* @param cachingParams Updated Hnsw index caching configuration.
* @param healerParams Updated Hnsw index cache.
* @param mergeParams Updated Hnsw batch index params for merging into main index.
* @param labels updated Hnsw index 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 {
Preconditions.checkArgument(indexId != null, "indexId must not be null.");
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.setIndexCachingParams(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());
}
/**
* {@inheritDoc}
*
* @param indexId Unique identifier for the index to be dropped.
* @param timeoutInMillis The maximum time in milliseconds that the method will wait for the
* index to get deleted.
* @throws RuntimeException if there is a runtime error during index deletion.
*/
@Override
public void indexDrop(IndexId indexId, long timeoutInMillis) {
Preconditions.checkArgument(indexId != null, "indexId must not be null.");
Preconditions.checkArgument(timeoutInMillis > 0, "timeoutInMillis must be >0.");
IndexServiceGrpc.IndexServiceBlockingStub indexService =
clusterTenderer.getIndexServiceBlockingStub();
indexService.drop(IndexDropRequest.newBuilder().setIndexId(indexId).build());
try {
waitForIndexDeletion(indexService, indexId, timeoutInMillis);
} 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) {
Preconditions.checkArgument(indexId != null, "indexId must not be null.");
Preconditions.checkArgument(cutoffTimestamp >= 0, "cutoffTimestamp must be >=0.");
IndexServiceGrpc.IndexServiceBlockingStub indexService =
clusterTenderer.getIndexServiceBlockingStub();
GcInvalidVerticesRequest request =
GcInvalidVerticesRequest.newBuilder()
.setIndexId(indexId)
.setCutoffTimestamp(cutoffTimestamp)
.build();
indexService.gcInvalidVertices(request);
}
/**
* List all available vector index.
*
* @param applyDefaults : Apply default values to parameters which are not set by user.
* @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) {
Preconditions.checkArgument(indexId != null, "indexId must not be null.");
IndexServiceGrpc.IndexServiceBlockingStub indexService =
clusterTenderer.getIndexServiceBlockingStub();
return indexService.getStatus(IndexStatusRequest.newBuilder().setIndexId(indexId).build());
}
/**
* {@inheritDoc}
*
* @param indexId Unique identifier for the index.
* @param applyDefaults
* @return index definition
*/
@Override
public IndexDefinition getIndex(IndexId indexId, boolean applyDefaults) {
Preconditions.checkArgument(indexId != null, "indexId must not be null.");
IndexServiceGrpc.IndexServiceBlockingStub indexService =
clusterTenderer.getIndexServiceBlockingStub();
return indexService.get(
IndexGetRequest.newBuilder()
.setIndexId(indexId)
.setApplyDefaults(applyDefaults)
.build());
}
/**
* Get information about a user.
*
* @param username username for which information should be fetched.
* @return user information
*/
@Override
public User getUser(String username) {
Preconditions.checkArgument(
username != null && !username.isEmpty(), "username must not be empty.");
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) {
Preconditions.checkArgument(credentials != null, "credentials must not be null.");
Preconditions.checkArgument(roles != null && !roles.isEmpty(), "roles must not be empty.");
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) {
Preconditions.checkArgument(credentials != null, "credentials must not be null.");
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) {
Preconditions.checkArgument(
username != null && !username.isEmpty(), "username must not be empty.");
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) {
Preconditions.checkArgument(
username != null && !username.isEmpty(), "username must not be empty.");
Preconditions.checkArgument(roles != null && !roles.isEmpty(), "roles must not be empty.");
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) {
Preconditions.checkArgument(
username != null && !username.isEmpty(), "username must not be empty.");
Preconditions.checkArgument(roles != null && !roles.isEmpty(), "roles must not be empty.");
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;
}
private void waitForIndexCreation(IndexId indexId, long timeoutInMillis)
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(WAIT_TIME_IN_MILLIS);
}
}
return (java.lang.Boolean) false;
};
Future future = clientExecutor.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)
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(WAIT_TIME_IN_MILLIS);
}
}
throw new RuntimeException(String.format("Timed out in %s index deletion", indexId));
}
}