From cc5145a0d067656be93a3f1655eea315dbb24895 Mon Sep 17 00:00:00 2001
From: Anush <anush@Anushs-Air-M2.local>
Date: Thu, 14 Dec 2023 13:59:27 +0530
Subject: [PATCH] feat: shard key ops

---
 .../java/io/qdrant/client/QdrantClient.java   | 54 +++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/src/main/java/io/qdrant/client/QdrantClient.java b/src/main/java/io/qdrant/client/QdrantClient.java
index 680f29e..918a434 100644
--- a/src/main/java/io/qdrant/client/QdrantClient.java
+++ b/src/main/java/io/qdrant/client/QdrantClient.java
@@ -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;
@@ -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
 
 	/**