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-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}") 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..1e0e137 --- /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().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) + } + + 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..025f998 --- /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 } + RequestMapping.withPayload(request.withPayload)?.let { withPayload = it } + RequestMapping.withVectors(request.withVector)?.let { withVectors = it } + }.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) } + RequestMapping.withPayload(request.withPayload)?.let { withPayload = it } + RequestMapping.withVectors(request.withVector)?.let { withVectors = it } + }.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..caa17e4 --- /dev/null +++ b/kdrant-transport-grpc/src/main/kotlin/dev/kdrant/transport/grpc/RequestMapping.kt @@ -0,0 +1,195 @@ +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 { + + /** + * `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 (it) { + WithPayload.None -> enable = false + WithPayload.All -> enable = true + is WithPayload.Include -> include = Points.PayloadIncludeSelector.newBuilder() + .addAllFields(it.fields) + .build() + is WithPayload.Exclude -> exclude = Points.PayloadExcludeSelector.newBuilder() + .addAllFields(it.fields) + .build() + } + }.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 { + 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(request.withPayload)?.let { withPayload = it } + withVectors(request.withVector)?.let { withVectors = it } + }.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..7e31830 --- /dev/null +++ b/kdrant-transport-grpc/src/test/kotlin/dev/kdrant/transport/grpc/GrpcQdrantTransportTest.kt @@ -0,0 +1,630 @@ +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 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))) + + 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 {