Skip to content

Commit

Permalink
feat: shard key ops
Browse files Browse the repository at this point in the history
  • Loading branch information
Anush authored and Anush committed Dec 14, 2023
1 parent 90a2c17 commit cc5145a
Showing 1 changed file with 54 additions and 0 deletions.
54 changes: 54 additions & 0 deletions src/main/java/io/qdrant/client/QdrantClient.java
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,12 @@
import static io.qdrant.client.grpc.Collections.CollectionOperationResponse;
import static io.qdrant.client.grpc.Collections.CreateAlias;
import static io.qdrant.client.grpc.Collections.CreateCollection;
import static io.qdrant.client.grpc.Collections.CreateShardKeyRequest;
import static io.qdrant.client.grpc.Collections.CreateShardKeyResponse;
import static io.qdrant.client.grpc.Collections.DeleteAlias;
import static io.qdrant.client.grpc.Collections.DeleteCollection;
import static io.qdrant.client.grpc.Collections.DeleteShardKeyRequest;
import static io.qdrant.client.grpc.Collections.DeleteShardKeyResponse;
import static io.qdrant.client.grpc.Collections.GetCollectionInfoRequest;
import static io.qdrant.client.grpc.Collections.GetCollectionInfoResponse;
import static io.qdrant.client.grpc.Collections.ListAliasesRequest;
Expand Down Expand Up @@ -665,6 +669,56 @@ public ListenableFuture<List<AliasDescription>> listAliasesAsync(@Nullable Durat

//endregion

//region ShardKey Management

/**
* Creates a shard key for a collection.
*
* @param createShardKey The request object for the operation.
* @param timeout The timeout for the call.
* @return a new instance of {@link CreateShardKeyResponse}
*/
public ListenableFuture<CreateShardKeyResponse> createShardKeyAsync(CreateShardKeyRequest createShardKey, @Nullable Duration timeout) {
String collectionName = createShardKey.getCollectionName();
Preconditions.checkArgument(!collectionName.isEmpty(), "Collection name must not be empty");
logger.debug("Create shard key'{}'", collectionName);

ListenableFuture<CreateShardKeyResponse> future = getCollections(timeout).createShardKey(createShardKey);
addLogFailureCallback(future, "Create shard key");
return Futures.transform(future, response -> {
if (!response.getResult()) {
logger.error("Shard key could not be created for '{}'", collectionName);
throw new QdrantException("Shard key could not be created for '" + collectionName);
}
return response;
}, MoreExecutors.directExecutor());
}

/**
* Deletes a shard key for a collection.
*
* @param createShardKey The request object for the operation.
* @param timeout The timeout for the call.
* @return a new instance of {@link DeleteShardKeyResponse}
*/
public ListenableFuture<DeleteShardKeyResponse> deleteShardKeyAsync(DeleteShardKeyRequest deleteShardKey, @Nullable Duration timeout) {
String collectionName = deleteShardKey.getCollectionName();
Preconditions.checkArgument(!collectionName.isEmpty(), "Collection name must not be empty");
logger.debug("Delete shard key'{}'", collectionName);

ListenableFuture<DeleteShardKeyResponse> future = getCollections(timeout).deleteShardKey(deleteShardKey);
addLogFailureCallback(future, "Delete shard key");
return Futures.transform(future, response -> {
if (!response.getResult()) {
logger.error("Shard key could not be deleted for '{}'", collectionName);
throw new QdrantException("Shard key could not be deleted for '" + collectionName);
}
return response;
}, MoreExecutors.directExecutor());
}

//endregion

//region Point Management

/**
Expand Down

0 comments on commit cc5145a

Please sign in to comment.