add user defined interceptor

Author: RequiemP

src/auth.rs

@@ -25,3 +25,31 @@ impl Interceptor for TokenInterceptor {
         Ok(req)
     }
 }
+
+#[derive(Clone)]
+pub struct WrappedInterceptor<I: Send + Sync + 'static + Clone + Interceptor = TokenInterceptor> {
+    pub inner: Option<I>,
+}
+
+impl<I: Send + Sync + 'static + Clone + Interceptor> Default for WrappedInterceptor<I> {
+    fn default() -> Self {
+        Self { inner: None }
+    }
+}
+
+impl<I: Send + Sync + 'static + Clone + Interceptor> WrappedInterceptor<I> {
+    pub fn new(interceptor: I) -> Self {
+        Self {
+            inner: Some(interceptor),
+        }
+    }
+}
+
+impl<I: Send + Sync + 'static + Clone + Interceptor> Interceptor for WrappedInterceptor<I> {
+    fn call(&mut self, req: Request<()>) -> anyhow::Result<Request<()>, Status> {
+        match self.inner.as_mut() {
+            Some(inter) => inter.call(req),
+            None => Ok(req),
+        }
+    }
+}

src/qdrant_client/collection.rs

@@ -5,6 +5,7 @@ use tonic::service::Interceptor;
 use tonic::transport::Channel;
 use tonic::Status;
 
+use crate::auth::WrappedInterceptor;
 use crate::qdrant::collections_client::CollectionsClient;
 use crate::qdrant::{
     alias_operations, AliasOperations, ChangeAliases, CollectionClusterInfoRequest,
@@ -14,18 +15,18 @@ use crate::qdrant::{
     ListCollectionAliasesRequest, ListCollectionsRequest, ListCollectionsResponse, RenameAlias,
     UpdateCollection, UpdateCollectionClusterSetupRequest, UpdateCollectionClusterSetupResponse,
 };
-use crate::qdrant_client::{Qdrant, QdrantResult};
+use crate::qdrant_client::{GenericQdrant, QdrantResult};
 
 /// # Collection operations
 ///
 /// Create, update and delete collections, manage collection aliases and collection cluster
 /// configuration.
 ///
 /// Documentation: <https://qdrant.tech/documentation/concepts/collections/>
-impl<I: Send + Sync + 'static + Clone + Interceptor> Qdrant<I> {
+impl<I: Send + Sync + 'static + Clone + Interceptor> GenericQdrant<I> {
     pub(super) async fn with_collections_client<T, O: Future<Output = Result<T, Status>>>(
         &self,
-        f: impl Fn(CollectionsClient<InterceptedService<Channel, I>>) -> O,
+        f: impl Fn(CollectionsClient<InterceptedService<Channel, WrappedInterceptor<I>>>) -> O,
     ) -> QdrantResult<T> {
         let result = self
             .channel
@@ -441,6 +442,7 @@ mod tests {
         CountPointsBuilder, Distance, PointStruct, SearchPointsBuilder, UpsertPointsBuilder,
         VectorParamsBuilder,
     };
+    use crate::Qdrant;
 
     #[tokio::test]
     async fn create_collection_and_do_the_search() -> QdrantResult<()> {

src/qdrant_client/config.rs

@@ -2,18 +2,21 @@ use std::time::Duration;
 
 use tonic::service::Interceptor;
 
-use crate::{auth::TokenInterceptor, Qdrant, QdrantError};
+use crate::auth::{TokenInterceptor, WrappedInterceptor};
+use crate::qdrant_client::GenericQdrant;
+use crate::QdrantError;
 
-struct DefaultConfigValues {
+struct DefaultConfigValues<I: Send + Sync + 'static + Clone + Interceptor = TokenInterceptor> {
     timeout: Duration,
     connect_timeout: Duration,
     keep_alive_while_idle: bool,
     compression: Option<CompressionEncoding>,
     check_compatibility: bool,
     pool_size: usize,
+    interceptor: WrappedInterceptor<I>,
 }
 
-impl Default for DefaultConfigValues {
+impl<I: Send + Sync + 'static + Clone + Interceptor> Default for DefaultConfigValues<I> {
     fn default() -> Self {
         Self {
             timeout: Duration::from_secs(5),
@@ -22,6 +25,7 @@ impl Default for DefaultConfigValues {
             compression: None,
             check_compatibility: true,
             pool_size: 3,
+            interceptor: WrappedInterceptor::<I>::default(),
         }
     }
 }
@@ -65,11 +69,11 @@ pub struct QdrantConfig<I: Send + Sync + 'static + Clone + Interceptor = TokenIn
     pub pool_size: usize,
 
     /// The interceptor to use for modifying requests
-    pub interceptor: I,
+    pub interceptor: WrappedInterceptor<I>,
 }
 
 impl<I: Send + Sync + 'static + Clone + Interceptor> QdrantConfig<I> {
-    fn with_defaults(url: &str, interceptor: I) -> Self {
+    fn with_defaults(url: &str) -> Self {
         let defaults = DefaultConfigValues::default();
         Self {
             uri: url.to_string(),
@@ -79,12 +83,20 @@ impl<I: Send + Sync + 'static + Clone + Interceptor> QdrantConfig<I> {
             compression: defaults.compression,
             check_compatibility: defaults.check_compatibility,
             pool_size: defaults.pool_size,
-            interceptor,
+            interceptor: defaults.interceptor,
         }
     }
 
-    pub fn from_url_with_interceptor(url: &str, interceptor: I) -> Self {
-        Self::with_defaults(url, interceptor)
+    /// 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)
     }
 
     /// Keep the connection alive while idle
@@ -136,6 +148,33 @@ impl<I: Send + Sync + 'static + Clone + Interceptor> QdrantConfig<I> {
         self
     }
 
+    /// Set the interceptor to use for this client
+    ///
+    /// ```no_run
+    /// use qdrant_client::GenericQdrant;
+    /// 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 = GenericQdrant<CustomInterceptor>::from_url("http://localhost:6334")
+    /// .interceptor(CustomInterceptor)
+    /// .build()?;
+    ///# Ok(())
+    ///# }
+    /// ```
+    pub fn interceptor(mut self, interceptor: I) -> Self {
+        self.interceptor = WrappedInterceptor::new(interceptor);
+        self
+    }
+
     /// Set the timeout for this client
     ///
     /// Also see [`timeout()`](fn@Self::timeout).
@@ -164,9 +203,16 @@ impl<I: Send + Sync + 'static + Clone + Interceptor> QdrantConfig<I> {
         self.compression = compression;
     }
 
+    /// Set the interceptor to use for this client
+    ///
+    /// Also see [`interceptor()`](fn@Self::interceptor).
+    pub fn set_interceptor(&mut self, interceptor: I) {
+        self.interceptor = WrappedInterceptor::new(interceptor);
+    }
+
     /// Build the configured [`Qdrant`] client
-    pub fn build(self) -> Result<Qdrant<I>, QdrantError> {
-        Qdrant::new(self)
+    pub fn build(self) -> Result<GenericQdrant<I>, QdrantError> {
+        GenericQdrant::new(self)
     }
 
     pub fn skip_compatibility_check(mut self) -> Self {
@@ -182,18 +228,6 @@ impl<I: Send + Sync + 'static + Clone + Interceptor> QdrantConfig<I> {
 }
 
 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.
@@ -209,15 +243,16 @@ impl QdrantConfig<TokenInterceptor> {
     ///     .build();
     /// ```
     pub fn api_key(mut self, api_key: impl AsOptionApiKey) -> Self {
-        self.interceptor = TokenInterceptor::new(api_key.api_key());
+        self.interceptor = WrappedInterceptor::new(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()));
+        self.interceptor =
+            WrappedInterceptor::new(TokenInterceptor::new(Some(api_key.to_string())));
     }
 }
 
@@ -226,7 +261,7 @@ impl QdrantConfig<TokenInterceptor> {
 /// Connects to `http://localhost:6334` without an API key.
 impl Default for QdrantConfig<TokenInterceptor> {
     fn default() -> Self {
-        Self::with_defaults("http://localhost:6334", TokenInterceptor::new(None))
+        Self::with_defaults("http://localhost:6334")
     }
 }
 

src/qdrant_client/index.rs

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

src/qdrant_client/mod.rs

@@ -21,7 +21,7 @@ use tonic::service::Interceptor;
 use tonic::transport::{Channel, Uri};
 use tonic::Status;
 
-use crate::auth::TokenInterceptor;
+use crate::auth::{TokenInterceptor, WrappedInterceptor};
 use crate::channel_pool::ChannelPool;
 use crate::qdrant::{qdrant_client, HealthCheckReply, HealthCheckRequest};
 use crate::qdrant_client::config::QdrantConfig;
@@ -33,6 +33,7 @@ pub type QdrantResult<T> = Result<T, QdrantError>;
 
 /// A builder for [`Qdrant`]
 pub type QdrantBuilder = QdrantConfig;
+pub type Qdrant = GenericQdrant<TokenInterceptor>;
 
 /// API client to interact with a [Qdrant](https://qdrant.tech/) server.
 ///
@@ -84,7 +85,7 @@ 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<I: Send + Sync + 'static + Clone + Interceptor = TokenInterceptor> {
+pub struct GenericQdrant<I: Send + Sync + 'static + Clone + Interceptor> {
     /// Client configuration
     pub config: QdrantConfig<I>,
 
@@ -95,7 +96,7 @@ pub struct Qdrant<I: Send + Sync + 'static + Clone + Interceptor = TokenIntercep
 /// # Construct and connect
 ///
 /// Methods to construct a new Qdrant client.
-impl<I: Send + Sync + 'static + Clone + Interceptor> Qdrant<I> {
+impl<I: Send + Sync + 'static + Clone + Interceptor> GenericQdrant<I> {
     /// Create a new Qdrant client.
     ///
     /// Constructs the client and connects based on the given [`QdrantConfig`](config::QdrantConfig).
@@ -163,42 +164,34 @@ impl<I: Send + Sync + 'static + Clone + Interceptor> Qdrant<I> {
         Ok(client)
     }
 
-    /// Build a new Qdrant client with the given URL and custom interceptor.
+    /// Build a new Qdrant client with the given URL.
     ///
     /// ```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_with_interceptor(
-    ///     "http://localhost:6334",
-    ///     CustomInterceptor
-    /// ).build()?;
+    /// let client = Qdrant::from_url("http://localhost:6334").build()?;
     ///# Ok(())
     ///# }
     /// ```
-    pub fn from_url_with_interceptor(url: &str, interceptor: I) -> QdrantConfig<I> {
-        QdrantConfig::<I>::from_url_with_interceptor(url, interceptor)
+    ///
+    /// See more ways to set up the client [here](Self#set-up).
+    pub fn from_url(url: &str) -> QdrantConfig<I> {
+        QdrantConfig::<I>::from_url(url)
     }
 
     /// Wraps a channel with the configured interceptor
-    fn with_interceptor(&self, channel: Channel) -> InterceptedService<Channel, I> {
+    fn with_interceptor(
+        &self,
+        channel: Channel,
+    ) -> InterceptedService<Channel, WrappedInterceptor<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, I>>) -> O,
+        f: impl Fn(qdrant_client::QdrantClient<InterceptedService<Channel, WrappedInterceptor<I>>>) -> O,
     ) -> QdrantResult<T> {
         let result = self
             .channel
@@ -240,21 +233,3 @@ impl<I: Send + Sync + 'static + Clone + Interceptor> Qdrant<I> {
         .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

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