feat: Client can be created with custom endpoint configuration (#39)

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "dgraph-tonic"
-version = "0.10.6"
+version = "0.11.0"
 authors = ["Selmeci <[email protected]>"]
 edition = "2018"
 description = "A rust async/sync client for Dgraph database build with Tonic crate"

README.md

@@ -41,7 +41,7 @@ Before using this client, it is highly recommended to go through [tour.dgraph.io
 
 ```toml
 [dependencies]
-dgraph-tonic = "0.10"
+dgraph-tonic = "0.11"
 ```
 
 Default feature is `dgraph-1-1`.
@@ -50,14 +50,14 @@ All avaiable features can be activeted with:
 
 ```toml
 [dependencies]
-dgraph-tonic = {version = "0.10", features = ["all"]}
+dgraph-tonic = {version = "0.11", features = ["all"]}
 ```
 
 If you want to use Dgraph v1.0.x, add this dependency:
 
 ```toml
 [dependencies]
-dgraph-tonic = { version = "0.10", features = ["dgraph-1-0"], default-features = false }
+dgraph-tonic = { version = "0.11", features = ["dgraph-1-0"], default-features = false }
 ```
 
 Supported features:
@@ -94,22 +94,52 @@ Note: Only API breakage from **dgraph-1-0* to *dgraph-1-1* is in the function `M
 The following code snippet shows it with just one endpoint.
 
 ```rust
-let client = Client::new("http://127.0.0.1:19080").expect("Dgraph client");
+use dgraph_tonic::Client;
+
+fn main() {
+  let client = Client::new("http://127.0.0.1:19080").expect("Dgraph client");
+}
 ```
 
 or you can initialize new client with a multiple endpoints. Client cannot be create with empty endpoints vector.
 
 ```rust
-let client = Client::new(vec!["http://127.0.0.1:9080","http://127.0.0.1:19080"]).expect("Dgraph client");
+use dgraph_tonic::Client;
+
+fn main() {
+  let client = Client::new(vec!["http://127.0.0.1:9080","http://127.0.0.1:19080"]).expect("Dgraph client");
+}
 ```
 
+Client can be also initialized with custom [endpoint configuration.](https://docs.rs/tonic/0.8.3/tonic/transport/struct.Endpoint.html)
+```rust
+use dgraph_tonic::{Endpoint, EndpointConfig, Client};
+
+use std::time::Duration;
+
+#[derive(Debug, Default)]
+struct EndpointWithTimeout {}
+
+impl EndpointConfig for EndpointWithTimeout {
+    fn configure_endpoint(&self, endpoint: Endpoint) -> Endpoint {
+        endpoint.timeout(Duration::from_secs(5))
+    }
+}
+
+fn main() {
+  let endpoint_config = EndpointWithTimeout::default();
+  let client = Client::new_with_endpoint_config("http://127.0.0.1:19080",endpoint_config).expect("Dgraph client");
+}
+```
+
+
 ### Create a TLS client
 
 Alternatively, secure tls client is avaible in `tls` feature:
 
 ```toml
 [dependencies]
-dgraph-tonic = { version = "0.10", features = ["tls"] }
+dgraph-tonic = { version = "0.11", features = ["tls"] }
 ```
 
 ```rust
@@ -159,7 +189,7 @@ Client is avaible in `slash-ql` feature:
 
 ```toml
 [dependencies]
-dgraph-tonic = { version = "0.10", features = ["slash-ql"] }
+dgraph-tonic = { version = "0.11", features = ["slash-ql"] }
 ```
 
 ```rust
@@ -180,7 +210,7 @@ Alternatively, synchronous clients (Tls, Acl) are avaible with `sync` feature in
 
 ```toml
 [dependencies]
-dgraph-tonic = { version = "0.10", features = ["sync"] }
+dgraph-tonic = { version = "0.11", features = ["sync"] }
 ```
 
 ```rust
@@ -586,7 +616,7 @@ This enterprise Dgraph feature which can be activated with:
 
 ```toml
 [dependencies]
-dgraph-tonic = { version = "0.10", features = ["acl"] }
+dgraph-tonic = { version = "0.11", features = ["acl"] }
 ```
 
 [Access Control List (ACL)](https://dgraph.io/docs/enterprise-features/#access-control-lists) provides access protection to your data stored in Dgraph. When the ACL feature is turned on, a client must authenticate with a username and password before executing any transactions, and is only allowed to access the data permitted by the ACL rules.

src/client/default.rs

@@ -4,13 +4,17 @@ use anyhow::Result;
 use async_trait::async_trait;
 use http::Uri;
 use std::fmt::Debug;
+use std::sync::Arc;
 use tonic::transport::Channel;
 use tracing::trace;
 use tracing_attributes::instrument;
 
 use crate::client::lazy::{ILazyChannel, LazyClient};
 use crate::client::{balance_list, rnd_item, ClientState, ClientVariant, IClient};
-use crate::{Endpoint, Endpoints, TxnBestEffortType, TxnMutatedType, TxnReadOnlyType, TxnType};
+use crate::{
+    Endpoint, EndpointConfig, Endpoints, TxnBestEffortType, TxnMutatedType, TxnReadOnlyType,
+    TxnType,
+};
 
 ///
 /// Lazy initialization of gRPC channel
@@ -19,11 +23,21 @@ use crate::{Endpoint, Endpoints, TxnBestEffortType, TxnMutatedType, TxnReadOnlyT
 pub struct LazyChannel {
     uri: Uri,
     channel: Option<Channel>,
+    endpoint_config: Option<Arc<dyn EndpointConfig>>,
 }
 
 impl LazyChannel {
     fn new(uri: Uri) -> Self {
-        Self { uri, channel: None }
+        Self {
+            uri,
+            channel: None,
+            endpoint_config: None,
+        }
+    }
+
+    fn with_endpoint_config(mut self, endpoint_config: Option<Arc<dyn EndpointConfig>>) -> Self {
+        self.endpoint_config = endpoint_config;
+        self
     }
 }
 
@@ -33,7 +47,10 @@ impl ILazyChannel for LazyChannel {
         if let Some(channel) = &self.channel {
             Ok(channel.to_owned())
         } else {
-            let endpoint: Endpoint = self.uri.to_owned().into();
+            let mut endpoint: Endpoint = self.uri.to_owned().into();
+            if let Some(endpoint_config) = &self.endpoint_config {
+                endpoint = endpoint_config.configure_endpoint(endpoint);
+            }
             let channel = endpoint.connect().await?;
             self.channel.replace(channel.to_owned());
             Ok(channel)
@@ -90,6 +107,18 @@ pub type TxnBestEffort = TxnBestEffortType<LazyClient<LazyChannel>>;
 pub type TxnMutated = TxnMutatedType<LazyClient<LazyChannel>>;
 
 impl Client {
+    fn init_clients<S: TryInto<Uri>, E: Into<Endpoints<S>> + Debug>(
+        endpoints: E,
+        endpoint_config: Option<Arc<dyn EndpointConfig>>,
+    ) -> Result<Vec<LazyClient<LazyChannel>>> {
+        Ok(balance_list(endpoints)?
+            .into_iter()
+            .map(|uri| {
+                LazyClient::new(LazyChannel::new(uri).with_endpoint_config(endpoint_config.clone()))
+            })
+            .collect())
+    }
+
     ///
     /// Create new Dgraph client for interacting v DB.
     ///
@@ -118,10 +147,63 @@ impl Client {
     #[instrument]
     pub fn new<S: TryInto<Uri>, E: Into<Endpoints<S>> + Debug>(endpoints: E) -> Result<Self> {
         let extra = Http {
-            clients: balance_list(endpoints)?
-                .into_iter()
-                .map(|uri| LazyClient::new(LazyChannel::new(uri)))
-                .collect(),
+            clients: Self::init_clients(endpoints, None)?,
+        };
+        let state = Box::new(ClientState::new());
+        trace!("New http client");
+        Ok(Self { state, extra })
+    }
+
+    ///
+    /// Create new Dgraph client with custom endpoint configuration for interacting with DB.
+    ///
+    /// The client can be backed by multiple endpoints (to the same server, or multiple servers in a cluster).
+    ///
+    /// # Arguments
+    ///
+    /// * `endpoints` - one endpoint or vector of endpoints
+    /// * `endpoint_config` - custom endpoint configuration
+    ///
+    /// # Errors
+    ///
+    /// * endpoints vector is empty
+    /// * item in vector cannot by converted into Uri
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use dgraph_tonic::{Endpoint, EndpointConfig, Client};
+    ///
+    /// use std::time::Duration;
+    ///
+    /// #[derive(Debug, Default, Clone)]
+    /// struct EndpointWithTimeout {}
+    ///
+    /// impl EndpointConfig for EndpointWithTimeout {
+    ///     fn configure_endpoint(&self, endpoint: Endpoint) -> Endpoint {
+    ///         endpoint.timeout(Duration::from_secs(5))
+    ///     }
+    /// }
+    ///
+    /// // custom configuration
+    /// let endpoint_config = EndpointWithTimeout::default();
+    /// // vector of endpoints
+    /// let client = Client::new_with_endpoint_config(vec!["http://127.0.0.1:19080", "http://127.0.0.1:19080"], endpoint_config.clone()).expect("Dgraph client");
+    /// // one endpoint
+    /// let client = Client::new_with_endpoint_config("http://127.0.0.1:19080", endpoint_config).expect("Dgraph client");
+    /// ```
+    ///
+    #[instrument]
+    pub fn new_with_endpoint_config<
+        S: TryInto<Uri>,
+        E: Into<Endpoints<S>> + Debug,
+        C: EndpointConfig + 'static,
+    >(
+        endpoints: E,
+        endpoint_config: C,
+    ) -> Result<Self> {
+        let extra = Http {
+            clients: Self::init_clients(endpoints, Some(Arc::new(endpoint_config)))?,
         };
         let state = Box::new(ClientState::new());
         trace!("New http client");

src/client/mod.rs

@@ -7,7 +7,7 @@ use http::Uri;
 use rand::Rng;
 #[cfg(any(feature = "acl", feature = "slash-ql"))]
 use tonic::codegen::InterceptedService;
-use tonic::transport::Channel;
+use tonic::transport::{Channel, Endpoint};
 
 use crate::api::dgraph_client::DgraphClient as DClient;
 use crate::api::Version;
@@ -111,6 +111,13 @@ pub enum DgraphClient {
 #[cfg(any(feature = "acl", feature = "slash-ql"))]
 pub type DgraphInterceptorClient<T> = DClient<InterceptedService<Channel, T>>;
 
+///
+/// Allow custom configuration of endpoint
+///
+pub trait EndpointConfig: Send + Sync + Debug {
+    fn configure_endpoint(&self, endpoint: Endpoint) -> Endpoint;
+}
+
 ///
 /// Marker for client variant implementation
 ///