add user defined interceptor

Author: RequiemP

src/auth.rs

@@ -1,6 +1,7 @@
 use tonic::service::Interceptor;
 use tonic::{Request, Status};
 
+#[derive(Clone)]
 pub struct TokenInterceptor {
     api_key: Option<String>,
 }

src/qdrant_client/collection.rs

@@ -1,10 +1,10 @@
 use std::future::Future;
 
 use tonic::codegen::InterceptedService;
+use tonic::service::Interceptor;
 use tonic::transport::Channel;
 use tonic::Status;
 
-use crate::auth::TokenInterceptor;
 use crate::qdrant::collections_client::CollectionsClient;
 use crate::qdrant::{
     alias_operations, AliasOperations, ChangeAliases, CollectionClusterInfoRequest,
@@ -22,16 +22,16 @@ use crate::qdrant_client::{Qdrant, QdrantResult};
 /// configuration.
 ///
 /// Documentation: <https://qdrant.tech/documentation/concepts/collections/>
-impl Qdrant {
+impl<I: Send + Sync + 'static + Clone + Interceptor> Qdrant<I> {
     pub(super) async fn with_collections_client<T, O: Future<Output = Result<T, Status>>>(
         &self,
-        f: impl Fn(CollectionsClient<InterceptedService<Channel, TokenInterceptor>>) -> O,
+        f: impl Fn(CollectionsClient<InterceptedService<Channel, I>>) -> O,
     ) -> QdrantResult<T> {
         let result = self
             .channel
             .with_channel(
                 |channel| {
-                    let service = self.with_api_key(channel);
+                    let service = self.with_interceptor(channel);
                     let mut client =
                         CollectionsClient::new(service).max_decoding_message_size(usize::MAX);
                     if let Some(compression) = self.config.compression {

src/qdrant_client/config.rs

@@ -1,6 +1,30 @@
 use std::time::Duration;
 
-use crate::{Qdrant, QdrantError};
+use tonic::service::Interceptor;
+
+use crate::{auth::TokenInterceptor, Qdrant, QdrantError};
+
+struct DefaultConfigValues {
+    timeout: Duration,
+    connect_timeout: Duration,
+    keep_alive_while_idle: bool,
+    compression: Option<CompressionEncoding>,
+    check_compatibility: bool,
+    pool_size: usize,
+}
+
+impl Default for DefaultConfigValues {
+    fn default() -> Self {
+        Self {
+            timeout: Duration::from_secs(5),
+            connect_timeout: Duration::from_secs(5),
+            keep_alive_while_idle: true,
+            compression: None,
+            check_compatibility: true,
+            pool_size: 3,
+        }
+    }
+}
 
 /// Qdrant client configuration
 ///
@@ -17,7 +41,7 @@ use crate::{Qdrant, QdrantError};
 ///     .build();
 /// ```
 #[derive(Clone)]
-pub struct QdrantConfig {
+pub struct QdrantConfig<I: Send + Sync + 'static + Clone + Interceptor = TokenInterceptor> {
     /// Qdrant server URI to connect to
     pub uri: String,
 
@@ -30,9 +54,6 @@ pub struct QdrantConfig {
     /// Whether to keep idle connections active
     pub keep_alive_while_idle: bool,
 
-    /// Optional API key or token to use for authorization
-    pub api_key: Option<String>,
-
     /// Optional compression schema to use for API requests
     pub compression: Option<CompressionEncoding>,
 
@@ -42,51 +63,28 @@ pub struct QdrantConfig {
     /// Amount of concurrent connections.
     /// If set to 0 or 1, connection pools will be disabled.
     pub pool_size: usize,
+
+    /// The interceptor to use for modifying requests
+    pub interceptor: I,
 }
 
-impl QdrantConfig {
-    /// Start configuring a Qdrant client with an URL
-    ///
-    /// ```rust,no_run
-    ///# use qdrant_client::config::QdrantConfig;
-    /// let client = QdrantConfig::from_url("http://localhost:6334").build();
-    /// ```
-    ///
-    /// This is normally done through [`Qdrant::from_url`](crate::Qdrant::from_url).
-    pub fn from_url(url: &str) -> Self {
-        QdrantConfig {
+impl<I: Send + Sync + 'static + Clone + Interceptor> QdrantConfig<I> {
+    fn with_defaults(url: &str, interceptor: I) -> Self {
+        let defaults = DefaultConfigValues::default();
+        Self {
             uri: url.to_string(),
-            ..Self::default()
+            timeout: defaults.timeout,
+            connect_timeout: defaults.connect_timeout,
+            keep_alive_while_idle: defaults.keep_alive_while_idle,
+            compression: defaults.compression,
+            check_compatibility: defaults.check_compatibility,
+            pool_size: defaults.pool_size,
+            interceptor,
         }
     }
 
-    /// Set an optional API key
-    ///
-    /// # Examples
-    ///
-    /// A typical use case might be getting the key from an environment variable:
-    ///
-    /// ```rust,no_run
-    /// use qdrant_client::Qdrant;
-    ///
-    /// let client = Qdrant::from_url("http://localhost:6334")
-    ///     .api_key(std::env::var("QDRANT_API_KEY"))
-    ///     .build();
-    /// ```
-    ///
-    /// Or you might get it from some configuration:
-    ///
-    /// ```rust,no_run
-    ///# use std::collections::HashMap;
-    ///# let config: HashMap<&str, String> = HashMap::new();
-    ///# use qdrant_client::Qdrant;
-    /// let client = Qdrant::from_url("http://localhost:6334")
-    ///     .api_key(config.get("api_key"))
-    ///     .build();
-    /// ```
-    pub fn api_key(mut self, api_key: impl AsOptionApiKey) -> Self {
-        self.api_key = api_key.api_key();
-        self
+    pub fn from_url_with_interceptor(url: &str, interceptor: I) -> Self {
+        Self::with_defaults(url, interceptor)
     }
 
     /// Keep the connection alive while idle
@@ -138,13 +136,6 @@ impl QdrantConfig {
         self
     }
 
-    /// Set an API key
-    ///
-    /// Also see [`api_key()`](fn@Self::api_key).
-    pub fn set_api_key(&mut self, api_key: &str) {
-        self.api_key = Some(api_key.to_string());
-    }
-
     /// Set the timeout for this client
     ///
     /// Also see [`timeout()`](fn@Self::timeout).
@@ -174,7 +165,7 @@ impl QdrantConfig {
     }
 
     /// Build the configured [`Qdrant`] client
-    pub fn build(self) -> Result<Qdrant, QdrantError> {
+    pub fn build(self) -> Result<Qdrant<I>, QdrantError> {
         Qdrant::new(self)
     }
 
@@ -190,21 +181,52 @@ impl QdrantConfig {
     }
 }
 
+impl QdrantConfig<TokenInterceptor> {
+    /// Start configuring a Qdrant client with an URL
+    ///
+    /// ```rust,no_run
+    ///# use qdrant_client::config::QdrantConfig;
+    /// let client = QdrantConfig::from_url("http://localhost:6334").build();
+    /// ```
+    ///
+    /// This is normally done through [`Qdrant::from_url`](crate::Qdrant::from_url).
+    pub fn from_url(url: &str) -> Self {
+        Self::with_defaults(url, TokenInterceptor::new(None))
+    }
+
+    /// Set an optional API key
+    ///
+    /// This method is only available when using the default TokenInterceptor.
+    /// When you set an API key, it automatically configures the TokenInterceptor.
+    ///
+    /// # Examples
+    ///
+    /// ```rust,no_run
+    /// use qdrant_client::Qdrant;
+    ///
+    /// let client = Qdrant::from_url("http://localhost:6334")
+    ///     .api_key(std::env::var("QDRANT_API_KEY"))
+    ///     .build();
+    /// ```
+    pub fn api_key(mut self, api_key: impl AsOptionApiKey) -> Self {
+        self.interceptor = TokenInterceptor::new(api_key.api_key());
+        self
+    }
+
+    /// Set an API key
+    ///
+    /// Also see [`api_key()`](fn@Self::api_key).
+    pub fn set_api_key(&mut self, api_key: &str) {
+        self.interceptor = TokenInterceptor::new(Some(api_key.to_string()));
+    }
+}
+
 /// Default Qdrant client configuration.
 ///
 /// Connects to `http://localhost:6334` without an API key.
-impl Default for QdrantConfig {
+impl Default for QdrantConfig<TokenInterceptor> {
     fn default() -> Self {
-        Self {
-            uri: String::from("http://localhost:6334"),
-            timeout: Duration::from_secs(5),
-            connect_timeout: Duration::from_secs(5),
-            keep_alive_while_idle: true,
-            api_key: None,
-            compression: None,
-            check_compatibility: true,
-            pool_size: 3,
-        }
+        Self::with_defaults("http://localhost:6334", TokenInterceptor::new(None))
     }
 }
 

src/qdrant_client/index.rs

@@ -1,3 +1,5 @@
+use tonic::service::Interceptor;
+
 use crate::qdrant::{
     CreateFieldIndexCollection, DeleteFieldIndexCollection, PointsOperationResponse,
 };
@@ -8,7 +10,7 @@ use crate::qdrant_client::{Qdrant, QdrantResult};
 /// Manage field and payload indices in collections.
 ///
 /// Documentation: <https://qdrant.tech/documentation/concepts/indexing/>
-impl Qdrant {
+impl<I: Send + Sync + 'static + Clone + Interceptor> Qdrant<I> {
     /// Create payload index in a collection.
     ///
     /// ```no_run

src/qdrant_client/mod.rs

@@ -17,6 +17,7 @@ use std::sync::Arc;
 use std::thread;
 
 use tonic::codegen::InterceptedService;
+use tonic::service::Interceptor;
 use tonic::transport::{Channel, Uri};
 use tonic::Status;
 
@@ -83,9 +84,9 @@ pub type QdrantBuilder = QdrantConfig;
 /// - [`upsert_points`](Self::upsert_points) - insert or update points
 /// - [`query`](Self::query) - query points with similarity search
 #[derive(Clone)]
-pub struct Qdrant {
+pub struct Qdrant<I: Send + Sync + 'static + Clone + Interceptor = TokenInterceptor> {
     /// Client configuration
-    pub config: QdrantConfig,
+    pub config: QdrantConfig<I>,
 
     /// Internal connection pool
     channel: Arc<ChannelPool>,
@@ -94,11 +95,11 @@ pub struct Qdrant {
 /// # Construct and connect
 ///
 /// Methods to construct a new Qdrant client.
-impl Qdrant {
+impl<I: Send + Sync + 'static + Clone + Interceptor> Qdrant<I> {
     /// Create a new Qdrant client.
     ///
     /// Constructs the client and connects based on the given [`QdrantConfig`](config::QdrantConfig).
-    pub fn new(config: QdrantConfig) -> QdrantResult<Self> {
+    pub fn new(config: QdrantConfig<I>) -> QdrantResult<Self> {
         if config.check_compatibility {
             // create a temporary client to check compatibility
             let channel = ChannelPool::new(
@@ -162,38 +163,48 @@ impl Qdrant {
         Ok(client)
     }
 
-    /// Build a new Qdrant client with the given URL.
+    /// Build a new Qdrant client with the given URL and custom interceptor.
     ///
     /// ```no_run
     /// use qdrant_client::Qdrant;
+    /// use tonic::service::Interceptor;
+    /// use tonic::{Request, Status};
+    ///
+    /// #[derive(Clone)]
+    /// struct CustomInterceptor;
+    /// impl Interceptor for CustomInterceptor {
+    ///     fn call(&mut self, req: Request<()>) -> Result<Request<()>, Status> {
+    ///         Ok(req)
+    ///     }
+    /// }
     ///
     ///# async fn connect() -> Result<(), qdrant_client::QdrantError> {
-    /// let client = Qdrant::from_url("http://localhost:6334").build()?;
+    /// let client = Qdrant::from_url_with_interceptor(
+    ///     "http://localhost:6334",
+    ///     CustomInterceptor
+    /// ).build()?;
     ///# Ok(())
     ///# }
     /// ```
-    ///
-    /// See more ways to set up the client [here](Self#set-up).
-    pub fn from_url(url: &str) -> QdrantBuilder {
-        QdrantBuilder::from_url(url)
+    pub fn from_url_with_interceptor(url: &str, interceptor: I) -> QdrantConfig<I> {
+        QdrantConfig::<I>::from_url_with_interceptor(url, interceptor)
     }
 
-    /// Wraps a channel with a token interceptor
-    fn with_api_key(&self, channel: Channel) -> InterceptedService<Channel, TokenInterceptor> {
-        let interceptor = TokenInterceptor::new(self.config.api_key.clone());
-        InterceptedService::new(channel, interceptor)
+    /// Wraps a channel with the configured interceptor
+    fn with_interceptor(&self, channel: Channel) -> InterceptedService<Channel, I> {
+        InterceptedService::new(channel, self.config.interceptor.clone())
     }
 
     // Access to raw root qdrant API
     async fn with_root_qdrant_client<T, O: Future<Output = Result<T, Status>>>(
         &self,
-        f: impl Fn(qdrant_client::QdrantClient<InterceptedService<Channel, TokenInterceptor>>) -> O,
+        f: impl Fn(qdrant_client::QdrantClient<InterceptedService<Channel, I>>) -> O,
     ) -> QdrantResult<T> {
         let result = self
             .channel
             .with_channel(
                 |channel| {
-                    let service = self.with_api_key(channel);
+                    let service = self.with_interceptor(channel);
                     let mut client = qdrant_client::QdrantClient::new(service)
                         .max_decoding_message_size(usize::MAX);
                     if let Some(compression) = self.config.compression {
@@ -229,3 +240,21 @@ impl Qdrant {
         .await
     }
 }
+
+impl Qdrant<TokenInterceptor> {
+    /// Build a new Qdrant client with the given URL.
+    ///
+    /// ```no_run
+    /// use qdrant_client::Qdrant;
+    ///
+    ///# async fn connect() -> Result<(), qdrant_client::QdrantError> {
+    /// let client = Qdrant::from_url("http://localhost:6334").build()?;
+    ///# Ok(())
+    ///# }
+    /// ```
+    ///
+    /// See more ways to set up the client [here](Self#set-up).
+    pub fn from_url(url: &str) -> QdrantBuilder {
+        QdrantBuilder::from_url(url)
+    }
+}

src/qdrant_client/payload.rs

@@ -1,3 +1,5 @@
+use tonic::service::Interceptor;
+
 use crate::qdrant::{
     ClearPayloadPoints, DeletePayloadPoints, PointsOperationResponse, SetPayloadPoints,
 };
@@ -8,7 +10,7 @@ use crate::qdrant_client::{Qdrant, QdrantResult};
 /// Manage point payloads.
 ///
 /// Documentation: <https://qdrant.tech/documentation/concepts/payload/>
-impl Qdrant {
+impl<I: Send + Sync + 'static + Clone + Interceptor> Qdrant<I> {
     /// Set payload of points.
     ///
     /// Sets only the given payload values on a point, leaving other existing payloads in place.