From ff18322bc5af0cc2f1bdd8d79edf1128fe8c740c Mon Sep 17 00:00:00 2001 From: TonyTonyCoder11 Date: Fri, 31 Jul 2026 15:15:00 +0200 Subject: [PATCH 1/3] =?UTF-8?q?M31=20=C2=B7=20Speak=20gRPC,=20without=20ch?= =?UTF-8?q?arging=20anyone=20who=20does=20not?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The seam has been in place since the first release and has had one implementation. This is the second, and what it proves is that the public API really did carry no wire-protocol knowledge: not one line of kdrant-core changed to add a protocol. KdrantGrpc(host) returns the same QdrantClient the REST factory does, over Qdrant's Collections, Points, Snapshots and Health services on 6334. The same shared contract runs against both engines, against the same server, so the choice between them is a footprint and throughput decision rather than a feature one. The seam is wider than the protocol, and that is the interesting part. QdrantTransport was shaped by Qdrant's REST API, and eleven of its operations exist over HTTP only: telemetry, metrics, issues, snapshot recovery, snapshot transfer, and the shard-scope snapshots. Each of them throws here, naming itself and naming REST. A snapshot download that returned an empty flow would be a backup that quietly does not exist, and a KdrantException would have been the wrong type: nothing failed on the wire, and no retry will make the call work. Health is the one place the two protocols disagree and both are right. Qdrant has no gRPC /healthz, /readyz or /livez; it has the standard gRPC health service, which answers one status for the node. The three probes answer from it, and keep the REST engine's false-on-failure contract rather than throwing. Retries mirror the REST engine's, because maxRetries is a client setting and an engine that ignored it would be a behaviour difference nobody asked for. Deadlines are per call rather than per channel, because requestTimeout is documented as per attempt. A REST-only build still resolves no gRPC, no protobuf and no Netty. That was the whole argument for making this opt-in, and it is asserted rather than assumed. --- .github/workflows/ci.yml | 5 +- CHANGELOG.md | 26 +- README.md | 22 +- docs/migrating-from-qdrant-client.md | 13 +- .../api/kdrant-transport-grpc.api | 66 ++ kdrant-transport-grpc/build.gradle.kts | 5 + .../transport/grpc/CollectionMapping.kt | 358 +++++++++ .../dev/kdrant/transport/grpc/GrpcErrors.kt | 84 ++ .../transport/grpc/GrpcQdrantTransport.kt | 759 ++++++++++++++++++ .../dev/kdrant/transport/grpc/KdrantGrpc.kt | 36 + .../dev/kdrant/transport/grpc/PointMapping.kt | 28 + .../dev/kdrant/transport/grpc/QueryMapping.kt | 278 +++++++ .../kdrant/transport/grpc/RequestMapping.kt | 185 +++++ .../grpc/GrpcClientContractIntegrationTest.kt | 18 + .../kdrant/transport/grpc/GrpcErrorsTest.kt | 112 +++ .../transport/grpc/GrpcQdrantTransportTest.kt | 615 ++++++++++++++ .../transport/grpc/RestOnlyOperationsTest.kt | 86 ++ kdrant-transport-rest/build.gradle.kts | 26 + 18 files changed, 2705 insertions(+), 17 deletions(-) create mode 100644 kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/CollectionMapping.kt create mode 100644 kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/GrpcErrors.kt create mode 100644 kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransport.kt create mode 100644 kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/KdrantGrpc.kt create mode 100644 kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/QueryMapping.kt create mode 100644 kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/RequestMapping.kt create mode 100644 kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcClientContractIntegrationTest.kt create mode 100644 kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcErrorsTest.kt create mode 100644 kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransportTest.kt create mode 100644 kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/RestOnlyOperationsTest.kt diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 868adfe..46ac92d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -91,10 +91,13 @@ jobs: - name: Set up Gradle uses: gradle/actions/setup-gradle@v6 + # Both engines run the same client contract from kdrant-testkit, against the same image. - name: Integration tests against ${{ matrix.qdrant }} env: QDRANT_IMAGE: ${{ matrix.qdrant }} - run: ./gradlew :kdrant-transport-rest:test --tests "*IntegrationTest*" --no-daemon --stacktrace + run: > + ./gradlew :kdrant-transport-rest:test :kdrant-transport-grpc:test + --tests "*IntegrationTest*" --no-daemon --stacktrace dependency-review: name: Dependency review diff --git a/CHANGELOG.md b/CHANGELOG.md index be59356..2adc489 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,9 +17,6 @@ All notable changes to this project are documented in this file. The format is b every query. Here Qdrant runs the search. The module depends only on Koog's stable `rag-base`, not on the `rag-vector` beta. Koog's `namespace` becomes a payload field and a filter, so one collection can hold several of them. - -### Added - - Cluster and sharding (M32), the gap the migration guide used to name as having no Kdrant equivalent. `collectionClusterInfo(name)` reads how a collection's shards are spread across peers, including the transfers in flight; `updateCollectionCluster(name, operation)` moves, replicates, aborts or drops a @@ -47,8 +44,31 @@ All notable changes to this project are documented in this file. The format is b between a backup that fits in a window and one that does not. Shard ids come from `collectionClusterInfo`. + +- **`kdrant-transport-grpc`** (M31), the opt-in gRPC engine. `KdrantGrpc(host)` returns the same + `QdrantClient` the REST factory does, over Qdrant's `Collections`, `Points`, `Snapshots` and `Health` + services on port **6334**. REST stays the recommended engine; reach for this one when throughput or + long-lived streaming is the bottleneck, which is the case the README used to concede to the official + client. Nothing changes for a REST user: the module is separate, and a build that does not ask for it + resolves no gRPC, no protobuf and no Netty. + The stubs are generated from Qdrant's own `.proto` files, vendored verbatim at v1.18.2, rather than + taken from `io.qdrant:client` — grpc-kotlin emits suspend functions and `Flow`s, which is the shape + the transport seam already has, and generating decides the dependency set instead of inheriting a + shaded Netty jar that is most of the official client's footprint. +- Both engines are held to one **shared client contract** (`kdrant-testkit`), which runs the same 30 + behavioural tests against a real Qdrant over each protocol. The REST tests that came before it + asserted HTTP bodies, which a gRPC engine cannot satisfy by construction. + ### Changed +- Eleven `QdrantTransport` operations have no gRPC equivalent, because the seam was shaped by Qdrant's + REST API and Qdrant serves these over HTTP only: `telemetry`, `metrics`, `listIssues`, `clearIssues`, + `recoverSnapshot`, snapshot download and upload, and the five shard-scope snapshot operations. On the + gRPC engine each throws an `UnsupportedOperationException` naming the operation and pointing at REST, + rather than degrading quietly — a snapshot download that returns nothing is a backup that does not + exist. The REST engine is unchanged. + + - Releases publish to Maven Central only. The secondary publication to GitHub Packages is gone: it carried the same artifacts to a registry that requires authentication even for public packages, so it was a second place to keep in sync and no second way for anyone to depend on Kdrant. Versions up to diff --git a/README.md b/README.md index 80a0d83..8ae2cdd 100644 --- a/README.md +++ b/README.md @@ -74,9 +74,10 @@ Dependency stacks verified against `io.qdrant:client:1.18.3`: | Models | `kotlinx-serialization` data classes | generated protobuf messages | | GraalVM native / cold start | friendly (no Netty/protobuf reflection config) | needs gRPC/Netty/protobuf native config; heavier cold start | -For raw throughput and streaming, gRPC/HTTP2 still wins. Reach for the official client when that is -your bottleneck. For typical RAG and embedding-search workloads, Kdrant trades it for a fraction of -the footprint and idiomatic Kotlin. +For raw throughput and streaming, gRPC/HTTP2 still wins. That case now has an answer inside Kdrant: +`kdrant-transport-grpc` is an opt-in engine behind the same `QdrantClient`, so the column above stays +true of the default and a build that does not ask for gRPC never pays for it. For typical RAG and +embedding-search workloads, REST trades the wire for a fraction of the footprint and idiomatic Kotlin. ## Installation @@ -273,15 +274,18 @@ val hits = catching { qdrant.search("articles") { query(queryVector) } }.getOrEl ## Architecture -Two modules keep protocol concerns out of the public API: +Three modules keep protocol concerns out of the public API: | Module | Contents | | --- | --- | | `kdrant-core` | Public API (`QdrantClient`), models, DSLs, error hierarchy, and the `QdrantTransport` seam, with no wire-protocol knowledge. | | `kdrant-transport-rest` | The default REST engine (Ktor CIO) implementing `QdrantTransport`, plus the `Kdrant(...)` factory. | +| `kdrant-transport-grpc` | The opt-in gRPC engine over Qdrant's own protobuf services, plus the `KdrantGrpc(...)` factory. Depend on it only if you want it. | -The DSLs and client logic live in `kdrant-core` and are independent of the protocol; only the -engine module knows about HTTP. +The DSLs and client logic live in `kdrant-core` and are independent of the protocol; only an engine +module knows about a wire. Both engines are held to the same behavioural test suite, so the choice is +a footprint and throughput decision rather than a feature one — with the exception of the operations +Qdrant serves over HTTP only, which the gRPC engine names rather than degrading. ## Roadmap @@ -301,10 +305,8 @@ from `1.1.0` is a recompile rather than a jar swap — see [STABILITY.md](STABILITY.md#what-may-still-change-in-a-minor); the [CHANGELOG](CHANGELOG.md) has the version-by-version detail. -**Next.** Nothing is claimed for the next release yet; the board is where it gets decided. - -**Later.** Kotlin Multiplatform (`commonMain`) and an optional opt-in gRPC engine, with REST staying -the default. +**Merged, unreleased.** The opt-in gRPC engine, `kdrant-transport-grpc`. It is in `main` and has not +shipped; see the [CHANGELOG](CHANGELOG.md#unreleased) for what it changes. The plan lives on the [Kdrant board](https://github.com/orgs/NaCode-Studios/projects/4) — one item per milestone, each with its exit criterion — and every tier is a [milestone](https://github.com/NaCode-Studios/Kdrant/milestones) in this repository. See diff --git a/docs/migrating-from-qdrant-client.md b/docs/migrating-from-qdrant-client.md index 9a59758..6d1db00 100644 --- a/docs/migrating-from-qdrant-client.md +++ b/docs/migrating-from-qdrant-client.md @@ -9,11 +9,18 @@ builders and the static factory imports (`id`, `value`, `vectors`, `range`) beco Your collections, points, payloads and filters are untouched: both clients talk to the same server, so you can run them side by side against the same Qdrant and move one call site at a time. +If you want to keep gRPC, only the second and third of those changes apply. Depend on +`kdrant-transport-grpc` and build the client with `KdrantGrpc(host)` instead of `Kdrant(host)`; the port +stays `6334` and everything else in this guide reads the same, because the API above the wire is the +same API. What that engine cannot do is listed under [What this is not](#what-this-is-not). + ## What this is not -Kdrant does not speak gRPC. If your bottleneck is raw throughput or long-lived streaming, HTTP/2 is -still the faster wire and the official client is the right tool — the trade-off is set out in the -[README](../README.md#footprint-vs-the-official-client). +The gRPC engine is not the whole of Qdrant. Eleven operations are served over HTTP only — telemetry, +Prometheus metrics, the issues endpoint, snapshot recovery, snapshot download and upload, and the +shard-scope snapshots — and on the gRPC engine each throws rather than pretending. If you need any of +them, use the REST engine for the whole client or for that one call. The footprint trade-off between +the two is set out in the [README](../README.md#footprint-vs-the-official-client). Cluster support covers a collection's shard distribution: reading it, moving and replicating shards, and creating or dropping custom sharding keys. The node-level calls that administer the raft cluster diff --git a/kdrant-transport-grpc/api/kdrant-transport-grpc.api b/kdrant-transport-grpc/api/kdrant-transport-grpc.api index e69de29..9fe24f3 100644 --- a/kdrant-transport-grpc/api/kdrant-transport-grpc.api +++ b/kdrant-transport-grpc/api/kdrant-transport-grpc.api @@ -0,0 +1,66 @@ +public final class dev/kdrant/transport/grpc/GrpcQdrantTransport : dev/kdrant/transport/QdrantTransport { + public fun batchUpdate (Ljava/lang/String;Ljava/util/List;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun clearIssues (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun clearPayload (Ljava/lang/String;Ldev/kdrant/model/DeleteSelector;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun close ()V + public fun collectionClusterInfo (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun collectionExists (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun count (Ljava/lang/String;Ldev/kdrant/model/Filter;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun createCollection (Ljava/lang/String;Ldev/kdrant/model/CreateCollectionRequest;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun createPayloadIndex (Ljava/lang/String;Ljava/lang/String;Ldev/kdrant/model/PayloadSchemaType;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun createShardKey (Ljava/lang/String;Ldev/kdrant/model/CreateShardKeyRequest;Ljava/lang/Integer;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun createShardSnapshot (Ljava/lang/String;IZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun createSnapshot (Ljava/lang/String;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun createStorageSnapshot (ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun delete (Ljava/lang/String;Ldev/kdrant/model/DeleteSelector;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun deleteCollection (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun deletePayload (Ljava/lang/String;Ljava/util/List;Ldev/kdrant/model/DeleteSelector;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun deletePayloadIndex (Ljava/lang/String;Ljava/lang/String;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun deleteShardKey (Ljava/lang/String;Ldev/kdrant/model/ShardKey;Ljava/lang/Integer;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun deleteShardSnapshot (Ljava/lang/String;ILjava/lang/String;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun deleteSnapshot (Ljava/lang/String;Ljava/lang/String;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun deleteStorageSnapshot (Ljava/lang/String;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun deleteVectors (Ljava/lang/String;Ljava/util/List;Ldev/kdrant/model/DeleteSelector;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun downloadShardSnapshot (Ljava/lang/String;ILjava/lang/String;)Lkotlinx/coroutines/flow/Flow; + public fun downloadSnapshot (Ljava/lang/String;Ljava/lang/String;)Lkotlinx/coroutines/flow/Flow; + public fun downloadStorageSnapshot (Ljava/lang/String;)Lkotlinx/coroutines/flow/Flow; + public fun facet (Ljava/lang/String;Ljava/lang/String;Ldev/kdrant/model/Filter;Ljava/lang/Integer;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun getCollection (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun healthz (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun listAliases (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun listCollectionAliases (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun listCollections (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun listIssues (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun listShardSnapshots (Ljava/lang/String;ILkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun listSnapshots (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun listStorageSnapshots (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun livez (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun metrics (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun overwritePayload (Ljava/lang/String;Lkotlinx/serialization/json/JsonObject;Ldev/kdrant/model/DeleteSelector;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun query (Ljava/lang/String;Ldev/kdrant/model/SearchRequest;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun queryBatch (Ljava/lang/String;Ljava/util/List;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun queryGroups (Ljava/lang/String;Ldev/kdrant/model/SearchGroupsRequest;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun readyz (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun recoverShardSnapshot (Ljava/lang/String;ILjava/lang/String;Ldev/kdrant/model/SnapshotPriority;Ljava/lang/String;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun recoverSnapshot (Ljava/lang/String;Ljava/lang/String;Ldev/kdrant/model/SnapshotPriority;Ljava/lang/String;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun retrieve (Ljava/lang/String;Ljava/util/List;Ldev/kdrant/model/WithPayload;Ljava/lang/Boolean;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun scroll (Ljava/lang/String;Ldev/kdrant/model/ScrollRequest;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun searchMatrixOffsets (Ljava/lang/String;Ldev/kdrant/model/SearchMatrixRequest;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun searchMatrixPairs (Ljava/lang/String;Ldev/kdrant/model/SearchMatrixRequest;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun setPayload (Ljava/lang/String;Lkotlinx/serialization/json/JsonObject;Ldev/kdrant/model/DeleteSelector;Ljava/lang/String;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun telemetry (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun updateAliases (Ljava/util/List;Ljava/lang/Integer;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun updateCollection (Ljava/lang/String;Ldev/kdrant/model/UpdateCollectionRequest;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun updateCollectionCluster (Ljava/lang/String;Ldev/kdrant/model/ClusterOperation;Ljava/lang/Integer;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun updateVectors (Ljava/lang/String;Ljava/util/List;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun uploadShardSnapshot (Ljava/lang/String;ILkotlinx/coroutines/flow/Flow;Ldev/kdrant/model/SnapshotPriority;Ljava/lang/String;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun uploadSnapshot (Ljava/lang/String;Lkotlinx/coroutines/flow/Flow;Ldev/kdrant/model/SnapshotPriority;Ljava/lang/String;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun upsert (Ljava/lang/String;Ljava/util/List;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun upsert (Ljava/lang/String;Lkotlinx/coroutines/flow/Flow;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object; +} + +public final class dev/kdrant/transport/grpc/KdrantGrpcKt { + public static final fun KdrantGrpc (Ljava/lang/String;IILkotlin/jvm/functions/Function1;)Ldev/kdrant/QdrantClient; + public static synthetic fun KdrantGrpc$default (Ljava/lang/String;IILkotlin/jvm/functions/Function1;ILjava/lang/Object;)Ldev/kdrant/QdrantClient; +} + diff --git a/kdrant-transport-grpc/build.gradle.kts b/kdrant-transport-grpc/build.gradle.kts index 8436afc..a00e91a 100644 --- a/kdrant-transport-grpc/build.gradle.kts +++ b/kdrant-transport-grpc/build.gradle.kts @@ -33,6 +33,11 @@ dependencies { testRuntimeOnly(libs.junit.platform.launcher) testImplementation(libs.kotlinx.coroutines.test) testImplementation(libs.grpc.inprocess) + + // The shared client contract, run here against the same real Qdrant the REST engine is held to. + testImplementation(project(":kdrant-testkit")) + testImplementation(platform(libs.testcontainers.bom)) + testImplementation(libs.testcontainers.qdrant) } // Resolved before the protobuf { } block: inside it the extension receiver shadows the catalog accessor. diff --git a/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/CollectionMapping.kt b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/CollectionMapping.kt new file mode 100644 index 0000000..364db1c --- /dev/null +++ b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/CollectionMapping.kt @@ -0,0 +1,358 @@ +package dev.kdrant.transport.grpc + +import dev.kdrant.model.AliasOperation +import dev.kdrant.model.ClusterOperation +import dev.kdrant.model.CollectionConfig +import dev.kdrant.model.CollectionInfo +import dev.kdrant.model.CollectionParams +import dev.kdrant.model.CollectionStatus +import dev.kdrant.model.CreateCollectionRequest +import dev.kdrant.model.Distance +import dev.kdrant.model.HnswConfig +import dev.kdrant.model.Modifier +import dev.kdrant.model.MultiVectorComparator +import dev.kdrant.model.MultiVectorConfig +import dev.kdrant.model.OptimizersConfig +import dev.kdrant.model.PayloadIndexInfo +import dev.kdrant.model.PayloadSchemaType +import dev.kdrant.model.QuantizationConfig +import dev.kdrant.model.ReplicaState +import dev.kdrant.model.SnapshotDescription +import dev.kdrant.model.SparseVectorParams +import dev.kdrant.model.UpdateCollectionRequest +import dev.kdrant.model.VectorDatatype +import dev.kdrant.model.VectorParams +import dev.kdrant.model.VectorsConfig +import qdrant.Collections +import qdrant.SnapshotsService +import java.time.Instant + +/** + * Collection configuration and the read-back of it. + * + * The asymmetry worth knowing about: `CreateCollection` takes whole config messages, while + * `UpdateCollection` takes `*Diff` messages of the same shape. Kdrant's models are already diffs — + * every field is nullable and an unset field means "leave it alone" — so the two directions differ + * only in which protobuf type they build, not in what they read. + * + * Product and Turbo quantization exist on the wire and not in the model, so a collection created + * elsewhere with either reads back as no quantization rather than as something it is not. The REST + * engine drops them the same way, because the model is the same. + */ +internal object CollectionMapping { + + fun createCollection(name: String, request: CreateCollectionRequest): Collections.CreateCollection = + Collections.CreateCollection.newBuilder().apply { + collectionName = name + request.vectors?.let { vectorsConfig = vectorsConfig(it) } + request.sparseVectors?.let { sparseVectorsConfig = sparseVectorConfig(it) } + request.hnswConfig?.let { hnswConfig = hnswConfig(it) } + request.onDiskPayload?.let { onDiskPayload = it } + request.shardNumber?.let { shardNumber = it } + request.replicationFactor?.let { replicationFactor = it } + request.optimizersConfig?.let { optimizersConfig = optimizersConfig(it) } + request.quantizationConfig?.let { quantizationConfig = quantizationConfig(it) } + }.build() + + fun updateCollection(name: String, request: UpdateCollectionRequest): Collections.UpdateCollection = + Collections.UpdateCollection.newBuilder().apply { + collectionName = name + request.optimizersConfig?.let { optimizersConfig = optimizersConfig(it) } + request.hnswConfig?.let { hnswConfig = hnswConfig(it) } + request.quantizationConfig?.let { quantizationConfig = quantizationConfigDiff(it) } + }.build() + + fun collectionInfo(info: Collections.CollectionInfo): CollectionInfo = CollectionInfo( + status = status(info.status), + pointsCount = info.takeIf { it.hasPointsCount() }?.pointsCount, + indexedVectorsCount = info.takeIf { it.hasIndexedVectorsCount() }?.indexedVectorsCount, + segmentsCount = info.segmentsCount.toInt(), + config = if (info.hasConfig()) collectionConfig(info.config) else null, + payloadSchema = info.payloadSchemaMap.mapValues { (_, schema) -> payloadIndexInfo(schema) }, + ) + + fun snapshotDescription(description: SnapshotsService.SnapshotDescription): SnapshotDescription = + SnapshotDescription( + name = description.name, + // The wire carries an instant; the model carries the RFC 3339 text REST returns. + creationTime = description.takeIf { it.hasCreationTime() } + ?.creationTime + ?.let { Instant.ofEpochSecond(it.seconds, it.nanos.toLong()).toString() }, + size = description.size, + checksum = description.takeIf { it.hasChecksum() }?.checksum, + ) + + fun payloadSchemaType(type: PayloadSchemaType): Collections.PayloadSchemaType = when (type) { + PayloadSchemaType.KEYWORD -> Collections.PayloadSchemaType.Keyword + PayloadSchemaType.INTEGER -> Collections.PayloadSchemaType.Integer + PayloadSchemaType.FLOAT -> Collections.PayloadSchemaType.Float + PayloadSchemaType.GEO -> Collections.PayloadSchemaType.Geo + PayloadSchemaType.TEXT -> Collections.PayloadSchemaType.Text + PayloadSchemaType.BOOL -> Collections.PayloadSchemaType.Bool + PayloadSchemaType.DATETIME -> Collections.PayloadSchemaType.Datetime + PayloadSchemaType.UUID -> Collections.PayloadSchemaType.Uuid + } + + fun aliasOperation(operation: AliasOperation): Collections.AliasOperations = + Collections.AliasOperations.newBuilder().apply { + when (operation) { + is AliasOperation.Create -> createAlias = Collections.CreateAlias.newBuilder() + .setCollectionName(operation.collectionName) + .setAliasName(operation.aliasName) + .build() + is AliasOperation.Delete -> deleteAlias = Collections.DeleteAlias.newBuilder() + .setAliasName(operation.aliasName) + .build() + is AliasOperation.Rename -> renameAlias = Collections.RenameAlias.newBuilder() + .setOldAliasName(operation.oldAliasName) + .setNewAliasName(operation.newAliasName) + .build() + } + }.build() + + /** [timeout] is in seconds, as everywhere else on the seam. */ + fun clusterSetup( + name: String, + operation: ClusterOperation, + timeout: Int?, + ): Collections.UpdateCollectionClusterSetupRequest = + Collections.UpdateCollectionClusterSetupRequest.newBuilder().apply { + collectionName = name + timeout?.let { this.timeout = it.toLong() } + when (operation) { + is ClusterOperation.MoveShard -> moveShard = Collections.MoveShard.newBuilder() + .setShardId(operation.shardId) + .setFromPeerId(operation.fromPeerId) + .setToPeerId(operation.toPeerId) + .build() + is ClusterOperation.ReplicateShard -> replicateShard = Collections.ReplicateShard.newBuilder() + .setShardId(operation.shardId) + .setFromPeerId(operation.fromPeerId) + .setToPeerId(operation.toPeerId) + .build() + is ClusterOperation.AbortTransfer -> abortTransfer = Collections.AbortShardTransfer.newBuilder() + .setShardId(operation.shardId) + .setFromPeerId(operation.fromPeerId) + .setToPeerId(operation.toPeerId) + .build() + is ClusterOperation.DropReplica -> dropReplica = Collections.Replica.newBuilder() + .setShardId(operation.shardId) + .setPeerId(operation.peerId) + .build() + } + }.build() + + fun replicaState(state: Collections.ReplicaState): ReplicaState = when (state) { + Collections.ReplicaState.Active -> ReplicaState.ACTIVE + Collections.ReplicaState.Dead -> ReplicaState.DEAD + Collections.ReplicaState.Partial -> ReplicaState.PARTIAL + Collections.ReplicaState.Initializing -> ReplicaState.INITIALIZING + Collections.ReplicaState.Listener -> ReplicaState.LISTENER + Collections.ReplicaState.PartialSnapshot -> ReplicaState.PARTIAL_SNAPSHOT + Collections.ReplicaState.Recovery -> ReplicaState.RECOVERY + Collections.ReplicaState.Resharding -> ReplicaState.RESHARDING + Collections.ReplicaState.ReshardingScaleDown -> ReplicaState.RESHARDING_SCALE_DOWN + Collections.ReplicaState.ActiveRead -> ReplicaState.ACTIVE_READ + // A state a newer Qdrant added. The model already decodes an unrecognized state to UNKNOWN + // rather than failing the whole cluster-info response; the same tolerance applies here. + else -> ReplicaState.UNKNOWN + } + + private fun collectionConfig(config: Collections.CollectionConfig): CollectionConfig = CollectionConfig( + params = if (config.hasParams()) collectionParams(config.params) else null, + ) + + private fun collectionParams(params: Collections.CollectionParams): CollectionParams = CollectionParams( + vectors = if (params.hasVectorsConfig()) vectorsConfigToModel(params.vectorsConfig) else null, + sparseVectors = params.takeIf { it.hasSparseVectorsConfig() } + ?.sparseVectorsConfig + ?.mapMap + ?.mapValues { (_, value) -> sparseVectorParamsToModel(value) }, + shardNumber = params.shardNumber, + replicationFactor = params.takeIf { it.hasReplicationFactor() }?.replicationFactor, + writeConsistencyFactor = params.takeIf { it.hasWriteConsistencyFactor() }?.writeConsistencyFactor, + onDiskPayload = params.onDiskPayload, + ) + + /** + * The index type is kept as its wire string rather than as the enum, matching the REST engine: an + * index type a newer Qdrant adds decodes to a name this client does not know instead of failing + * the whole `getCollection` response. + */ + private fun payloadIndexInfo(schema: Collections.PayloadSchemaInfo): PayloadIndexInfo = PayloadIndexInfo( + dataType = schema.dataType.name.lowercase().takeUnless { it == "unknowntype" }, + points = schema.takeIf { it.hasPoints() }?.points, + ) + + private fun status(status: Collections.CollectionStatus): CollectionStatus = when (status) { + Collections.CollectionStatus.Green -> CollectionStatus.GREEN + Collections.CollectionStatus.Yellow -> CollectionStatus.YELLOW + Collections.CollectionStatus.Red -> CollectionStatus.RED + Collections.CollectionStatus.Grey -> CollectionStatus.GREY + else -> CollectionStatus.UNKNOWN + } + + private fun vectorsConfig(config: VectorsConfig): Collections.VectorsConfig = + Collections.VectorsConfig.newBuilder().apply { + when (config) { + is VectorsConfig.Single -> params = vectorParams(config.params) + is VectorsConfig.Named -> paramsMap = Collections.VectorParamsMap.newBuilder() + .putAllMap(config.vectors.mapValues { (_, params) -> vectorParams(params) }) + .build() + } + }.build() + + private fun vectorsConfigToModel(config: Collections.VectorsConfig): VectorsConfig? = + when (config.configCase) { + Collections.VectorsConfig.ConfigCase.PARAMS -> VectorsConfig.Single(vectorParamsToModel(config.params)) + Collections.VectorsConfig.ConfigCase.PARAMS_MAP -> VectorsConfig.Named( + config.paramsMap.mapMap.mapValues { (_, params) -> vectorParamsToModel(params) }, + ) + Collections.VectorsConfig.ConfigCase.CONFIG_NOT_SET, null -> null + } + + private fun vectorParams(params: VectorParams): Collections.VectorParams = + Collections.VectorParams.newBuilder().apply { + size = params.size + distance = distance(params.distance) + params.onDisk?.let { onDisk = it } + params.datatype?.let { datatype = datatype(it) } + params.hnswConfig?.let { hnswConfig = hnswConfig(it) } + params.multivectorConfig?.let { multivectorConfig = multiVectorConfig(it) } + }.build() + + private fun vectorParamsToModel(params: Collections.VectorParams): VectorParams = VectorParams( + size = params.size, + distance = distanceToModel(params.distance), + onDisk = params.takeIf { it.hasOnDisk() }?.onDisk, + datatype = params.takeIf { it.hasDatatype() }?.datatype?.let(::datatypeToModel), + hnswConfig = params.takeIf { it.hasHnswConfig() }?.hnswConfig?.let(::hnswConfigToModel), + multivectorConfig = params.takeIf { it.hasMultivectorConfig() }?.let { + MultiVectorConfig(MultiVectorComparator.MAX_SIM) + }, + ) + + private fun sparseVectorConfig(vectors: Map): Collections.SparseVectorConfig = + Collections.SparseVectorConfig.newBuilder() + .putAllMap( + vectors.mapValues { (_, params) -> + Collections.SparseVectorParams.newBuilder().apply { + params.modifier?.let { modifier = modifier(it) } + }.build() + }, + ) + .build() + + private fun sparseVectorParamsToModel(params: Collections.SparseVectorParams): SparseVectorParams = + SparseVectorParams( + modifier = params.takeIf { it.hasModifier() }?.modifier?.let { + if (it == Collections.Modifier.Idf) Modifier.IDF else Modifier.NONE + }, + ) + + private fun modifier(modifier: Modifier): Collections.Modifier = when (modifier) { + Modifier.NONE -> Collections.Modifier.None + Modifier.IDF -> Collections.Modifier.Idf + } + + private fun multiVectorConfig(config: MultiVectorConfig): Collections.MultiVectorConfig = + Collections.MultiVectorConfig.newBuilder() + .setComparator( + when (config.comparator) { + MultiVectorComparator.MAX_SIM -> Collections.MultiVectorComparator.MaxSim + }, + ) + .build() + + private fun distance(distance: Distance): Collections.Distance = when (distance) { + Distance.COSINE -> Collections.Distance.Cosine + Distance.DOT -> Collections.Distance.Dot + Distance.EUCLID -> Collections.Distance.Euclid + Distance.MANHATTAN -> Collections.Distance.Manhattan + } + + private fun distanceToModel(distance: Collections.Distance): Distance = when (distance) { + Collections.Distance.Cosine -> Distance.COSINE + Collections.Distance.Dot -> Distance.DOT + Collections.Distance.Euclid -> Distance.EUCLID + Collections.Distance.Manhattan -> Distance.MANHATTAN + // A distance this client does not know reads as cosine rather than failing the response. The + // model has no unknown variant, and the field is a read-back of a collection someone else made. + else -> Distance.COSINE + } + + private fun datatype(datatype: VectorDatatype): Collections.Datatype = when (datatype) { + VectorDatatype.FLOAT32 -> Collections.Datatype.Float32 + VectorDatatype.UINT8 -> Collections.Datatype.Uint8 + VectorDatatype.FLOAT16 -> Collections.Datatype.Float16 + } + + private fun datatypeToModel(datatype: Collections.Datatype): VectorDatatype? = when (datatype) { + Collections.Datatype.Float32 -> VectorDatatype.FLOAT32 + Collections.Datatype.Uint8 -> VectorDatatype.UINT8 + Collections.Datatype.Float16 -> VectorDatatype.FLOAT16 + else -> null + } + + private fun hnswConfig(config: HnswConfig): Collections.HnswConfigDiff = + Collections.HnswConfigDiff.newBuilder().apply { + config.m?.let { m = it.toLong() } + config.efConstruct?.let { efConstruct = it.toLong() } + config.fullScanThreshold?.let { fullScanThreshold = it.toLong() } + config.maxIndexingThreads?.let { maxIndexingThreads = it.toLong() } + config.onDisk?.let { onDisk = it } + config.payloadM?.let { payloadM = it.toLong() } + }.build() + + private fun hnswConfigToModel(config: Collections.HnswConfigDiff): HnswConfig = HnswConfig( + m = config.takeIf { it.hasM() }?.m?.toInt(), + efConstruct = config.takeIf { it.hasEfConstruct() }?.efConstruct?.toInt(), + fullScanThreshold = config.takeIf { it.hasFullScanThreshold() }?.fullScanThreshold?.toInt(), + maxIndexingThreads = config.takeIf { it.hasMaxIndexingThreads() }?.maxIndexingThreads?.toInt(), + onDisk = config.takeIf { it.hasOnDisk() }?.onDisk, + payloadM = config.takeIf { it.hasPayloadM() }?.payloadM?.toInt(), + ) + + private fun optimizersConfig(config: OptimizersConfig): Collections.OptimizersConfigDiff = + Collections.OptimizersConfigDiff.newBuilder().apply { + config.deletedThreshold?.let { deletedThreshold = it } + config.vacuumMinVectorNumber?.let { vacuumMinVectorNumber = it.toLong() } + config.defaultSegmentNumber?.let { defaultSegmentNumber = it.toLong() } + config.maxSegmentSize?.let { maxSegmentSize = it.toLong() } + config.memmapThreshold?.let { memmapThreshold = it.toLong() } + config.indexingThreshold?.let { indexingThreshold = it.toLong() } + config.flushIntervalSec?.let { flushIntervalSec = it.toLong() } + config.maxOptimizationThreads?.let { + maxOptimizationThreads = Collections.MaxOptimizationThreads.newBuilder().setValue(it.toLong()).build() + } + }.build() + + private fun quantizationConfig(config: QuantizationConfig): Collections.QuantizationConfig = + Collections.QuantizationConfig.newBuilder().apply { + when (config) { + is QuantizationConfig.Scalar -> scalar = scalarQuantization(config) + is QuantizationConfig.Binary -> binary = binaryQuantization(config) + } + }.build() + + private fun quantizationConfigDiff(config: QuantizationConfig): Collections.QuantizationConfigDiff = + Collections.QuantizationConfigDiff.newBuilder().apply { + when (config) { + is QuantizationConfig.Scalar -> scalar = scalarQuantization(config) + is QuantizationConfig.Binary -> binary = binaryQuantization(config) + } + }.build() + + private fun scalarQuantization(config: QuantizationConfig.Scalar): Collections.ScalarQuantization = + Collections.ScalarQuantization.newBuilder().apply { + // The model has one scalar type, the same `int8` the REST serializer writes. + type = Collections.QuantizationType.Int8 + config.quantile?.let { quantile = it } + config.alwaysRam?.let { alwaysRam = it } + }.build() + + private fun binaryQuantization(config: QuantizationConfig.Binary): Collections.BinaryQuantization = + Collections.BinaryQuantization.newBuilder().apply { + config.alwaysRam?.let { alwaysRam = it } + }.build() +} diff --git a/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/GrpcErrors.kt b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/GrpcErrors.kt new file mode 100644 index 0000000..263ba01 --- /dev/null +++ b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/GrpcErrors.kt @@ -0,0 +1,84 @@ +package dev.kdrant.transport.grpc + +import dev.kdrant.KdrantException +import io.grpc.Status +import io.grpc.StatusException +import io.grpc.StatusRuntimeException +import kotlinx.coroutines.CancellationException + +/** + * gRPC status codes onto the same [KdrantException] hierarchy the REST engine raises, so a caller can + * swap engines without rewriting a `when`. + * + * Two codes need more than a table. `NOT_FOUND` is what Qdrant answers for a missing collection, but + * it is also what an unimplemented path would answer, so the collection name is carried in rather than + * parsed out of the message. And `INVALID_ARGUMENT` is what Qdrant returns for a collection that does + * not exist on some write paths, which is why the message is checked as well as the code. + * + * [CancellationException] is re-thrown untouched. gRPC reports a cancelled call as `Status.CANCELLED`, + * and turning that into a `KdrantException` would break structured concurrency: the coroutine that was + * cancelled would see a failure instead of its own cancellation. + */ +internal object GrpcErrors { + + /** Runs [block], translating any gRPC failure. [collection] names the collection the call was about. */ + suspend fun mapping(collection: String?, block: suspend () -> T): T = + try { + block() + } catch (e: CancellationException) { + throw e + } catch (e: StatusRuntimeException) { + throw translate(e.status, e.status.description, collection, e) + } catch (e: StatusException) { + throw translate(e.status, e.status.description, collection, e) + } + + private fun translate( + status: Status, + description: String?, + collection: String?, + cause: Throwable, + ): KdrantException { + val message = description.orEmpty() + return when (status.code) { + Status.Code.NOT_FOUND -> notFound(collection, message) + Status.Code.UNAUTHENTICATED, Status.Code.PERMISSION_DENIED -> + KdrantException.Unauthorized(message.ifBlank { "Unauthorized" }) + Status.Code.INVALID_ARGUMENT -> + if (collection != null && looksLikeMissingCollection(message)) { + notFound(collection, message) + } else { + KdrantException.InvalidRequest(message.ifBlank { "Qdrant rejected the request" }) + } + Status.Code.ALREADY_EXISTS -> KdrantException.AlreadyExists(message.ifBlank { "Already exists" }) + Status.Code.DEADLINE_EXCEEDED -> KdrantException.Timeout(message.ifBlank { "Deadline exceeded" }, cause) + Status.Code.RESOURCE_EXHAUSTED -> KdrantException.RateLimited(message = message.ifBlank { RATE_LIMITED }) + Status.Code.UNAVAILABLE -> KdrantException.ServiceUnavailable(message.ifBlank { UNAVAILABLE }) + Status.Code.UNIMPLEMENTED -> KdrantException.InvalidRequest( + "Qdrant does not implement this call over gRPC${if (message.isBlank()) "" else ": $message"}", + ) + else -> KdrantException.ServerError(message.ifBlank { "Qdrant returned ${status.code}" }) + } + } + + private fun notFound(collection: String?, message: String): KdrantException = + if (collection != null) { + KdrantException.CollectionNotFound(collection, message) + } else { + KdrantException.InvalidRequest(message.ifBlank { "Not found" }) + } + + /** + * Qdrant answers some write paths on a missing collection with `INVALID_ARGUMENT` rather than + * `NOT_FOUND`, and the only thing separating that from a genuinely malformed request is the text. + * Matching on it is unpleasant, and the alternative is reporting a missing collection as a bad + * request on exactly the operations where the caller most needs to tell them apart. + */ + private fun looksLikeMissingCollection(message: String): Boolean = + message.contains("doesn't exist", ignoreCase = true) || + message.contains("does not exist", ignoreCase = true) || + message.contains("not found", ignoreCase = true) + + private const val RATE_LIMITED = "Rate limited by Qdrant" + private const val UNAVAILABLE = "Qdrant is temporarily unavailable" +} diff --git a/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransport.kt b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransport.kt new file mode 100644 index 0000000..e25b073 --- /dev/null +++ b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransport.kt @@ -0,0 +1,759 @@ +package dev.kdrant.transport.grpc + +import dev.kdrant.KdrantConfig +import dev.kdrant.model.AliasDescription +import dev.kdrant.model.AliasOperation +import dev.kdrant.model.ClusterOperation +import dev.kdrant.model.CollectionClusterInfo +import dev.kdrant.model.CollectionDescription +import dev.kdrant.model.CollectionInfo +import dev.kdrant.model.CreateCollectionRequest +import dev.kdrant.model.CreateShardKeyRequest +import dev.kdrant.model.DeleteSelector +import dev.kdrant.model.FacetHit +import dev.kdrant.model.Filter +import dev.kdrant.model.LocalShardInfo +import dev.kdrant.model.Payload +import dev.kdrant.model.PayloadSchemaType +import dev.kdrant.model.PointGroup +import dev.kdrant.model.PointId +import dev.kdrant.model.PointStruct +import dev.kdrant.model.PointVectors +import dev.kdrant.model.PointsUpdateOperation +import dev.kdrant.model.Record +import dev.kdrant.model.RemoteShardInfo +import dev.kdrant.model.ScoredPoint +import dev.kdrant.model.ScrollPage +import dev.kdrant.model.ScrollRequest +import dev.kdrant.model.SearchGroupsRequest +import dev.kdrant.model.SearchMatrixOffsets +import dev.kdrant.model.SearchMatrixPair +import dev.kdrant.model.SearchMatrixPairs +import dev.kdrant.model.SearchMatrixRequest +import dev.kdrant.model.SearchRequest +import dev.kdrant.model.ShardKey +import dev.kdrant.model.ShardTransferInfo +import dev.kdrant.model.SnapshotDescription +import dev.kdrant.model.SnapshotPriority +import dev.kdrant.model.UpdateCollectionRequest +import dev.kdrant.model.WithPayload +import dev.kdrant.transport.QdrantTransport +import grpc.health.v1.HealthCheck +import grpc.health.v1.HealthGrpcKt +import io.grpc.ManagedChannel +import io.grpc.Status +import io.grpc.StatusException +import io.grpc.StatusRuntimeException +import io.grpc.kotlin.AbstractCoroutineStub +import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.delay +import kotlinx.coroutines.flow.Flow +import kotlinx.coroutines.withContext +import kotlinx.serialization.json.JsonElement +import kotlinx.serialization.json.JsonObject +import qdrant.Collections +import qdrant.CollectionsGrpcKt +import qdrant.Points +import qdrant.PointsGrpcKt +import qdrant.SnapshotsGrpcKt +import qdrant.SnapshotsService +import java.util.concurrent.TimeUnit +import kotlin.math.min +import kotlin.random.Random + +/** + * The gRPC engine: [QdrantTransport] over Qdrant's `Collections`, `Points`, `Snapshots` and `Health` + * services. + * + * **The seam is wider than the protocol.** `QdrantTransport` was shaped by Qdrant's REST API, and + * eleven of its operations have no gRPC equivalent: telemetry, Prometheus metrics, the issues + * endpoint, snapshot recovery, snapshot download and upload, and the whole shard-scope snapshot group. + * They are not silently degraded here. Each throws an [UnsupportedOperationException] naming the + * operation and pointing at the REST engine, because a snapshot download that quietly returns nothing + * is a backup that quietly does not exist. A `KdrantException` would have been the wrong type: nothing + * failed on the wire, and no retry or fallback will make the call work. + * + * **Health probes** map to the standard gRPC health checking service rather than to `/healthz`, + * `/readyz` and `/livez`, which have no gRPC counterpart. Qdrant serves one health status, so the three + * probes answer from it: a serving node is healthy, ready and alive, and a node that answers anything + * else — or does not answer — is none of the three. That is the same false-on-failure contract the REST + * probes have. + * + * **Retries** mirror the REST engine's, because [KdrantConfig.maxRetries] is a client setting and an + * engine that ignored it would be a behaviour difference the caller did not ask for. + */ +public class GrpcQdrantTransport internal constructor( + private val config: KdrantConfig, + private val channel: ManagedChannel, + private val upsertBatchSize: Int, +) : QdrantTransport { + + private val points = PointsGrpcKt.PointsCoroutineStub(channel) + private val collections = CollectionsGrpcKt.CollectionsCoroutineStub(channel) + private val snapshots = SnapshotsGrpcKt.SnapshotsCoroutineStub(channel) + private val health = HealthGrpcKt.HealthCoroutineStub(channel) + + init { + require(upsertBatchSize > 0) { "upsertBatchSize must be positive, was $upsertBatchSize" } + } + + // --- Collections ------------------------------------------------------------------------- + + override suspend fun createCollection(name: String, request: CreateCollectionRequest) { + call(name) { collections.deadlined().create(CollectionMapping.createCollection(name, request)) } + } + + override suspend fun updateCollection(name: String, request: UpdateCollectionRequest) { + call(name) { collections.deadlined().update(CollectionMapping.updateCollection(name, request)) } + } + + override suspend fun deleteCollection(name: String) { + call(name) { + collections.deadlined().delete( + Collections.DeleteCollection.newBuilder().setCollectionName(name).build(), + ) + } + } + + override suspend fun collectionExists(name: String): Boolean = call(name) { + collections.deadlined() + .collectionExists(Collections.CollectionExistsRequest.newBuilder().setCollectionName(name).build()) + .result + .exists + } + + override suspend fun getCollection(name: String): CollectionInfo = call(name) { + CollectionMapping.collectionInfo( + collections.deadlined() + .get(Collections.GetCollectionInfoRequest.newBuilder().setCollectionName(name).build()) + .result, + ) + } + + override suspend fun listCollections(): List = call(null) { + collections.deadlined() + .list(Collections.ListCollectionsRequest.getDefaultInstance()) + .collectionsList + .map { CollectionDescription(it.name) } + } + + // --- Points ------------------------------------------------------------------------------ + + override suspend fun upsert(name: String, points: List, wait: Boolean) { + points.chunked(upsertBatchSize).forEach { batch -> upsertBatch(name, batch, wait) } + } + + override suspend fun upsert(name: String, points: Flow, wait: Boolean) { + val batch = ArrayList(upsertBatchSize) + points.collect { point -> + batch += point + if (batch.size == upsertBatchSize) { + upsertBatch(name, batch.toList(), wait) + batch.clear() + } + } + if (batch.isNotEmpty()) upsertBatch(name, batch, wait) + } + + private suspend fun upsertBatch(name: String, batch: List, wait: Boolean) { + call(name) { + points.deadlined().upsert( + Points.UpsertPoints.newBuilder() + .setCollectionName(name) + .setWait(wait) + .addAllPoints(batch.map(PointMapping::pointToProto)) + .build(), + ) + } + } + + override suspend fun query(name: String, request: SearchRequest): List = call(name) { + points.deadlined() + .query(QueryMapping.queryPoints(name, request)) + .resultList + .map(PointMapping::scoredPointToModel) + } + + override suspend fun queryBatch(name: String, requests: List): List> = + call(name) { + points.deadlined() + .queryBatch( + Points.QueryBatchPoints.newBuilder() + .setCollectionName(name) + .addAllQueryPoints(requests.map { QueryMapping.queryPoints(name, it) }) + .build(), + ) + .resultList + .map { batch -> batch.resultList.map(PointMapping::scoredPointToModel) } + } + + override suspend fun queryGroups(name: String, request: SearchGroupsRequest): List = call(name) { + points.deadlined() + .queryGroups(QueryMapping.queryGroups(name, request)) + .result + .groupsList + .map { group -> + PointGroup( + id = PointMapping.groupId(group.id), + hits = group.hitsList.map(PointMapping::scoredPointToModel), + lookup = if (group.hasLookup()) PointMapping.recordToModel(group.lookup) else null, + ) + } + } + + override suspend fun scroll(name: String, request: ScrollRequest): ScrollPage = call(name) { + val response = points.deadlined().scroll(RequestMapping.scrollPoints(name, request)) + ScrollPage( + points = response.resultList.map(PointMapping::recordToModel), + nextPageOffset = if (response.hasNextPageOffset()) { + PointMapping.idToModel(response.nextPageOffset) + } else { + null + }, + ) + } + + override suspend fun count(name: String, filter: Filter?, exact: Boolean): Long = call(name) { + points.deadlined() + .count( + Points.CountPoints.newBuilder().apply { + collectionName = name + this.exact = exact + filter?.let { this.filter = FilterMapping.toProto(it) } + }.build(), + ) + .result + .count + } + + override suspend fun retrieve( + name: String, + ids: List, + withPayload: WithPayload?, + withVector: Boolean?, + ): List = call(name) { + points.deadlined() + .get( + Points.GetPoints.newBuilder() + .setCollectionName(name) + .addAllIds(ids.map(PointMapping::idToProto)) + .setWithPayload(RequestMapping.withPayload(withPayload)) + .setWithVectors(RequestMapping.withVectors(withVector)) + .build(), + ) + .resultList + .map(PointMapping::recordToModel) + } + + override suspend fun delete(name: String, selector: DeleteSelector, wait: Boolean) { + call(name) { + points.deadlined().delete( + Points.DeletePoints.newBuilder() + .setCollectionName(name) + .setWait(wait) + .setPoints(RequestMapping.pointsSelector(selector)) + .build(), + ) + } + } + + // --- Payload and vectors ----------------------------------------------------------------- + + override suspend fun createPayloadIndex( + name: String, + field: String, + schema: PayloadSchemaType, + wait: Boolean, + ) { + call(name) { + points.deadlined().createFieldIndex( + Points.CreateFieldIndexCollection.newBuilder() + .setCollectionName(name) + .setWait(wait) + .setFieldName(field) + .setFieldType(RequestMapping.fieldType(schema)) + .build(), + ) + } + } + + override suspend fun deletePayloadIndex(name: String, field: String, wait: Boolean) { + call(name) { + points.deadlined().deleteFieldIndex( + Points.DeleteFieldIndexCollection.newBuilder() + .setCollectionName(name) + .setWait(wait) + .setFieldName(field) + .build(), + ) + } + } + + override suspend fun setPayload( + name: String, + payload: Payload, + selector: DeleteSelector, + key: String?, + wait: Boolean, + ) { + call(name) { + points.deadlined().setPayload( + Points.SetPayloadPoints.newBuilder().apply { + collectionName = name + this.wait = wait + putAllPayload(PayloadMapping.toProto(payload)) + pointsSelector = RequestMapping.pointsSelector(selector) + key?.let { this.key = it } + }.build(), + ) + } + } + + override suspend fun overwritePayload(name: String, payload: Payload, selector: DeleteSelector, wait: Boolean) { + call(name) { + points.deadlined().overwritePayload( + Points.SetPayloadPoints.newBuilder() + .setCollectionName(name) + .setWait(wait) + .putAllPayload(PayloadMapping.toProto(payload)) + .setPointsSelector(RequestMapping.pointsSelector(selector)) + .build(), + ) + } + } + + override suspend fun deletePayload(name: String, keys: List, selector: DeleteSelector, wait: Boolean) { + call(name) { + points.deadlined().deletePayload( + Points.DeletePayloadPoints.newBuilder() + .setCollectionName(name) + .setWait(wait) + .addAllKeys(keys) + .setPointsSelector(RequestMapping.pointsSelector(selector)) + .build(), + ) + } + } + + override suspend fun clearPayload(name: String, selector: DeleteSelector, wait: Boolean) { + call(name) { + points.deadlined().clearPayload( + Points.ClearPayloadPoints.newBuilder() + .setCollectionName(name) + .setWait(wait) + .setPoints(RequestMapping.pointsSelector(selector)) + .build(), + ) + } + } + + override suspend fun updateVectors(name: String, points: List, wait: Boolean) { + call(name) { + this.points.deadlined().updateVectors( + Points.UpdatePointVectors.newBuilder() + .setCollectionName(name) + .setWait(wait) + .addAllPoints(points.map(RequestMapping::pointVectors)) + .build(), + ) + } + } + + override suspend fun deleteVectors( + name: String, + vectors: List, + selector: DeleteSelector, + wait: Boolean, + ) { + call(name) { + points.deadlined().deleteVectors( + Points.DeletePointVectors.newBuilder() + .setCollectionName(name) + .setWait(wait) + .setPointsSelector(RequestMapping.pointsSelector(selector)) + .setVectors(Points.VectorsSelector.newBuilder().addAllNames(vectors)) + .build(), + ) + } + } + + override suspend fun batchUpdate(name: String, operations: List, wait: Boolean) { + call(name) { + points.deadlined().updateBatch( + Points.UpdateBatchPoints.newBuilder() + .setCollectionName(name) + .setWait(wait) + .addAllOperations(operations.map(RequestMapping::updateOperation)) + .build(), + ) + } + } + + // --- Aliases ----------------------------------------------------------------------------- + + override suspend fun updateAliases(operations: List, timeout: Int?) { + call(null) { + collections.deadlined().updateAliases( + Collections.ChangeAliases.newBuilder().apply { + addAllActions(operations.map(CollectionMapping::aliasOperation)) + timeout?.let { this.timeout = it.toLong() } + }.build(), + ) + } + } + + override suspend fun listAliases(): List = call(null) { + collections.deadlined() + .listAliases(Collections.ListAliasesRequest.getDefaultInstance()) + .aliasesList + .map { AliasDescription(aliasName = it.aliasName, collectionName = it.collectionName) } + } + + override suspend fun listCollectionAliases(name: String): List = call(name) { + collections.deadlined() + .listCollectionAliases( + Collections.ListCollectionAliasesRequest.newBuilder().setCollectionName(name).build(), + ) + .aliasesList + .map { AliasDescription(aliasName = it.aliasName, collectionName = it.collectionName) } + } + + // --- Service & health -------------------------------------------------------------------- + + override suspend fun healthz(): Boolean = serving() + + override suspend fun readyz(): Boolean = serving() + + override suspend fun livez(): Boolean = serving() + + /** + * The gRPC health service answers one status for the whole node, so the three REST probes answer + * from it. Any failure is `false`, never an exception: the REST probes report a non-2xx the same + * way, and a probe that throws is a probe every caller has to wrap. + */ + private suspend fun serving(): Boolean = + try { + withContext(config.dispatcher) { + health.deadlined().check(HealthCheck.HealthCheckRequest.getDefaultInstance()) + .status == HealthCheck.HealthCheckResponse.ServingStatus.SERVING + } + } catch (e: CancellationException) { + throw e + } catch (@Suppress("TooGenericExceptionCaught") e: Exception) { + false + } + + override suspend fun telemetry(): JsonObject = restOnly("telemetry") + + override suspend fun metrics(): String = restOnly("metrics") + + override suspend fun listIssues(): JsonElement = restOnly("listIssues") + + override suspend fun clearIssues(): Unit = restOnly("clearIssues") + + // --- Analytics --------------------------------------------------------------------------- + + override suspend fun facet( + name: String, + key: String, + filter: Filter?, + limit: Int?, + exact: Boolean, + ): List = call(name) { + points.deadlined() + .facet( + Points.FacetCounts.newBuilder().apply { + collectionName = name + this.key = key + this.exact = exact + filter?.let { this.filter = FilterMapping.toProto(it) } + limit?.let { this.limit = it.toLong() } + }.build(), + ) + .hitsList + .map { FacetHit(value = PointMapping.facetValue(it.value), count = it.count) } + } + + override suspend fun searchMatrixPairs(name: String, request: SearchMatrixRequest): SearchMatrixPairs = + call(name) { + SearchMatrixPairs( + pairs = points.deadlined() + .searchMatrixPairs(QueryMapping.searchMatrix(name, request)) + .result + .pairsList + .map { + SearchMatrixPair( + a = PointMapping.idToModel(it.a), + b = PointMapping.idToModel(it.b), + score = it.score, + ) + }, + ) + } + + override suspend fun searchMatrixOffsets(name: String, request: SearchMatrixRequest): SearchMatrixOffsets = + call(name) { + val result = points.deadlined() + .searchMatrixOffsets(QueryMapping.searchMatrix(name, request)) + .result + SearchMatrixOffsets( + offsetsRow = result.offsetsRowList, + offsetsCol = result.offsetsColList, + scores = result.scoresList, + ids = result.idsList.map(PointMapping::idToModel), + ) + } + + // --- Cluster & sharding ------------------------------------------------------------------ + + override suspend fun collectionClusterInfo(name: String): CollectionClusterInfo = call(name) { + val response = collections.deadlined().collectionClusterInfo( + Collections.CollectionClusterInfoRequest.newBuilder().setCollectionName(name).build(), + ) + CollectionClusterInfo( + peerId = response.peerId, + shardCount = response.shardCount.toInt(), + localShards = response.localShardsList.map { + LocalShardInfo( + shardId = it.shardId, + shardKey = if (it.hasShardKey()) RequestMapping.shardKeyToModel(it.shardKey) else null, + pointsCount = it.pointsCount, + state = CollectionMapping.replicaState(it.state), + ) + }, + remoteShards = response.remoteShardsList.map { + RemoteShardInfo( + shardId = it.shardId, + shardKey = if (it.hasShardKey()) RequestMapping.shardKeyToModel(it.shardKey) else null, + peerId = it.peerId, + state = CollectionMapping.replicaState(it.state), + ) + }, + shardTransfers = response.shardTransfersList.map { + ShardTransferInfo( + shardId = it.shardId, + toShardId = it.takeIf { transfer -> transfer.hasToShardId() }?.toShardId, + from = it.from, + to = it.to, + sync = it.sync, + ) + }, + ) + } + + override suspend fun updateCollectionCluster(name: String, operation: ClusterOperation, timeout: Int?) { + call(name) { + collections.deadlined() + .updateCollectionClusterSetup(CollectionMapping.clusterSetup(name, operation, timeout)) + } + } + + override suspend fun createShardKey(name: String, request: CreateShardKeyRequest, timeout: Int?) { + call(name) { + collections.deadlined().createShardKey( + Collections.CreateShardKeyRequest.newBuilder().apply { + collectionName = name + this.request = Collections.CreateShardKey.newBuilder().apply { + shardKey = RequestMapping.shardKey(request.shardKey) + request.shardsNumber?.let { shardsNumber = it } + request.replicationFactor?.let { replicationFactor = it } + request.placement?.let { addAllPlacement(it) } + }.build() + timeout?.let { this.timeout = it.toLong() } + }.build(), + ) + } + } + + override suspend fun deleteShardKey(name: String, shardKey: ShardKey, timeout: Int?) { + call(name) { + collections.deadlined().deleteShardKey( + Collections.DeleteShardKeyRequest.newBuilder().apply { + collectionName = name + request = Collections.DeleteShardKey.newBuilder() + .setShardKey(RequestMapping.shardKey(shardKey)) + .build() + timeout?.let { this.timeout = it.toLong() } + }.build(), + ) + } + } + + // --- Snapshots --------------------------------------------------------------------------- + + override suspend fun createSnapshot(name: String, wait: Boolean): SnapshotDescription = call(name) { + CollectionMapping.snapshotDescription( + snapshots.deadlined() + .create(SnapshotsService.CreateSnapshotRequest.newBuilder().setCollectionName(name).build()) + .snapshotDescription, + ) + } + + override suspend fun listSnapshots(name: String): List = call(name) { + snapshots.deadlined() + .list(SnapshotsService.ListSnapshotsRequest.newBuilder().setCollectionName(name).build()) + .snapshotDescriptionsList + .map(CollectionMapping::snapshotDescription) + } + + override suspend fun deleteSnapshot(name: String, snapshotName: String, wait: Boolean) { + call(name) { + snapshots.deadlined().delete( + SnapshotsService.DeleteSnapshotRequest.newBuilder() + .setCollectionName(name) + .setSnapshotName(snapshotName) + .build(), + ) + } + } + + override suspend fun createStorageSnapshot(wait: Boolean): SnapshotDescription = call(null) { + CollectionMapping.snapshotDescription( + snapshots.deadlined() + .createFull(SnapshotsService.CreateFullSnapshotRequest.getDefaultInstance()) + .snapshotDescription, + ) + } + + override suspend fun listStorageSnapshots(): List = call(null) { + snapshots.deadlined() + .listFull(SnapshotsService.ListFullSnapshotsRequest.getDefaultInstance()) + .snapshotDescriptionsList + .map(CollectionMapping::snapshotDescription) + } + + override suspend fun deleteStorageSnapshot(snapshotName: String, wait: Boolean) { + call(null) { + snapshots.deadlined().deleteFull( + SnapshotsService.DeleteFullSnapshotRequest.newBuilder().setSnapshotName(snapshotName).build(), + ) + } + } + + // --- What gRPC does not carry ------------------------------------------------------------ + + override suspend fun recoverSnapshot( + name: String, + location: String, + priority: SnapshotPriority?, + checksum: String?, + wait: Boolean, + ): Unit = restOnly("recoverSnapshot") + + override fun downloadSnapshot(name: String, snapshotName: String): Flow = + restOnly("downloadSnapshot") + + override suspend fun uploadSnapshot( + name: String, + data: Flow, + priority: SnapshotPriority?, + checksum: String?, + wait: Boolean, + ): Unit = restOnly("uploadSnapshot") + + override suspend fun createShardSnapshot(name: String, shardId: Int, wait: Boolean): SnapshotDescription = + restOnly("createShardSnapshot") + + override suspend fun listShardSnapshots(name: String, shardId: Int): List = + restOnly("listShardSnapshots") + + override suspend fun deleteShardSnapshot(name: String, shardId: Int, snapshotName: String, wait: Boolean): Unit = + restOnly("deleteShardSnapshot") + + override suspend fun recoverShardSnapshot( + name: String, + shardId: Int, + location: String, + priority: SnapshotPriority?, + checksum: String?, + wait: Boolean, + ): Unit = restOnly("recoverShardSnapshot") + + override fun downloadShardSnapshot(name: String, shardId: Int, snapshotName: String): Flow = + restOnly("downloadShardSnapshot") + + override suspend fun uploadShardSnapshot( + name: String, + shardId: Int, + data: Flow, + priority: SnapshotPriority?, + checksum: String?, + wait: Boolean, + ): Unit = restOnly("uploadShardSnapshot") + + override fun downloadStorageSnapshot(snapshotName: String): Flow = + restOnly("downloadStorageSnapshot") + + override fun close() { + channel.shutdown() + if (!channel.awaitTermination(SHUTDOWN_GRACE_SECONDS, TimeUnit.SECONDS)) channel.shutdownNow() + } + + // --- Plumbing ---------------------------------------------------------------------------- + + private fun restOnly(operation: String): Nothing = throw UnsupportedOperationException( + "$operation has no gRPC equivalent: Qdrant serves it over HTTP only. Use the REST engine " + + "(kdrant-transport-rest) for it, or for the whole client if you need it often.", + ) + + /** + * One attempt's deadline. It is set per call rather than on the channel because + * [KdrantConfig.requestTimeout] is documented as applying to each attempt, and a channel-level + * deadline would be shared across the retries of one logical request. + */ + private fun > S.deadlined(): S = + withDeadlineAfter(config.requestTimeout.inWholeMilliseconds, TimeUnit.MILLISECONDS) + + /** + * Runs one operation on the configured dispatcher, retrying what the REST engine retries and + * translating what is left. [collection] is the collection the call concerns, so a `NOT_FOUND` + * can name it; `null` for the cluster-wide calls. + */ + private suspend fun call(collection: String?, block: suspend () -> T): T = + withContext(config.dispatcher) { + GrpcErrors.mapping(collection) { withRetries(block) } + } + + /** + * The same policy the REST engine gets from Ktor: retry the transient statuses with exponential + * backoff and jitter, up to [KdrantConfig.maxRetries], and surface anything else at once. + */ + private suspend fun withRetries(block: suspend () -> T): T { + var attempt = 0 + while (true) { + try { + return block() + } catch (e: Exception) { + if (!retryable(e, attempt)) throw e + } + delay(backoff(attempt)) + attempt++ + } + } + + /** + * A cancellation is never retried and never swallowed: it falls through to the `throw` above, which + * is what keeps structured concurrency working. + */ + private fun retryable(failure: Exception, attempt: Int): Boolean { + if (failure is CancellationException || attempt >= config.maxRetries) return false + val code = when (failure) { + is StatusRuntimeException -> failure.status.code + is StatusException -> failure.status.code + else -> return false + } + return code == Status.Code.UNAVAILABLE || code == Status.Code.RESOURCE_EXHAUSTED + } + + private fun backoff(attempt: Int): Long { + val base = config.retryBaseDelay.inWholeMilliseconds + val exponential = base shl min(attempt, MAX_BACKOFF_SHIFT) + val capped = min(exponential, config.retryMaxDelay.inWholeMilliseconds) + return capped + Random.nextLong(base) + } +} + +private const val SHUTDOWN_GRACE_SECONDS = 5L + +/** Caps the doubling at 2^6, so a long-running retry cannot overflow before retryMaxDelay clamps it. */ +private const val MAX_BACKOFF_SHIFT = 6 diff --git a/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/KdrantGrpc.kt b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/KdrantGrpc.kt new file mode 100644 index 0000000..9eb8857 --- /dev/null +++ b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/KdrantGrpc.kt @@ -0,0 +1,36 @@ +package dev.kdrant.transport.grpc + +import dev.kdrant.KdrantConfigBuilder +import dev.kdrant.QdrantClient +import dev.kdrant.kdrantConfig + +/** + * Entry point: creates a [QdrantClient] backed by the gRPC engine. + * + * ```kotlin + * KdrantGrpc(host = "localhost") { requestTimeout = 5.seconds }.use { qdrant -> + * qdrant.upsert("docs", wait = true) { point(1) { vector(0.1f, 0.2f) } } + * } + * ``` + * + * REST stays the recommended engine. Reach for this one when throughput is the bottleneck: HTTP/2 + * multiplexing and protobuf framing win on large upserts and on many concurrent small reads. What it + * costs is the eleven operations Qdrant serves over HTTP only — telemetry, metrics, issues, snapshot + * transfer and shard-scope snapshots — which throw here rather than pretending; see + * [GrpcQdrantTransport]. + * + * @param port defaults to **6334**, Qdrant's gRPC port. It is not 6333: a config carried over from the + * REST engine will point at the wrong listener, and nothing rewrites it silently. + * @param upsertBatchSize maximum points per upsert request; larger lists and flows are split. gRPC's + * default message limit is 4 MiB, well under REST's 32 MiB, so this default is lower than the REST + * engine's for the same reason that one exists. + */ +public fun KdrantGrpc( + host: String, + port: Int = 6334, + upsertBatchSize: Int = 256, + configure: KdrantConfigBuilder.() -> Unit = {}, +): QdrantClient { + val config = kdrantConfig(host, port, configure) + return QdrantClient(GrpcQdrantTransport(config, managedChannel(config), upsertBatchSize)) +} diff --git a/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/PointMapping.kt b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/PointMapping.kt index 089af93..c89d6a4 100644 --- a/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/PointMapping.kt +++ b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/PointMapping.kt @@ -1,10 +1,12 @@ package dev.kdrant.transport.grpc +import dev.kdrant.model.FacetValue import dev.kdrant.model.PointId import dev.kdrant.model.PointStruct import dev.kdrant.model.Record import dev.kdrant.model.ScoredPoint import dev.kdrant.model.VectorData +import kotlinx.serialization.json.JsonPrimitive import qdrant.Common import qdrant.Points @@ -109,8 +111,34 @@ internal object PointMapping { id = idToModel(point.id), payload = PayloadMapping.toJson(point.payloadMap).takeIf { point.payloadMap.isNotEmpty() }, vector = if (point.hasVectors()) vectorsToModel(point.vectors) else null, + orderValue = if (point.hasOrderValue()) orderValueToModel(point.orderValue) else null, ) + /** + * The value an ordered scroll sorted this point by. It is what the client pages on — Qdrant returns + * no cursor for an ordered scroll — so dropping it here would silently turn a multi-page ordered + * scroll into its first page. + */ + private fun orderValueToModel(value: Points.OrderValue): JsonPrimitive? = when (value.variantCase) { + Points.OrderValue.VariantCase.INT -> JsonPrimitive(value.int) + Points.OrderValue.VariantCase.FLOAT -> JsonPrimitive(value.float) + Points.OrderValue.VariantCase.VARIANT_NOT_SET, null -> null + } + + fun groupId(id: Points.GroupId): JsonPrimitive = when (id.kindCase) { + Points.GroupId.KindCase.STRING_VALUE -> JsonPrimitive(id.stringValue) + Points.GroupId.KindCase.INTEGER_VALUE -> JsonPrimitive(id.integerValue) + Points.GroupId.KindCase.UNSIGNED_VALUE -> JsonPrimitive(id.unsignedValue.toULong().toString()) + Points.GroupId.KindCase.KIND_NOT_SET, null -> JsonPrimitive("") + } + + fun facetValue(value: Points.FacetValue): FacetValue = when (value.variantCase) { + Points.FacetValue.VariantCase.STRING_VALUE -> FacetValue.StringValue(value.stringValue) + Points.FacetValue.VariantCase.INTEGER_VALUE -> FacetValue.IntValue(value.integerValue) + Points.FacetValue.VariantCase.BOOL_VALUE -> FacetValue.BoolValue(value.boolValue) + Points.FacetValue.VariantCase.VARIANT_NOT_SET, null -> FacetValue.StringValue("") + } + fun scoredPointToModel(point: Points.ScoredPoint): ScoredPoint = ScoredPoint( id = idToModel(point.id), score = point.score, diff --git a/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/QueryMapping.kt b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/QueryMapping.kt new file mode 100644 index 0000000..0eb794a --- /dev/null +++ b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/QueryMapping.kt @@ -0,0 +1,278 @@ +package dev.kdrant.transport.grpc + +import dev.kdrant.model.ContextPair +import dev.kdrant.model.DecayParams +import dev.kdrant.model.Expression +import dev.kdrant.model.Filter +import dev.kdrant.model.FusionAlgorithm +import dev.kdrant.model.LookupLocation +import dev.kdrant.model.Mmr +import dev.kdrant.model.Prefetch +import dev.kdrant.model.QueryInterface +import dev.kdrant.model.RecommendStrategy +import dev.kdrant.model.SearchGroupsRequest +import dev.kdrant.model.SearchMatrixRequest +import dev.kdrant.model.SearchRequest +import dev.kdrant.model.VectorInput +import qdrant.Common +import qdrant.Points + +/** + * The query surface: `SearchRequest` and its `QueryInterface` onto `QueryPoints` and `Query`. + * + * REST puts several shapes in one `query` field and tells them apart by their JSON — a bare array is a + * dense vector, a bare number or string is a point id, an object with `indices` is sparse, an object + * with `nearest` is the long form. Protobuf has a variant per shape, so the ambiguity that the REST + * serializer resolves on the way out is resolved here on the way in, once, in [vectorInput]. + * + * `RelevanceFeedbackInput`, the eleventh `Query` variant, has no model: it is newer than the client's + * query surface and reaching it would mean adding it to the REST engine too. + */ +internal object QueryMapping { + + fun queryPoints(collection: String, request: SearchRequest): Points.QueryPoints = + Points.QueryPoints.newBuilder().apply { + collectionName = collection + request.query?.let { query = query(it) } + request.prefetch?.let { addAllPrefetch(it.map(::prefetch)) } + request.using?.let { using = it } + request.filter?.let { filter = FilterMapping.toProto(it) } + request.limit?.let { limit = it.toLong() } + request.offset?.let { offset = it.toLong() } + request.scoreThreshold?.let { scoreThreshold = it.toFloat() } + RequestMapping.searchParams(request.params)?.let { params = it } + request.lookupFrom?.let { lookupFrom = lookupLocation(it) } + RequestMapping.shardKeySelector(request.shardKey)?.let { shardKeySelector = it } + withPayload = RequestMapping.withPayload(request.withPayload) + withVectors = RequestMapping.withVectors(request.withVector) + }.build() + + fun queryGroups(collection: String, request: SearchGroupsRequest): Points.QueryPointGroups = + Points.QueryPointGroups.newBuilder().apply { + collectionName = collection + groupBy = request.groupBy + request.query?.let { query = query(it) } + request.prefetch?.let { addAllPrefetch(it.map(::prefetch)) } + request.using?.let { using = it } + request.filter?.let { filter = FilterMapping.toProto(it) } + request.limit?.let { limit = it.toLong() } + request.groupSize?.let { groupSize = it.toLong() } + request.scoreThreshold?.let { scoreThreshold = it.toFloat() } + RequestMapping.searchParams(request.params)?.let { params = it } + request.lookupFrom?.let { lookupFrom = lookupLocation(it) } + withPayload = RequestMapping.withPayload(request.withPayload) + withVectors = RequestMapping.withVectors(request.withVector) + }.build() + + fun searchMatrix(collection: String, request: SearchMatrixRequest): Points.SearchMatrixPoints = + Points.SearchMatrixPoints.newBuilder().apply { + collectionName = collection + request.filter?.let { filter = FilterMapping.toProto(it) } + request.sample?.let { sample = it.toLong() } + request.limit?.let { limit = it.toLong() } + request.using?.let { using = it } + }.build() + + private fun prefetch(prefetch: Prefetch): Points.PrefetchQuery = + Points.PrefetchQuery.newBuilder().apply { + prefetch.query?.let { query = query(it) } + prefetch.prefetch?.let { addAllPrefetch(it.map(::prefetch)) } + prefetch.using?.let { using = it } + prefetch.filter?.let { filter = FilterMapping.toProto(it) } + prefetch.limit?.let { limit = it.toLong() } + prefetch.scoreThreshold?.let { scoreThreshold = it.toFloat() } + RequestMapping.searchParams(prefetch.params)?.let { params = it } + prefetch.lookupFrom?.let { lookupFrom = lookupLocation(it) } + }.build() + + private fun lookupLocation(location: LookupLocation): Points.LookupLocation = + Points.LookupLocation.newBuilder().apply { + collectionName = location.collection + location.vector?.let { vectorName = it } + }.build() + + private fun query(query: QueryInterface): Points.Query { + val builder = Points.Query.newBuilder() + when (query) { + is VectorInput -> builder.nearest = vectorInput(query) + is QueryInterface.Nearest -> builder.setNearest(query) + is QueryInterface.Fusion -> builder.setFusion(query) + is QueryInterface.OrderBy -> builder.orderBy = orderBy(query) + QueryInterface.Sample -> builder.sample = Points.Sample.Random + is QueryInterface.Formula -> builder.formula = formula(query) + is QueryInterface.Recommend -> builder.recommend = recommend(query) + is QueryInterface.Discover -> builder.discover = Points.DiscoverInput.newBuilder() + .setTarget(vectorInput(query.target)) + .setContext(contextInput(query.context)) + .build() + is QueryInterface.Context -> builder.context = contextInput(query.pairs) + } + return builder.build() + } + + private fun orderBy(query: QueryInterface.OrderBy): Points.OrderBy = + Points.OrderBy.newBuilder().apply { + key = query.key + query.direction?.let { direction = RequestMapping.direction(it) } + }.build() + + private fun recommend(query: QueryInterface.Recommend): Points.RecommendInput = + Points.RecommendInput.newBuilder().apply { + addAllPositive(query.positive.map(::vectorInput)) + addAllNegative(query.negative.map(::vectorInput)) + query.strategy?.let { strategy = strategy(it) } + }.build() + + /** + * The long nearest form. Without MMR it is the same request as the bare vector, so it goes out as + * `nearest` rather than as a `NearestInputWithMmr` carrying no MMR — those are different messages, + * and only one of them is what a caller who wrote no reranking asked for. + */ + private fun Points.Query.Builder.setNearest(query: QueryInterface.Nearest) { + val input = vectorInput(query.input) + val reranking = query.mmr + if (reranking == null) { + nearest = input + } else { + nearestWithMmr = Points.NearestInputWithMmr.newBuilder() + .setNearest(input) + .setMmr(mmr(reranking)) + .build() + } + } + + /** + * Plain RRF and DBSF are values of the `Fusion` enum; RRF with a `k` or with per-source weights is + * a separate `Rrf` message. Same distinction the REST serializer makes between `"fusion": "rrf"` + * and an `"rrf": { ... }` object. + */ + private fun Points.Query.Builder.setFusion(query: QueryInterface.Fusion) { + val parameterized = query.algorithm == FusionAlgorithm.RRF && + (query.rrfK != null || query.rrfWeights != null) + if (parameterized) { + rrf = Points.Rrf.newBuilder().apply { + query.rrfK?.let { k = it } + query.rrfWeights?.let { addAllWeights(it) } + }.build() + } else { + fusion = when (query.algorithm) { + FusionAlgorithm.RRF -> Points.Fusion.RRF + FusionAlgorithm.DBSF -> Points.Fusion.DBSF + } + } + } + + private fun mmr(mmr: Mmr): Points.Mmr = Points.Mmr.newBuilder().apply { + mmr.diversity?.let { diversity = it } + mmr.candidatesLimit?.let { candidatesLimit = it } + }.build() + + private fun contextInput(pairs: List): Points.ContextInput = + Points.ContextInput.newBuilder() + .addAllPairs( + pairs.map { + Points.ContextInputPair.newBuilder() + .setPositive(vectorInput(it.positive)) + .setNegative(vectorInput(it.negative)) + .build() + }, + ) + .build() + + private fun strategy(strategy: RecommendStrategy): Points.RecommendStrategy = when (strategy) { + RecommendStrategy.AVERAGE_VECTOR -> Points.RecommendStrategy.AverageVector + RecommendStrategy.BEST_SCORE -> Points.RecommendStrategy.BestScore + RecommendStrategy.SUM_SCORES -> Points.RecommendStrategy.SumScores + } + + private fun vectorInput(input: VectorInput): Points.VectorInput = + Points.VectorInput.newBuilder().apply { + when (input) { + is QueryInterface.Vector -> dense = denseOf(input.values) + is QueryInterface.VectorArray -> dense = denseOf(input.values.asList()) + is QueryInterface.ById -> id = PointMapping.idToProto(input.id) + is QueryInterface.Sparse -> sparse = Points.SparseVector.newBuilder() + .addAllIndices(input.indices) + .addAllValues(input.values) + .build() + is QueryInterface.MultiVector -> multiDense = Points.MultiDenseVector.newBuilder() + .addAllVectors(input.vectors.map(::denseOf)) + .build() + } + }.build() + + private fun denseOf(values: List): Points.DenseVector = + Points.DenseVector.newBuilder().addAllData(values).build() + + private fun formula(formula: QueryInterface.Formula): Points.Formula = + Points.Formula.newBuilder() + .setExpression(expression(formula.formula)) + .putAllDefaults(formula.defaults.mapValues { (_, value) -> PayloadMapping.valueToProto(value) }) + .build() + + private fun expression(expression: Expression): Points.Expression { + val builder = Points.Expression.newBuilder() + when (expression) { + is Expression.Value -> builder.constant = expression.value.toFloat() + is Expression.Variable -> builder.variable = expression.name + is Expression.Condition -> builder.condition = conditionOf(expression) + is Expression.GeoDistance -> builder.geoDistance = geoDistance(expression) + is Expression.Datetime -> builder.datetime = expression.value + is Expression.DatetimeKey -> builder.datetimeKey = expression.key + is Expression.Mult -> builder.mult = Points.MultExpression.newBuilder() + .addAllMult(expression.operands.map(::expression)) + .build() + is Expression.Sum -> builder.sum = Points.SumExpression.newBuilder() + .addAllSum(expression.operands.map(::expression)) + .build() + is Expression.Div -> builder.div = div(expression) + is Expression.Pow -> builder.pow = Points.PowExpression.newBuilder() + .setBase(expression(expression.base)) + .setExponent(expression(expression.exponent)) + .build() + is Expression.Neg -> builder.neg = expression(expression.operand) + is Expression.Abs -> builder.abs = expression(expression.operand) + is Expression.Sqrt -> builder.sqrt = expression(expression.operand) + is Expression.Exp -> builder.exp = expression(expression.operand) + is Expression.Log10 -> builder.log10 = expression(expression.operand) + is Expression.Ln -> builder.ln = expression(expression.operand) + is Expression.ExpDecay -> builder.expDecay = decay(expression.params) + is Expression.GaussDecay -> builder.gaussDecay = decay(expression.params) + is Expression.LinDecay -> builder.linDecay = decay(expression.params) + } + return builder.build() + } + + private fun geoDistance(expression: Expression.GeoDistance): Points.GeoDistance = + Points.GeoDistance.newBuilder() + .setOrigin( + Common.GeoPoint.newBuilder() + .setLon(expression.origin.lon) + .setLat(expression.origin.lat), + ) + .setTo(expression.to) + .build() + + private fun div(expression: Expression.Div): Points.DivExpression = + Points.DivExpression.newBuilder().apply { + left = expression(expression.left) + right = expression(expression.right) + expression.byZeroDefault?.let { byZeroDefault = it.toFloat() } + }.build() + + /** + * A filter condition used as a number: true is 1.0 and false is 0.0. The wire field is a bare + * `Condition`, so the condition is wrapped in a one-clause filter and unwrapped again, which is + * the only way to reach [FilterMapping]'s per-condition mapping from outside it. + */ + private fun conditionOf(expression: Expression.Condition): Common.Condition = + FilterMapping.toProto(Filter(must = listOf(expression.condition))).getMust(0) + + private fun decay(params: DecayParams): Points.DecayParamsExpression = + Points.DecayParamsExpression.newBuilder().apply { + x = expression(params.x) + params.target?.let { target = expression(it) } + params.scale?.let { scale = it.toFloat() } + params.midpoint?.let { midpoint = it.toFloat() } + }.build() +} diff --git a/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/RequestMapping.kt b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/RequestMapping.kt new file mode 100644 index 0000000..4674869 --- /dev/null +++ b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/RequestMapping.kt @@ -0,0 +1,185 @@ +package dev.kdrant.transport.grpc + +import dev.kdrant.model.DeleteSelector +import dev.kdrant.model.Direction +import dev.kdrant.model.OrderBy +import dev.kdrant.model.PayloadSchemaType +import dev.kdrant.model.PointVectors +import dev.kdrant.model.PointsUpdateOperation +import dev.kdrant.model.ScrollRequest +import dev.kdrant.model.SearchParams +import dev.kdrant.model.ShardKey +import dev.kdrant.model.WithPayload +import kotlinx.serialization.json.JsonPrimitive +import kotlinx.serialization.json.booleanOrNull +import kotlinx.serialization.json.longOrNull +import qdrant.Collections +import qdrant.Points + +/** + * The small request shapes shared across the seam's operations: what to return, which points to act + * on, which shards to read, and how to order a scroll. + * + * These are where REST's "one field, several JSON shapes" style meets protobuf's "one message per + * shape". `with_payload` is `true` or a list of names or an exclusion object in JSON, and three + * distinct wire variants here; `points` is an id array or a filter object, and a oneof here. + */ +internal object RequestMapping { + + fun withPayload(selector: WithPayload?): Points.WithPayloadSelector = + Points.WithPayloadSelector.newBuilder().apply { + when (selector) { + null, WithPayload.None -> enable = false + WithPayload.All -> enable = true + is WithPayload.Include -> include = Points.PayloadIncludeSelector.newBuilder() + .addAllFields(selector.fields) + .build() + is WithPayload.Exclude -> exclude = Points.PayloadExcludeSelector.newBuilder() + .addAllFields(selector.fields) + .build() + } + }.build() + + fun withVectors(include: Boolean?): Points.WithVectorsSelector = + Points.WithVectorsSelector.newBuilder().setEnable(include ?: false).build() + + fun pointsSelector(selector: DeleteSelector): Points.PointsSelector = + Points.PointsSelector.newBuilder().apply { + when (selector) { + is DeleteSelector.Ids -> points = Points.PointsIdsList.newBuilder() + .addAllIds(selector.ids.map(PointMapping::idToProto)) + .build() + is DeleteSelector.ByFilter -> filter = FilterMapping.toProto(selector.filter) + } + }.build() + + fun shardKeySelector(key: ShardKey?): Points.ShardKeySelector? = key?.let { + Points.ShardKeySelector.newBuilder().addShardKeys(shardKey(it)).build() + } + + fun shardKey(key: ShardKey): Collections.ShardKey = Collections.ShardKey.newBuilder().apply { + when (key) { + is ShardKey.Name -> keyword = key.value + is ShardKey.Num -> number = key.value.toLong() + } + }.build() + + fun shardKeyToModel(key: Collections.ShardKey): ShardKey? = when (key.keyCase) { + Collections.ShardKey.KeyCase.KEYWORD -> ShardKey.Name(key.keyword) + Collections.ShardKey.KeyCase.NUMBER -> ShardKey.Num(key.number.toULong()) + Collections.ShardKey.KeyCase.KEY_NOT_SET, null -> null + } + + fun searchParams(params: SearchParams?): Points.SearchParams? = params?.let { + Points.SearchParams.newBuilder().apply { + it.hnswEf?.let { ef -> hnswEf = ef.toLong() } + it.exact?.let { value -> exact = value } + it.indexedOnly?.let { value -> indexedOnly = value } + }.build() + } + + fun direction(direction: Direction): Points.Direction = when (direction) { + Direction.ASC -> Points.Direction.Asc + Direction.DESC -> Points.Direction.Desc + } + + fun orderBy(orderBy: OrderBy): Points.OrderBy = Points.OrderBy.newBuilder().apply { + key = orderBy.key + orderBy.direction?.let { direction = direction(it) } + orderBy.startFrom?.let { startFrom = startFrom(it, orderBy.key) } + }.build() + + fun scrollPoints(name: String, request: ScrollRequest): Points.ScrollPoints = + Points.ScrollPoints.newBuilder().apply { + collectionName = name + limit = request.limit + request.filter?.let { filter = FilterMapping.toProto(it) } + request.offset?.let { offset = PointMapping.idToProto(it) } + request.orderBy?.let { orderBy = RequestMapping.orderBy(it) } + RequestMapping.shardKeySelector(request.shardKey)?.let { shardKeySelector = it } + withPayload = RequestMapping.withPayload(request.withPayload) + withVectors = RequestMapping.withVectors(request.withVector) + }.build() + + fun pointVectors(vectors: PointVectors): Points.PointVectors = Points.PointVectors.newBuilder() + .setId(PointMapping.idToProto(vectors.id)) + .setVectors(PointMapping.vectorsToProto(vectors.vector)) + .build() + + fun updateOperation(operation: PointsUpdateOperation): Points.PointsUpdateOperation { + val builder = Points.PointsUpdateOperation.newBuilder() + when (operation) { + is PointsUpdateOperation.Upsert -> + builder.upsert = Points.PointsUpdateOperation.PointStructList.newBuilder() + .addAllPoints(operation.points.map(PointMapping::pointToProto)) + .build() + is PointsUpdateOperation.Delete -> + builder.deletePoints = Points.PointsUpdateOperation.DeletePoints.newBuilder() + .setPoints(pointsSelector(operation.selector)) + .build() + is PointsUpdateOperation.SetPayload -> + builder.setPayload = Points.PointsUpdateOperation.SetPayload.newBuilder().apply { + putAllPayload(PayloadMapping.toProto(operation.payload)) + pointsSelector = pointsSelector(operation.selector) + operation.key?.let { key = it } + }.build() + is PointsUpdateOperation.OverwritePayload -> + builder.overwritePayload = Points.PointsUpdateOperation.OverwritePayload.newBuilder() + .putAllPayload(PayloadMapping.toProto(operation.payload)) + .setPointsSelector(pointsSelector(operation.selector)) + .build() + is PointsUpdateOperation.DeletePayload -> + builder.deletePayload = Points.PointsUpdateOperation.DeletePayload.newBuilder() + .addAllKeys(operation.keys) + .setPointsSelector(pointsSelector(operation.selector)) + .build() + is PointsUpdateOperation.ClearPayload -> + builder.clearPayload = Points.PointsUpdateOperation.ClearPayload.newBuilder() + .setPoints(pointsSelector(operation.selector)) + .build() + is PointsUpdateOperation.UpdateVectors -> + builder.updateVectors = Points.PointsUpdateOperation.UpdateVectors.newBuilder() + .addAllPoints(operation.points.map(::pointVectors)) + .build() + is PointsUpdateOperation.DeleteVectors -> + builder.deleteVectors = Points.PointsUpdateOperation.DeleteVectors.newBuilder() + .setPointsSelector(pointsSelector(operation.selector)) + .setVectors(Points.VectorsSelector.newBuilder().addAllNames(operation.vectors)) + .build() + } + return builder.build() + } + + fun fieldType(schema: PayloadSchemaType): Points.FieldType = when (schema) { + PayloadSchemaType.KEYWORD -> Points.FieldType.FieldTypeKeyword + PayloadSchemaType.INTEGER -> Points.FieldType.FieldTypeInteger + PayloadSchemaType.FLOAT -> Points.FieldType.FieldTypeFloat + PayloadSchemaType.GEO -> Points.FieldType.FieldTypeGeo + PayloadSchemaType.TEXT -> Points.FieldType.FieldTypeText + PayloadSchemaType.BOOL -> Points.FieldType.FieldTypeBool + PayloadSchemaType.DATETIME -> Points.FieldType.FieldTypeDatetime + PayloadSchemaType.UUID -> Points.FieldType.FieldTypeUuid + } + + /** + * `start_from` is one field of an untyped JSON value in REST and a four-way oneof here, so the + * type has to be decided from the value rather than passed along. + * + * Integral before floating point, and by the text: a value too large for `Long` still parses as a + * `Double`, so testing that first would round the boundary the caller is resuming from and skip + * whatever fell in the gap. The same reasoning as the payload mapping, for the same reason. + */ + private fun startFrom(value: JsonPrimitive, key: String): Points.StartFrom = + Points.StartFrom.newBuilder().apply { + when { + value.isString -> datetime = value.content + value.booleanOrNull != null -> throw IllegalArgumentException( + "startFrom on '$key' is a boolean, and a scroll orders by a number or a datetime", + ) + value.longOrNull != null -> integer = value.longOrNull!! + else -> float = value.content.toDoubleOrNull() ?: throw IllegalArgumentException( + "startFrom on '$key' is ${value.content}, which is neither a number nor a datetime", + ) + } + }.build() +} diff --git a/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcClientContractIntegrationTest.kt b/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcClientContractIntegrationTest.kt new file mode 100644 index 0000000..4df214f --- /dev/null +++ b/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcClientContractIntegrationTest.kt @@ -0,0 +1,18 @@ +package dev.kdrant.transport.grpc + +import dev.kdrant.QdrantClient +import dev.kdrant.testkit.QdrantClientContract +import org.testcontainers.qdrant.QdrantContainer + +/** + * The shared client contract, run over the gRPC engine. Same file, same server, different protocol — + * which is the whole point of extracting it. + * + * The port is the one difference the subclass has to know about: Qdrant serves gRPC on 6334 and the + * container maps both. + */ +class GrpcClientContractIntegrationTest : QdrantClientContract() { + + override fun connect(container: QdrantContainer): QdrantClient = + KdrantGrpc(host = container.host, port = container.grpcPort) +} diff --git a/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcErrorsTest.kt b/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcErrorsTest.kt new file mode 100644 index 0000000..bc06267 --- /dev/null +++ b/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcErrorsTest.kt @@ -0,0 +1,112 @@ +package dev.kdrant.transport.grpc + +import dev.kdrant.KdrantException +import io.grpc.Status +import io.grpc.StatusException +import io.grpc.StatusRuntimeException +import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.test.runTest +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.assertThrows + +/** + * gRPC statuses onto the same exception hierarchy the REST engine raises. A caller that switches + * engines keeps its `when`, which only holds if the two agree on which failure is which. + */ +class GrpcErrorsTest { + + @Test + fun `NOT_FOUND on a collection call names the collection`() = runTest { + val error = assertThrows { + GrpcErrors.mapping("docs") { throw Status.NOT_FOUND.withDescription("nope").asRuntimeException() } + } + + assertEquals("docs", error.collection) + assertTrue(error.message!!.contains("nope"), error.message) + } + + @Test + fun `NOT_FOUND with no collection in scope is a bad request, not a missing collection`() = runTest { + // listAliases and the storage snapshots are not about a collection, so there is no name to + // report and inventing one would be worse than saying the request was rejected. + assertThrows { + GrpcErrors.mapping(null) { throw Status.NOT_FOUND.asRuntimeException() } + } + } + + @Test + fun `INVALID_ARGUMENT saying the collection does not exist is read as a missing collection`() = runTest { + // Qdrant answers some write paths this way rather than with NOT_FOUND, and the only thing + // separating it from a genuinely malformed request is the text. + val error = assertThrows { + GrpcErrors.mapping("docs") { + throw Status.INVALID_ARGUMENT.withDescription("Collection `docs` doesn't exist!").asRuntimeException() + } + } + + assertEquals("docs", error.collection) + } + + @Test + fun `INVALID_ARGUMENT about anything else stays a bad request`() = runTest { + assertThrows { + GrpcErrors.mapping("docs") { + throw Status.INVALID_ARGUMENT.withDescription("Wrong input: vector dim mismatch").asRuntimeException() + } + } + } + + @Test + fun `a cancellation propagates untouched, so structured concurrency still works`() = runTest { + // gRPC reports a cancelled call as Status.CANCELLED; turning that into a KdrantException would + // hand the coroutine a failure instead of its own cancellation. + assertThrows { + GrpcErrors.mapping("docs") { throw CancellationException("cancelled") } + } + } + + @Test + fun `the remaining statuses land on the exception the REST engine raises for the same case`() = runTest { + val cases = listOf>>( + Status.UNAUTHENTICATED to KdrantException.Unauthorized::class.java, + Status.PERMISSION_DENIED to KdrantException.Unauthorized::class.java, + Status.ALREADY_EXISTS to KdrantException.AlreadyExists::class.java, + Status.DEADLINE_EXCEEDED to KdrantException.Timeout::class.java, + Status.RESOURCE_EXHAUSTED to KdrantException.RateLimited::class.java, + Status.UNAVAILABLE to KdrantException.ServiceUnavailable::class.java, + Status.INTERNAL to KdrantException.ServerError::class.java, + Status.DATA_LOSS to KdrantException.ServerError::class.java, + ) + + cases.forEach { (status, expected) -> + val error = runCatching { + GrpcErrors.mapping("docs") { throw status.asRuntimeException() } + }.exceptionOrNull() + + assertTrue(expected.isInstance(error), "${status.code} mapped to ${error?.let { it::class.simpleName }}") + } + } + + @Test + fun `both of gRPC's status exception types are translated, not only the runtime one`() = runTest { + // The stubs throw StatusException; the older blocking API throws StatusRuntimeException. Both + // reach this code, and catching only one would leak a raw gRPC type out of the seam. + assertThrows { + GrpcErrors.mapping("docs") { throw StatusException(Status.UNAVAILABLE) } + } + assertThrows { + GrpcErrors.mapping("docs") { throw StatusRuntimeException(Status.UNAVAILABLE) } + } + } + + @Test + fun `UNIMPLEMENTED says the server does not have the call, rather than blaming the request`() = runTest { + val error = assertThrows { + GrpcErrors.mapping("docs") { throw Status.UNIMPLEMENTED.asRuntimeException() } + } + + assertTrue(error.message!!.contains("does not implement"), error.message) + } +} diff --git a/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransportTest.kt b/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransportTest.kt new file mode 100644 index 0000000..bd6a239 --- /dev/null +++ b/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransportTest.kt @@ -0,0 +1,615 @@ +package dev.kdrant.transport.grpc + +import dev.kdrant.dsl.filter +import dev.kdrant.dsl.payloadOf +import dev.kdrant.kdrantConfig +import dev.kdrant.model.CollectionStatus +import dev.kdrant.model.CreateCollectionRequest +import dev.kdrant.model.DeleteSelector +import dev.kdrant.model.Direction +import dev.kdrant.model.Distance +import dev.kdrant.model.FacetValue +import dev.kdrant.model.OrderBy +import dev.kdrant.model.PayloadSchemaType +import dev.kdrant.model.PointId +import dev.kdrant.model.PointStruct +import dev.kdrant.model.PointVectors +import dev.kdrant.model.PointsUpdateOperation +import dev.kdrant.model.QueryInterface +import dev.kdrant.model.ScrollRequest +import dev.kdrant.model.SearchGroupsRequest +import dev.kdrant.model.SearchMatrixRequest +import dev.kdrant.model.SearchRequest +import dev.kdrant.model.VectorData +import dev.kdrant.model.VectorParams +import dev.kdrant.model.VectorsConfig +import dev.kdrant.model.WithPayload +import io.grpc.Server +import io.grpc.inprocess.InProcessChannelBuilder +import io.grpc.inprocess.InProcessServerBuilder +import kotlinx.coroutines.flow.flowOf +import kotlinx.coroutines.test.runTest +import kotlinx.serialization.json.JsonPrimitive +import org.junit.jupiter.api.AfterEach +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertFalse +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.BeforeEach +import org.junit.jupiter.api.Test +import qdrant.Collections +import qdrant.CollectionsGrpcKt +import qdrant.Points +import qdrant.PointsGrpcKt +import qdrant.SnapshotsGrpcKt +import qdrant.SnapshotsService + +/** + * What the transport puts on the wire, asserted against a real gRPC server over the in-process + * transport. The fake services record every request and answer with canned responses, so both halves + * of each operation are covered: the message built going out, and the model read coming back. + * + * This is the gRPC counterpart of the REST engine's `MockEngine` tests. It is not the client contract — + * that runs against a real Qdrant and lives in `kdrant-testkit` — it is the layer below, where a field + * set on the wrong message is visible without a server that has to agree. + */ +class GrpcQdrantTransportTest { + + private lateinit var server: Server + private lateinit var transport: GrpcQdrantTransport + private val points = RecordingPoints() + private val collections = RecordingCollections() + private val snapshots = RecordingSnapshots() + + @BeforeEach + fun start() { + val name = InProcessServerBuilder.generateName() + server = InProcessServerBuilder.forName(name) + .directExecutor() + .addService(points) + .addService(collections) + .addService(snapshots) + .build() + .start() + transport = GrpcQdrantTransport( + config = kdrantConfig("localhost", 6334), + channel = InProcessChannelBuilder.forName(name).directExecutor().build(), + upsertBatchSize = 2, + ) + } + + @AfterEach + fun stop() { + transport.close() + server.shutdownNow() + } + + // --- Collections ------------------------------------------------------------------------- + + @Test + fun `createCollection carries the vector config it was given`() = runTest { + transport.createCollection( + "docs", + CreateCollectionRequest( + vectors = VectorsConfig.Named(mapOf("text" to VectorParams(size = 8, distance = Distance.DOT))), + onDiskPayload = true, + shardNumber = 3, + ), + ) + + val request = collections.created.single() + assertEquals("docs", request.collectionName) + assertEquals(true, request.onDiskPayload) + assertEquals(3, request.shardNumber) + val params = request.vectorsConfig.paramsMap.mapMap.getValue("text") + assertEquals(8L, params.size) + assertEquals(Collections.Distance.Dot, params.distance) + } + + @Test + fun `getCollection reads the status, counts and config back`() = runTest { + val info = transport.getCollection("docs") + + assertEquals(CollectionStatus.GREEN, info.status) + assertEquals(7L, info.pointsCount) + assertEquals( + VectorsConfig.Single(VectorParams(size = 4, distance = Distance.COSINE)), + info.config?.params?.vectors, + ) + assertEquals("keyword", info.payloadSchema.getValue("lang").dataType) + assertEquals(PayloadSchemaType.KEYWORD, info.payloadSchema.getValue("lang").schemaType) + } + + @Test + fun `collectionExists reads the nested boolean rather than the envelope`() = runTest { + assertTrue(transport.collectionExists("docs")) + } + + // --- Points ------------------------------------------------------------------------------ + + @Test + fun `a list longer than the batch size is split, and every point is sent once`() = runTest { + val batch = (1L..5L).map { PointStruct(PointId.num(it), VectorData.Dense(listOf(0.1f))) } + + transport.upsert("docs", batch, wait = true) + + // upsertBatchSize is 2, so five points are three requests: 2, 2, 1. + assertEquals(listOf(2, 2, 1), points.upserts.map { it.pointsCount }) + assertEquals((1L..5L).toList(), points.upserts.flatMap { it.pointsList.map { p -> p.id.num } }) + assertTrue(points.upserts.all { it.wait }) + } + + @Test + fun `a flow is chunked the same way a list is`() = runTest { + val batch = (1L..3L).map { PointStruct(PointId.num(it), VectorData.Dense(listOf(0.1f))) } + + transport.upsert("docs", flowOf(*batch.toTypedArray()), wait = false) + + assertEquals(listOf(2, 1), points.upserts.map { it.pointsCount }) + assertFalse(points.upserts.first().wait) + } + + @Test + fun `retrieve asks for the payload and vectors it was told to, and reads the points back`() = runTest { + val records = transport.retrieve( + "docs", + ids = listOf(PointId.num(1), PointId.uuid(UUID)), + withPayload = WithPayload.include("lang"), + withVector = true, + ) + + val request = points.gets.single() + assertEquals(listOf(1L), request.idsList.filter { it.hasNum() }.map { it.num }) + assertEquals(listOf(UUID), request.idsList.filter { !it.hasNum() }.map { it.uuid }) + assertEquals(listOf("lang"), request.withPayload.include.fieldsList) + assertTrue(request.withVectors.enable) + + assertEquals(PointId.num(1), records.single().id) + assertEquals(JsonPrimitive("it"), records.single().payload?.get("lang")) + } + + @Test + fun `count sends the filter and reads the nested result`() = runTest { + val total = transport.count("docs", filter { must { "lang" eq "it" } }, exact = true) + + assertEquals(42L, total) + val request = points.counts.single() + assertTrue(request.exact) + assertEquals("lang", request.filter.mustList.single().field.key) + } + + @Test + fun `delete by id and delete by filter reach different halves of the selector`() = runTest { + transport.delete("docs", DeleteSelector.Ids(listOf(PointId.num(1))), wait = true) + transport.delete("docs", DeleteSelector.ByFilter(filter { must { "stale" eq true } }), wait = true) + + assertEquals( + listOf(Points.PointsSelector.PointsSelectorOneOfCase.POINTS, Points.PointsSelector.PointsSelectorOneOfCase.FILTER), + points.deletes.map { it.points.pointsSelectorOneOfCase }, + ) + } + + // --- Query ------------------------------------------------------------------------------- + + @Test + fun `a query carries its vector, limit, filter and selectors`() = runTest { + val hits = transport.query( + "docs", + SearchRequest( + query = QueryInterface.Vector(listOf(0.1f, 0.2f)), + limit = 5, + offset = 2, + filter = filter { must { "lang" eq "it" } }, + withPayload = WithPayload.All, + withVector = true, + scoreThreshold = 0.5, + ), + ) + + val request = points.queries.single() + assertEquals(listOf(0.1f, 0.2f), request.query.nearest.dense.dataList) + assertEquals(5L, request.limit) + assertEquals(2L, request.offset) + assertEquals(0.5f, request.scoreThreshold) + assertTrue(request.withPayload.enable) + assertTrue(request.withVectors.enable) + assertEquals(PointId.num(9), hits.single().id) + assertEquals(0.75f, hits.single().score) + } + + @Test + fun `a prefetch with fusion becomes a prefetch list and a fusion query`() = runTest { + transport.query( + "docs", + SearchRequest( + prefetch = listOf( + dev.kdrant.model.Prefetch(query = QueryInterface.Vector(listOf(0.1f)), limit = 20), + dev.kdrant.model.Prefetch(query = QueryInterface.Sparse(listOf(1), listOf(0.5f)), limit = 20), + ), + query = QueryInterface.Fusion.rrf(k = 60), + limit = 5, + ), + ) + + val request = points.queries.single() + assertEquals(2, request.prefetchCount) + assertEquals(listOf(1), request.getPrefetch(1).query.nearest.sparse.indicesList) + // A parameterized RRF is the Rrf message, not the Fusion enum. + assertEquals(Points.Query.VariantCase.RRF, request.query.variantCase) + assertEquals(60, request.query.rrf.k) + } + + @Test + fun `a plain fusion stays the enum, so it is not confused with a parameterized one`() = runTest { + transport.query("docs", SearchRequest(query = QueryInterface.Fusion.dbsf, limit = 5)) + + assertEquals(Points.Query.VariantCase.FUSION, points.queries.single().query.variantCase) + assertEquals(Points.Fusion.DBSF, points.queries.single().query.fusion) + } + + @Test + fun `queryBatch sends one QueryPoints per request and reads one result list per request`() = runTest { + val results = transport.queryBatch( + "docs", + listOf( + SearchRequest(query = QueryInterface.Vector(listOf(0.1f)), limit = 1), + SearchRequest(query = QueryInterface.Vector(listOf(0.2f)), limit = 1), + ), + ) + + assertEquals(2, points.queryBatches.single().queryPointsCount) + assertEquals(2, results.size) + } + + @Test + fun `queryGroups carries the group field and reads the groups back`() = runTest { + val groups = transport.queryGroups( + "docs", + SearchGroupsRequest(groupBy = "lang", groupSize = 2, limit = 3, query = QueryInterface.Vector(listOf(0.1f))), + ) + + assertEquals("lang", points.groupQueries.single().groupBy) + assertEquals(JsonPrimitive("it"), groups.single().id) + assertEquals(1, groups.single().hits.size) + } + + // --- Scroll ------------------------------------------------------------------------------ + + @Test + fun `an ordered scroll sends the order and reads back the value it sorted on`() = runTest { + val page = transport.scroll( + "docs", + ScrollRequest(limit = 10, orderBy = OrderBy(key = "n", direction = Direction.DESC)), + ) + + val request = points.scrolls.single() + assertEquals("n", request.orderBy.key) + assertEquals(Points.Direction.Desc, request.orderBy.direction) + // Qdrant returns no cursor for an ordered scroll; the client pages on this value instead. + assertEquals(JsonPrimitive(3L), page.points.single().orderValue) + } + + @Test + fun `a scroll resuming from an id sends it as the offset`() = runTest { + transport.scroll("docs", ScrollRequest(limit = 10, offset = PointId.num(7))) + + assertEquals(7L, points.scrolls.single().offset.num) + } + + // --- Payload, vectors and batches -------------------------------------------------------- + + @Test + fun `setPayload sends the payload, the selector and the nested key`() = runTest { + transport.setPayload( + "docs", + payloadOf("reviewed" to true), + DeleteSelector.Ids(listOf(PointId.num(1))), + key = "meta", + wait = true, + ) + + val request = points.setPayloads.single() + assertEquals(true, request.payloadMap.getValue("reviewed").boolValue) + assertEquals("meta", request.key) + assertEquals(1L, request.pointsSelector.points.getIds(0).num) + } + + @Test + fun `updateVectors sends one message per point, with its named vectors`() = runTest { + transport.updateVectors( + "docs", + listOf( + PointVectors( + PointId.num(1), + VectorData.Named(mapOf("text" to VectorData.Dense(listOf(0.9f)))), + ), + ), + wait = true, + ) + + val request = points.vectorUpdates.single() + assertEquals(listOf(0.9f), request.getPoints(0).vectors.vectors.vectorsMap.getValue("text").dense.dataList) + } + + @Test + fun `batchUpdate keeps the order it was given, because that is the only guarantee it has`() = runTest { + transport.batchUpdate( + "docs", + listOf( + PointsUpdateOperation.Upsert(listOf(PointStruct(PointId.num(1), VectorData.Dense(listOf(0.1f))))), + PointsUpdateOperation.SetPayload(payloadOf("a" to 1), DeleteSelector.Ids(listOf(PointId.num(1)))), + PointsUpdateOperation.ClearPayload(DeleteSelector.Ids(listOf(PointId.num(1)))), + ), + wait = true, + ) + + assertEquals( + listOf( + Points.PointsUpdateOperation.OperationCase.UPSERT, + Points.PointsUpdateOperation.OperationCase.SET_PAYLOAD, + Points.PointsUpdateOperation.OperationCase.CLEAR_PAYLOAD, + ), + points.batches.single().operationsList.map { it.operationCase }, + ) + } + + @Test + fun `a payload index carries the field type Qdrant needs to build it`() = runTest { + transport.createPayloadIndex("docs", "lang", PayloadSchemaType.KEYWORD, wait = true) + + assertEquals(Points.FieldType.FieldTypeKeyword, points.indexes.single().fieldType) + assertEquals("lang", points.indexes.single().fieldName) + } + + // --- Analytics and service --------------------------------------------------------------- + + @Test + fun `facet reads each hit's value out of the variant it arrived in`() = runTest { + val hits = transport.facet("docs", key = "lang", filter = null, limit = 5, exact = true) + + assertEquals(FacetValue.StringValue("it"), hits[0].value) + assertEquals(2L, hits[0].count) + assertEquals(FacetValue.IntValue(2024), hits[1].value) + assertEquals(5L, points.facets.single().limit) + } + + @Test + fun `both matrix forms read their own result shape`() = runTest { + val pairs = transport.searchMatrixPairs("docs", SearchMatrixRequest(sample = 10, limit = 2)) + val offsets = transport.searchMatrixOffsets("docs", SearchMatrixRequest(sample = 10, limit = 2)) + + assertEquals(PointId.num(1), pairs.pairs.single().a) + assertEquals(listOf(0L), offsets.offsetsRow) + assertEquals(listOf(PointId.num(1)), offsets.ids) + } + + @Test + fun `aliases go out as one atomic batch of actions`() = runTest { + transport.updateAliases( + listOf( + dev.kdrant.model.AliasOperation.Delete("docs"), + dev.kdrant.model.AliasOperation.Create(collectionName = "docs-v2", aliasName = "docs"), + ), + timeout = 5, + ) + + val request = collections.aliasChanges.single() + assertEquals(2, request.actionsCount) + assertEquals("docs", request.getActions(0).deleteAlias.aliasName) + assertEquals("docs-v2", request.getActions(1).createAlias.collectionName) + assertEquals(5L, request.timeout) + } + + @Test + fun `a snapshot description keeps its name, size and creation instant`() = runTest { + val snapshot = transport.createSnapshot("docs", wait = true) + + assertEquals("docs-2024.snapshot", snapshot.name) + assertEquals(1024L, snapshot.size) + assertEquals("2024-06-01T00:00:00Z", snapshot.creationTime) + } + + // --- Fixtures ---------------------------------------------------------------------------- + + private class RecordingPoints : PointsGrpcKt.PointsCoroutineImplBase() { + val upserts = mutableListOf() + val gets = mutableListOf() + val counts = mutableListOf() + val deletes = mutableListOf() + val queries = mutableListOf() + val queryBatches = mutableListOf() + val groupQueries = mutableListOf() + val scrolls = mutableListOf() + val setPayloads = mutableListOf() + val vectorUpdates = mutableListOf() + val batches = mutableListOf() + val indexes = mutableListOf() + val facets = mutableListOf() + + override suspend fun upsert(request: Points.UpsertPoints) = accepted { upserts += request } + + override suspend fun delete(request: Points.DeletePoints) = accepted { deletes += request } + + override suspend fun setPayload(request: Points.SetPayloadPoints) = accepted { setPayloads += request } + + override suspend fun updateVectors(request: Points.UpdatePointVectors) = accepted { vectorUpdates += request } + + override suspend fun createFieldIndex(request: Points.CreateFieldIndexCollection) = + accepted { indexes += request } + + override suspend fun get(request: Points.GetPoints): Points.GetResponse { + gets += request + return Points.GetResponse.newBuilder().addResult(retrievedPoint()).build() + } + + override suspend fun count(request: Points.CountPoints): Points.CountResponse { + counts += request + return Points.CountResponse.newBuilder() + .setResult(Points.CountResult.newBuilder().setCount(42)) + .build() + } + + override suspend fun query(request: Points.QueryPoints): Points.QueryResponse { + queries += request + return Points.QueryResponse.newBuilder().addResult(scoredPoint()).build() + } + + override suspend fun queryBatch(request: Points.QueryBatchPoints): Points.QueryBatchResponse { + queryBatches += request + val batch = Points.BatchResult.newBuilder().addResult(scoredPoint()).build() + return Points.QueryBatchResponse.newBuilder().addResult(batch).addResult(batch).build() + } + + override suspend fun queryGroups(request: Points.QueryPointGroups): Points.QueryGroupsResponse { + groupQueries += request + val group = Points.PointGroup.newBuilder() + .setId(Points.GroupId.newBuilder().setStringValue("it")) + .addHits(scoredPoint()) + .build() + return Points.QueryGroupsResponse.newBuilder() + .setResult(Points.GroupsResult.newBuilder().addGroups(group)) + .build() + } + + override suspend fun scroll(request: Points.ScrollPoints): Points.ScrollResponse { + scrolls += request + return Points.ScrollResponse.newBuilder() + .addResult( + retrievedPoint().toBuilder() + .setOrderValue(Points.OrderValue.newBuilder().setInt(3)) + .build(), + ) + .build() + } + + override suspend fun updateBatch(request: Points.UpdateBatchPoints): Points.UpdateBatchResponse { + batches += request + return Points.UpdateBatchResponse.getDefaultInstance() + } + + override suspend fun facet(request: Points.FacetCounts): Points.FacetResponse { + facets += request + return Points.FacetResponse.newBuilder() + .addHits(facetHit(Points.FacetValue.newBuilder().setStringValue("it").build(), 2)) + .addHits(facetHit(Points.FacetValue.newBuilder().setIntegerValue(2024).build(), 1)) + .build() + } + + override suspend fun searchMatrixPairs(request: Points.SearchMatrixPoints): Points.SearchMatrixPairsResponse = + Points.SearchMatrixPairsResponse.newBuilder() + .setResult( + Points.SearchMatrixPairs.newBuilder().addPairs( + Points.SearchMatrixPair.newBuilder() + .setA(id(1)) + .setB(id(2)) + .setScore(0.5f), + ), + ) + .build() + + override suspend fun searchMatrixOffsets( + request: Points.SearchMatrixPoints, + ): Points.SearchMatrixOffsetsResponse = + Points.SearchMatrixOffsetsResponse.newBuilder() + .setResult( + Points.SearchMatrixOffsets.newBuilder() + .addOffsetsRow(0) + .addOffsetsCol(1) + .addScores(0.5f) + .addIds(id(1)), + ) + .build() + + private inline fun accepted(record: () -> Unit): Points.PointsOperationResponse { + record() + return Points.PointsOperationResponse.newBuilder() + .setResult(Points.UpdateResult.newBuilder().setStatus(Points.UpdateStatus.Completed)) + .build() + } + + private fun facetHit(value: Points.FacetValue, count: Long): Points.FacetHit = + Points.FacetHit.newBuilder().setValue(value).setCount(count).build() + + private fun retrievedPoint(): Points.RetrievedPoint = Points.RetrievedPoint.newBuilder() + .setId(id(1)) + .putPayload("lang", qdrant.JsonWithInt.Value.newBuilder().setStringValue("it").build()) + .build() + + private fun scoredPoint(): Points.ScoredPoint = Points.ScoredPoint.newBuilder() + .setId(id(9)) + .setScore(0.75f) + .build() + + private fun id(value: Long) = qdrant.Common.PointId.newBuilder().setNum(value).build() + } + + private class RecordingCollections : CollectionsGrpcKt.CollectionsCoroutineImplBase() { + val created = mutableListOf() + val aliasChanges = mutableListOf() + + override suspend fun create(request: Collections.CreateCollection): Collections.CollectionOperationResponse { + created += request + return Collections.CollectionOperationResponse.newBuilder().setResult(true).build() + } + + override suspend fun updateAliases( + request: Collections.ChangeAliases, + ): Collections.CollectionOperationResponse { + aliasChanges += request + return Collections.CollectionOperationResponse.newBuilder().setResult(true).build() + } + + override suspend fun collectionExists( + request: Collections.CollectionExistsRequest, + ): Collections.CollectionExistsResponse = + Collections.CollectionExistsResponse.newBuilder() + .setResult(Collections.CollectionExists.newBuilder().setExists(true)) + .build() + + override suspend fun get( + request: Collections.GetCollectionInfoRequest, + ): Collections.GetCollectionInfoResponse { + val params = Collections.CollectionParams.newBuilder() + .setVectorsConfig( + Collections.VectorsConfig.newBuilder().setParams( + Collections.VectorParams.newBuilder() + .setSize(4) + .setDistance(Collections.Distance.Cosine), + ), + ) + .setShardNumber(1) + .build() + val info = Collections.CollectionInfo.newBuilder() + .setStatus(Collections.CollectionStatus.Green) + .setPointsCount(7) + .setConfig(Collections.CollectionConfig.newBuilder().setParams(params)) + .putPayloadSchema( + "lang", + Collections.PayloadSchemaInfo.newBuilder() + .setDataType(Collections.PayloadSchemaType.Keyword) + .build(), + ) + .build() + return Collections.GetCollectionInfoResponse.newBuilder().setResult(info).build() + } + } + + private class RecordingSnapshots : SnapshotsGrpcKt.SnapshotsCoroutineImplBase() { + override suspend fun create( + request: SnapshotsService.CreateSnapshotRequest, + ): SnapshotsService.CreateSnapshotResponse = + SnapshotsService.CreateSnapshotResponse.newBuilder() + .setSnapshotDescription( + SnapshotsService.SnapshotDescription.newBuilder() + .setName("docs-2024.snapshot") + .setSize(1024) + .setCreationTime( + com.google.protobuf.Timestamp.newBuilder().setSeconds(1_717_200_000).build(), + ), + ) + .build() + } + + private companion object { + const val UUID = "550e8400-e29b-41d4-a716-446655440000" + } +} diff --git a/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/RestOnlyOperationsTest.kt b/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/RestOnlyOperationsTest.kt new file mode 100644 index 0000000..c2c6613 --- /dev/null +++ b/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/RestOnlyOperationsTest.kt @@ -0,0 +1,86 @@ +package dev.kdrant.transport.grpc + +import dev.kdrant.kdrantConfig +import io.grpc.inprocess.InProcessChannelBuilder +import kotlinx.coroutines.flow.Flow +import kotlinx.coroutines.flow.emptyFlow +import kotlinx.coroutines.test.runTest +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.DynamicTest +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.TestFactory +import org.junit.jupiter.api.assertThrows + +/** + * The eleven operations `QdrantTransport` carries that Qdrant serves over HTTP only. + * + * The seam was shaped by the REST API, so it is wider than the gRPC protocol, and the interesting + * question is not whether these work — they cannot — but what happens when one is called. Each fails + * loudly, naming itself and naming REST. A snapshot download that returned an empty flow, or a + * telemetry call that returned an empty object, would be a client agreeing to something it did not do. + * + * No server is needed: the failure is decided before anything reaches the wire, which is also what + * these assertions prove. + */ +class RestOnlyOperationsTest { + + private val transport = GrpcQdrantTransport( + config = kdrantConfig("localhost", 6334), + channel = InProcessChannelBuilder.forName("no-server").build(), + upsertBatchSize = 64, + ) + + @TestFactory + fun `every operation gRPC does not carry fails by naming itself and REST`(): List { + val operations: Map Unit> = mapOf( + "telemetry" to { transport.telemetry() }, + "metrics" to { transport.metrics() }, + "listIssues" to { transport.listIssues() }, + "clearIssues" to { transport.clearIssues() }, + "recoverSnapshot" to { transport.recoverSnapshot("c", "file://x", null, null, true) }, + "uploadSnapshot" to { transport.uploadSnapshot("c", emptyFlow(), null, null, true) }, + "createShardSnapshot" to { transport.createShardSnapshot("c", 0, true) }, + "listShardSnapshots" to { transport.listShardSnapshots("c", 0) }, + "deleteShardSnapshot" to { transport.deleteShardSnapshot("c", 0, "s", true) }, + "recoverShardSnapshot" to { transport.recoverShardSnapshot("c", 0, "file://x", null, null, true) }, + "uploadShardSnapshot" to { transport.uploadShardSnapshot("c", 0, emptyFlow(), null, null, true) }, + ) + return operations.map { (name, call) -> + DynamicTest.dynamicTest(name) { + runTest { + val error = assertThrows { call() } + assertTrue(error.message!!.contains(name), error.message) + assertTrue(error.message!!.contains("REST"), error.message) + } + } + } + } + + @TestFactory + fun `the streaming downloads fail at the call, not on the first collect`(): List { + // These return a Flow rather than suspending, so a lazy failure would surface far from the + // mistake — at the collector, in whatever coroutine happened to consume it. + val downloads: Map Flow> = mapOf( + "downloadSnapshot" to { transport.downloadSnapshot("c", "s") }, + "downloadShardSnapshot" to { transport.downloadShardSnapshot("c", 0, "s") }, + "downloadStorageSnapshot" to { transport.downloadStorageSnapshot("s") }, + ) + return downloads.map { (name, call) -> + DynamicTest.dynamicTest(name) { + val error = assertThrows { call() } + assertTrue(error.message!!.contains(name), error.message) + } + } + } + + @Test + fun `an upsert batch size that cannot hold a point is refused at construction`() { + assertThrows { + GrpcQdrantTransport( + config = kdrantConfig("localhost", 6334), + channel = InProcessChannelBuilder.forName("no-server").build(), + upsertBatchSize = 0, + ) + } + } +} diff --git a/kdrant-transport-rest/build.gradle.kts b/kdrant-transport-rest/build.gradle.kts index 10d618e..885b43b 100644 --- a/kdrant-transport-rest/build.gradle.kts +++ b/kdrant-transport-rest/build.gradle.kts @@ -1,3 +1,4 @@ +import org.gradle.api.artifacts.component.ModuleComponentIdentifier import org.gradle.api.tasks.testing.logging.TestExceptionFormat plugins { @@ -35,6 +36,31 @@ dependencies { testImplementation(libs.testcontainers.qdrant) } +// The other half of the opt-in argument: the gRPC engine costs a REST user nothing only if a build +// that depends on this module resolves none of it. Checked on every `check` rather than argued in the +// README, because a transitive dependency added later would make the claim quietly false and the +// footprint table in the README rests on it. +val forbiddenGroups = setOf("io.grpc", "com.google.protobuf", "io.netty", "com.google.guava") +val runtimeArtifacts = configurations.named("runtimeClasspath").flatMap { it.incoming.artifacts.resolvedArtifacts } + +val verifyRestEngineFootprint by tasks.registering { + description = "Fails if the REST engine's runtime classpath resolves gRPC, protobuf, Netty or Guava." + val artifacts = runtimeArtifacts + val groups = forbiddenGroups + doLast { + val found = artifacts.get() + .mapNotNull { it.id.componentIdentifier as? ModuleComponentIdentifier } + .filter { it.group in groups } + .map { "${it.group}:${it.module}:${it.version}" } + .sorted() + require(found.isEmpty()) { + "the REST engine must not resolve the gRPC stack, and it resolved:\n" + found.joinToString("\n") { " $it" } + } + } +} + +tasks.named("check") { dependsOn(verifyRestEngineFootprint) } + tasks.test { useJUnitPlatform() testLogging { From f275433076ab6f27e1a0a976f98b0bf52956f2b5 Mon Sep 17 00:00:00 2001 From: TonyTonyCoder11 Date: Fri, 31 Jul 2026 15:45:43 +0200 Subject: [PATCH 2/3] Constrain the gRPC engine from the BOM The BOM exists so the modules cannot drift apart, and a module missing from it drifts by default. Importing the platform and then naming kdrant-transport-grpc without a version would resolve nothing. --- kdrant-bom/build.gradle.kts | 1 + 1 file changed, 1 insertion(+) diff --git a/kdrant-bom/build.gradle.kts b/kdrant-bom/build.gradle.kts index b713540..a9ea80c 100644 --- a/kdrant-bom/build.gradle.kts +++ b/kdrant-bom/build.gradle.kts @@ -9,6 +9,7 @@ dependencies { constraints { api("io.github.nacode-studios:kdrant-core:${project.version}") api("io.github.nacode-studios:kdrant-transport-rest:${project.version}") + api("io.github.nacode-studios:kdrant-transport-grpc:${project.version}") api("io.github.nacode-studios:kdrant-spring-boot-starter:${project.version}") api("io.github.nacode-studios:kdrant-spring-ai:${project.version}") api("io.github.nacode-studios:kdrant-langchain4j:${project.version}") From 0e891811300a65e7d4dd38df9288d67118a9afea Mon Sep 17 00:00:00 2001 From: TonyTonyCoder11 Date: Fri, 31 Jul 2026 15:51:19 +0200 Subject: [PATCH 3/3] Leave a selector off when the caller did not choose one The shared contract caught this on its first run against a real server, in the one test that reads a payload back from a scroll nobody configured. A scroll and a retrieve default with_payload to true on the server. The engine was translating "the caller said nothing" into an explicit enable = false, which is a different request: every scroll and every retrieve that did not ask for a payload silently came back without one. The REST engine omits the field and gets the default, which is why the two disagreed and why only a test that ran against both could see it. Absent now stays absent, on the payload selector and the vector selector alike, and a unit test asserts the fields are unset rather than false. --- .../transport/grpc/GrpcQdrantTransport.kt | 12 ++++---- .../dev/kdrant/transport/grpc/QueryMapping.kt | 8 +++--- .../kdrant/transport/grpc/RequestMapping.kt | 28 +++++++++++++------ .../transport/grpc/GrpcQdrantTransportTest.kt | 15 ++++++++++ 4 files changed, 44 insertions(+), 19 deletions(-) diff --git a/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransport.kt b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransport.kt index e25b073..1e0e137 100644 --- a/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransport.kt +++ b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransport.kt @@ -234,12 +234,12 @@ public class GrpcQdrantTransport internal constructor( ): List = call(name) { points.deadlined() .get( - Points.GetPoints.newBuilder() - .setCollectionName(name) - .addAllIds(ids.map(PointMapping::idToProto)) - .setWithPayload(RequestMapping.withPayload(withPayload)) - .setWithVectors(RequestMapping.withVectors(withVector)) - .build(), + Points.GetPoints.newBuilder().apply { + collectionName = name + addAllIds(ids.map(PointMapping::idToProto)) + RequestMapping.withPayload(withPayload)?.let { setWithPayload(it) } + RequestMapping.withVectors(withVector)?.let { setWithVectors(it) } + }.build(), ) .resultList .map(PointMapping::recordToModel) diff --git a/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/QueryMapping.kt b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/QueryMapping.kt index 0eb794a..025f998 100644 --- a/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/QueryMapping.kt +++ b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/QueryMapping.kt @@ -43,8 +43,8 @@ internal object QueryMapping { RequestMapping.searchParams(request.params)?.let { params = it } request.lookupFrom?.let { lookupFrom = lookupLocation(it) } RequestMapping.shardKeySelector(request.shardKey)?.let { shardKeySelector = it } - withPayload = RequestMapping.withPayload(request.withPayload) - withVectors = RequestMapping.withVectors(request.withVector) + RequestMapping.withPayload(request.withPayload)?.let { withPayload = it } + RequestMapping.withVectors(request.withVector)?.let { withVectors = it } }.build() fun queryGroups(collection: String, request: SearchGroupsRequest): Points.QueryPointGroups = @@ -60,8 +60,8 @@ internal object QueryMapping { request.scoreThreshold?.let { scoreThreshold = it.toFloat() } RequestMapping.searchParams(request.params)?.let { params = it } request.lookupFrom?.let { lookupFrom = lookupLocation(it) } - withPayload = RequestMapping.withPayload(request.withPayload) - withVectors = RequestMapping.withVectors(request.withVector) + RequestMapping.withPayload(request.withPayload)?.let { withPayload = it } + RequestMapping.withVectors(request.withVector)?.let { withVectors = it } }.build() fun searchMatrix(collection: String, request: SearchMatrixRequest): Points.SearchMatrixPoints = diff --git a/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/RequestMapping.kt b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/RequestMapping.kt index 4674869..caa17e4 100644 --- a/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/RequestMapping.kt +++ b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/RequestMapping.kt @@ -26,22 +26,32 @@ import qdrant.Points */ internal object RequestMapping { - fun withPayload(selector: WithPayload?): Points.WithPayloadSelector = + /** + * `null` means the caller did not choose, and the field is left off the request so Qdrant applies + * its own default. Sending `enable = false` instead is not the same request: a scroll and a + * retrieve default `with_payload` to **true**, so an explicit false would silently drop the payload + * from every call that did not ask for one. The REST engine omits the field for the same reason, + * and this is the difference the shared client contract caught. + */ + fun withPayload(selector: WithPayload?): Points.WithPayloadSelector? = selector?.let { Points.WithPayloadSelector.newBuilder().apply { - when (selector) { - null, WithPayload.None -> enable = false + when (it) { + WithPayload.None -> enable = false WithPayload.All -> enable = true is WithPayload.Include -> include = Points.PayloadIncludeSelector.newBuilder() - .addAllFields(selector.fields) + .addAllFields(it.fields) .build() is WithPayload.Exclude -> exclude = Points.PayloadExcludeSelector.newBuilder() - .addAllFields(selector.fields) + .addAllFields(it.fields) .build() } }.build() + } - fun withVectors(include: Boolean?): Points.WithVectorsSelector = - Points.WithVectorsSelector.newBuilder().setEnable(include ?: false).build() + /** As [withPayload]: `null` leaves the field off rather than sending an explicit false. */ + fun withVectors(include: Boolean?): Points.WithVectorsSelector? = include?.let { + Points.WithVectorsSelector.newBuilder().setEnable(it).build() + } fun pointsSelector(selector: DeleteSelector): Points.PointsSelector = Points.PointsSelector.newBuilder().apply { @@ -97,8 +107,8 @@ internal object RequestMapping { request.offset?.let { offset = PointMapping.idToProto(it) } request.orderBy?.let { orderBy = RequestMapping.orderBy(it) } RequestMapping.shardKeySelector(request.shardKey)?.let { shardKeySelector = it } - withPayload = RequestMapping.withPayload(request.withPayload) - withVectors = RequestMapping.withVectors(request.withVector) + withPayload(request.withPayload)?.let { withPayload = it } + withVectors(request.withVector)?.let { withVectors = it } }.build() fun pointVectors(vectors: PointVectors): Points.PointVectors = Points.PointVectors.newBuilder() diff --git a/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransportTest.kt b/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransportTest.kt index bd6a239..7e31830 100644 --- a/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransportTest.kt +++ b/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransportTest.kt @@ -288,6 +288,21 @@ class GrpcQdrantTransportTest { assertEquals(JsonPrimitive(3L), page.points.single().orderValue) } + @Test + fun `a request that did not choose leaves the selectors off, so Qdrant applies its own defaults`() = + runTest { + // Scroll and retrieve default with_payload to true on the server. Sending an explicit + // `enable = false` for "the caller said nothing" is a different request, and it silently + // drops the payload from every call that did not ask for one. + transport.scroll("docs", ScrollRequest(limit = 10)) + transport.retrieve("docs", listOf(PointId.num(1)), withPayload = null, withVector = null) + + assertFalse(points.scrolls.single().hasWithPayload()) + assertFalse(points.scrolls.single().hasWithVectors()) + assertFalse(points.gets.single().hasWithPayload()) + assertFalse(points.gets.single().hasWithVectors()) + } + @Test fun `a scroll resuming from an id sends it as the offset`() = runTest { transport.scroll("docs", ScrollRequest(limit = 10, offset = PointId.num(7)))