Repository: scylladb/scylla-rust-driver
Pull request #796 (2023-08-22 10:23:40)

session: add docstrings for AddressTranslator
Both `AddressTranslator` trait and `address_translator` `pub` field in `SessionConfig` were missing docstrings. They were added. The docstring for the trait is adapted from Java driver 3.

## Pre-review checklist

<!--
    Make sure you took care of the issues on the list.
    Put 'x' into those boxes which apply.
    You can also create the PR now and click on all relevant checkboxes.
    See CONTRIBUTING.md for more details.
-->

- [x] I have split my patch into logically separate commits.
- [x] All commit messages clearly explain what they change and why.
- ~~[ ] I added relevant tests for new features and bug fixes.~~
- [x] All commits compile, pass static checks and pass test.
- [x] PR description sums up the changes and reasons why they should be introduced.
- ~~[ ] I have provided docstrings for the public items that I want to introduce.~~
- ~~[ ] I have adjusted the documentation in `./docs/source/`.~~
- ~~[ ] I added appropriate `Fixes:` annotations to PR description.~~



Patch:
diff --git a/docs/source/queries/schema-agreement.md b/docs/source/queries/schema-agreement.md
index 2df38743..b9a519b4 100644
--- a/docs/source/queries/schema-agreement.md
+++ b/docs/source/queries/schema-agreement.md
@@ -1,33 +1,25 @@
 # Schema agreement
 
-Sometimes after performing queries some nodes have not been updated, so we need a mechanism that checks if every node have agreed on schema version.
+Sometimes after performing queries some nodes have not been updated so we need a mechanism that checks if every node have agreed schema version.
+There are four methods in `Session` that assist us. 
+Every method raise `QueryError` if something goes wrong, but they should never raise any errors, unless there is a DB or connection malfunction.
 
-### Automated awaiting schema agreement
-
-The driver automatically awaits schema agreement after a schema-altering query is executed.
-Waiting for schema agreement more than necessary is never a bug, but might slow down applications which do a lot of schema changes (e.g. a migration).
-For instance, in case where somebody wishes to create a keyspace and then a lot of tables in it, it makes sense only to wait after creating a keyspace
-and after creating all the tables rather than after every query. Therefore, the said behaviour can be disabled:
+### Checking schema version
+`Session::fetch_schema_version` returns an `Uuid` of local node's schema version. 
 
 ```rust
 # extern crate scylla;
-# use scylla::SessionBuilder;
+# use scylla::Session;
 # use std::error::Error;
-# async fn check_only_compiles() -> Result<(), Box<dyn Error>> {
-let session = SessionBuilder::new()
-    .known_node("127.0.0.1:9042")
-    .auto_await_schema_agreement(false)
-    .build()
-    .await?;
+# async fn check_only_compiles(session: &Session) -> Result<(), Box<dyn Error>> {
+println!("Local schema version is: {}", session.fetch_schema_version().await?);
 # Ok(())
 # }
 ```
 
-### Manually awaiting schema agreement
+### Awaiting schema agreement
 
-`Session::await_schema_agreement` returns a `Future` that can be `await`ed as long as schema is not in an agreement.
-However, it won't wait forever; `SessionConfig` defines a timeout that limits the time of waiting. If the timeout elapses,
-the return value is `Err(QueryError::RequestTimeout)`, otherwise it is `Ok(schema_version)`.
+`Session::await_schema_agreement` returns a `Future` that can be `await`ed on as long as schema is not in an agreement.
 
 ```rust
 # extern crate scylla;
@@ -39,10 +31,28 @@ session.await_schema_agreement().await?;
 # }
 ```
 
-### Interval of checking for schema agreement
+### Awaiting with timeout
+We can also set timeout in milliseconds with `Session::await_timed_schema_agreement`.
+It takes one argument, an `std::time::Duration` value that tells how long our driver should await for schema agreement. If the timeout is met the return value is `false` otherwise it is `true`.
+
+```rust
+# extern crate scylla;
+# use scylla::Session;
+# use std::error::Error;
+# use std::time::Duration;
+# async fn check_only_compiles(session: &Session) -> Result<(), Box<dyn Error>> {
+if session.await_timed_schema_agreement(Duration::from_secs(5)).await? { // wait for 5 seconds
+    println!("SCHEMA AGREED");
+} else {
+    println!("SCHEMA IS NOT IN AGREEMENT - TIMED OUT");
+}
+# Ok(())
+# }
+```
+
+### Checking for schema interval
+If schema is not agreed driver sleeps for a duration before checking it again. Default value is 200 milliseconds but it can be changed with `SessionBuilder::schema_agreement_interval`.
 
-If the schema is not agreed upon, the driver sleeps for a duration before checking it again. The default value is 200 milliseconds,
-but it can be changed with `SessionBuilder::schema_agreement_interval`.
 
 ```rust
 # extern crate scylla;
@@ -60,15 +70,15 @@ SessionBuilder::new()
 ```
 
 ### Checking if schema is in agreement now
+If you want to check if schema is in agreement now without retrying after failure you can use `Session::check_schema_agreement` function.
 
-If you want to check if schema is in agreement now, without retrying after failure, you can use `Session::check_schema_agreement` function.
 
 ```rust
 # extern crate scylla;
 # use scylla::Session;
 # use std::error::Error;
 # async fn check_only_compiles(session: &Session) -> Result<(), Box<dyn Error>> {
-if session.check_schema_agreement().await?.is_some() {
+if session.check_schema_agreement().await? { 
     println!("SCHEMA AGREED");
 } else {
     println!("SCHEMA IS NOT IN AGREEMENT");
@@ -76,3 +86,5 @@ if session.check_schema_agreement().await?.is_some() {
 # Ok(())
 # }
 ```
+
+
diff --git a/examples/schema_agreement.rs b/examples/schema_agreement.rs
index 08e5a593..b1825ddf 100644
--- a/examples/schema_agreement.rs
+++ b/examples/schema_agreement.rs
@@ -1,5 +1,4 @@
-use anyhow::{bail, Result};
-use scylla::transport::errors::QueryError;
+use anyhow::Result;
 use scylla::transport::session::{IntoTypedRows, Session};
 use scylla::SessionBuilder;
 use std::env;
@@ -18,17 +17,21 @@ async fn main() -> Result<()> {
         .build()
         .await?;
 
-    let schema_version = session.await_schema_agreement().await?;
-
+    let schema_version = session.fetch_schema_version().await?;
     println!("Schema version: {}", schema_version);
 
+    session.await_schema_agreement().await?; // without timeout example
     session.query("CREATE KEYSPACE IF NOT EXISTS ks WITH REPLICATION = {'class' : 'NetworkTopologyStrategy', 'replication_factor' : 1}", &[]).await?;
 
-    match session.await_schema_agreement().await {
-        Ok(_schema_version) => println!("Schema is in agreement in time"),
-        Err(QueryError::RequestTimeout(_)) => println!("Schema is NOT in agreement in time"),
-        Err(err) => bail!(err),
-    };
+    if session
+        .await_timed_schema_agreement(Duration::from_secs(5))
+        .await?
+    {
+        // with timeout example
+        println!("Timed schema is in agreement");
+    } else {
+        println!("Timed schema is NOT in agreement");
+    }
     session
         .query(
             "CREATE TABLE IF NOT EXISTS ks.t (a int, b int, c text, primary key (a, b))",
@@ -68,7 +71,7 @@ async fn main() -> Result<()> {
     }
     println!("Ok.");
 
-    let schema_version = session.await_schema_agreement().await?;
+    let schema_version = session.fetch_schema_version().await?;
     println!("Schema version: {}", schema_version);
 
     Ok(())
diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
index cc9c489c..de97b07d 100644
--- a/scylla/src/transport/session.rs
+++ b/scylla/src/transport/session.rs
@@ -33,7 +33,7 @@ use std::time::Duration;
 use tokio::net::lookup_host;
 use tokio::time::timeout;
 use tracing::warn;
-use tracing::{debug, trace, trace_span, Instrument};
+use tracing::{debug, error, trace, trace_span, Instrument};
 use uuid::Uuid;
 
 use super::cluster::ContactPoint;
@@ -55,7 +55,7 @@ use crate::frame::value::{
 use crate::prepared_statement::PreparedStatement;
 use crate::query::Query;
 use crate::routing::Token;
-use crate::statement::Consistency;
+use crate::statement::{Consistency, SerialConsistency};
 use crate::tracing::{TracingEvent, TracingInfo};
 use crate::transport::cluster::{Cluster, ClusterData, ClusterNeatDebug};
 use crate::transport::connection::{Connection, ConnectionConfig, VerifiedKeyspaceName};
@@ -80,6 +80,22 @@ use crate::authentication::AuthenticatorProvider;
 #[cfg(feature = "ssl")]
 use openssl::ssl::SslContext;
 
+/// Translates IP addresses received from ScyllaDB nodes into locally reachable addresses.
+///
+/// The driver auto-detects new ScyllaDB nodes added to the cluster through server side pushed
+/// notifications and through checking the system tables. For each node, the address the driver
+/// receives corresponds to the address set as `rpc_address` in the node yaml file. In most
+/// cases, this is the correct address to use by the driver and that is what is used by default.
+/// However, sometimes the addresses received through this mechanism will either not be reachable
+/// directly by the driver or should not be the preferred address to use to reach the node (for
+/// instance, the `rpc_address` set on ScyllaDB nodes might be a private IP, but some clients
+/// may have to use a public IP, or pass by a router, e.g. through NAT, to reach that node).
+/// This interface allows to deal with such cases, by allowing to translate an address as sent
+/// by a ScyllaDB node to another address to be used by the driver for connection.
+///
+/// Please note that the "known nodes" addresses provided while creating the [`Session`]
+/// instance are not translated, only IP address retrieved from or sent by Cassandra nodes
+/// to the driver are.
 #[async_trait]
 pub trait AddressTranslator: Send + Sync {
     async fn translate_address(
@@ -127,8 +143,7 @@ pub struct Session {
     default_execution_profile_handle: ExecutionProfileHandle,
     schema_agreement_interval: Duration,
     metrics: Arc<Metrics>,
-    schema_agreement_timeout: Duration,
-    schema_agreement_automatic_waiting: bool,
+    auto_await_schema_agreement_timeout: Option<Duration>,
     refresh_metadata_on_auto_schema_agreement: bool,
     keyspace_name: ArcSwapOption<String>,
     tracing_info_fetch_attempts: NonZeroU32,
@@ -150,7 +165,7 @@ impl std::fmt::Debug for Session {
             .field("metrics", &self.metrics)
             .field(
                 "auto_await_schema_agreement_timeout",
-                &self.schema_agreement_timeout,
+                &self.auto_await_schema_agreement_timeout,
             )
             .finish()
     }
@@ -184,6 +199,7 @@ pub struct SessionConfig {
 
     pub authenticator: Option<Arc<dyn AuthenticatorProvider>>,
 
+    pub schema_agreement_interval: Duration,
     pub connect_timeout: Duration,
 
     /// Size of the per-node connection pool, i.e. how many connections the driver should keep to each node.
@@ -208,22 +224,14 @@ pub struct SessionConfig {
     /// If `None`, connections are never closed due to lack of response to a keepalive message.
     pub keepalive_timeout: Option<Duration>,
 
-    /// How often the driver should ask if schema is in agreement.
-    pub schema_agreement_interval: Duration,
-
-    /// Controls the timeout for waiting for schema agreement.
-    /// This works both for manual awaiting schema agreement and for
-    /// automatic waiting after a schema-altering statement is sent.
-    pub schema_agreement_timeout: Duration,
-
-    /// Controls whether schema agreement is automatically awaited
-    /// after sending a schema-altering statement.
-    pub schema_agreement_automatic_waiting: bool,
-
-    /// If true, full schema metadata is fetched after successfully reaching a schema agreement.
-    /// It is true by default but can be disabled if successive schema-altering statements should be performed.
-    pub refresh_metadata_on_auto_schema_agreement: bool,
+    /// Controls the timeout for the automatic wait for schema agreement after sending a schema-altering statement.
+    /// If `None`, the automatic schema agreement is disabled.
+    pub auto_await_schema_agreement_timeout: Option<Duration>,
 
+    /// The address translator is used to translate addresses received from ScyllaDB nodes
+    /// (either with cluster metadata or with an event) to addresses that can be used to
+    /// actually connect to those nodes. This may be needed e.g. when there is NAT
+    /// between the nodes and the driver.
     pub address_translator: Option<Arc<dyn AddressTranslator>>,
 
     /// The host filter decides whether any connections should be opened
@@ -231,6 +239,10 @@ pub struct SessionConfig {
     /// re-establishing the control connection.
     pub host_filter: Option<Arc<dyn HostFilter>>,
 
+    /// If true, full schema metadata is fetched after successfully reaching a schema agreement.
+    /// It is true by default but can be disabled if successive schema-altering statements should be performed.
+    pub refresh_metadata_on_auto_schema_agreement: bool,
+
     /// If the driver is to connect to ScyllaCloud, there is a config for it.
     #[cfg(feature = "cloud")]
     pub cloud_config: Option<Arc<CloudConfig>>,
@@ -314,8 +326,7 @@ impl SessionConfig {
             fetch_schema_metadata: true,
             keepalive_interval: Some(Duration::from_secs(30)),
             keepalive_timeout: Some(Duration::from_secs(30)),
-            schema_agreement_timeout: Duration::from_secs(60),
-            schema_agreement_automatic_waiting: true,
+            auto_await_schema_agreement_timeout: Some(std::time::Duration::from_secs(60)),
             address_translator: None,
             host_filter: None,
             refresh_metadata_on_auto_schema_agreement: true,
@@ -571,8 +582,7 @@ impl Session {
             default_execution_profile_handle,
             schema_agreement_interval: config.schema_agreement_interval,
             metrics: Arc::new(Metrics::new()),
-            schema_agreement_timeout: config.schema_agreement_timeout,
-            schema_agreement_automatic_waiting: config.schema_agreement_automatic_waiting,
+            auto_await_schema_agreement_timeout: config.auto_await_schema_agreement_timeout,
             refresh_metadata_on_auto_schema_agreement: config
                 .refresh_metadata_on_auto_schema_agreement,
             keyspace_name: ArcSwapOption::default(), // will be set by use_keyspace
@@ -722,7 +732,8 @@ impl Session {
         };
 
         self.handle_set_keyspace_response(&response).await?;
-        self.handle_auto_await_schema_agreement(&response).await?;
+        self.handle_auto_await_schema_agreement(&query.contents, &response)
+            .await?;
 
         let result = response.into_query_result()?;
         span.record_result_fields(&result);
@@ -747,11 +758,20 @@ impl Session {
 
     async fn handle_auto_await_schema_agreement(
         &self,
+        contents: &str,
         response: &NonErrorQueryResponse,
     ) -> Result<(), QueryError> {
-        if self.schema_agreement_automatic_waiting {
-            if response.as_schema_change().is_some() {
-                self.await_schema_agreement().await?;
+        if let Some(timeout) = self.auto_await_schema_agreement_timeout {
+            if response.as_schema_change().is_some()
+                && !self.await_timed_schema_agreement(timeout).await?
+            {
+                // TODO: The TimeoutError should allow to provide more context.
+                // For now, print an error to the logs
+                error!(
+                    "Failed to reach schema agreement after a schema-altering statement: {}",
+                    contents,
+                );
+                return Err(QueryError::TimeoutError);
             }
 
             if self.refresh_metadata_on_auto_schema_agreement
@@ -1060,7 +1080,8 @@ impl Session {
         };
 
         self.handle_set_keyspace_response(&response).await?;
-        self.handle_auto_await_schema_agreement(&response).await?;
+        self.handle_auto_await_schema_agreement(prepared.get_statement(), &response)
+            .await?;
 
         let result = response.into_query_result()?;
         span.record_result_fields(&result);
@@ -1802,27 +1823,58 @@ impl Session {
         last_error.map(Result::Err)
     }
 
-    async fn await_schema_agreement_indefinitely(&self) -> Result<Uuid, QueryError> {
-        loop {
-            tokio::time::sleep(self.schema_agreement_interval).await;
-            if let Some(agreed_version) = self.check_schema_agreement().await? {
-                return Ok(agreed_version);
-            }
+    pub async fn await_schema_agreement(&self) -> Result<(), QueryError> {
+        while !self.check_schema_agreement().await? {
+            tokio::time::sleep(self.schema_agreement_interval).await
         }
+        Ok(())
     }
 
-    pub async fn await_schema_agreement(&self) -> Result<Uuid, QueryError> {
-        timeout(
-            self.schema_agreement_timeout,
-            self.await_schema_agreement_indefinitely(),
-        )
-        .await
-        .unwrap_or(Err(QueryError::RequestTimeout(
-            "schema agreement not reached in time".to_owned(),
-        )))
+    pub async fn await_timed_schema_agreement(
+        &self,
+        timeout_duration: Duration,
+    ) -> Result<bool, QueryError> {
+        timeout(timeout_duration, self.await_schema_agreement())
+            .await
+            .map_or(Ok(false), |res| res.and(Ok(true)))
     }
 
-    pub async fn check_schema_agreement(&self) -> Result<Option<Uuid>, QueryError> {
+    async fn schema_agreement_auxiliary<ResT, QueryFut>(
+        &self,
+        do_query: impl Fn(Arc<Connection>, Consistency, &ExecutionProfileInner) -> QueryFut,
+    ) -> Result<ResT, QueryError>
+    where
+        QueryFut: Future<Output = Result<ResT, QueryError>>,
+        ResT: AllowedRunQueryResTType,
+    {
+        let info = RoutingInfo::default();
+        let config = StatementConfig {
+            is_idempotent: true,
+            serial_consistency: Some(Some(SerialConsistency::LocalSerial)),
+            ..Default::default()
+        };
+
+        let span = RequestSpan::new_none();
+
+        match self
+            .run_query(
+                info,
+                &config,
+                self.get_default_execution_profile_handle().access(),
+                |node: Arc<Node>| async move { node.random_connection().await },
+                do_query,
+                &span,
+            )
+            .await?
+        {
+            RunQueryResult::IgnoredWriteError => Err(QueryError::ProtocolError(
+                "Retry policy has made the driver ignore schema's agreement query.",
+            )),
+            RunQueryResult::Completed(result) => Ok(result),
+        }
+    }
+
+    pub async fn check_schema_agreement(&self) -> Result<bool, QueryError> {
         let cluster_data = self.get_cluster_data();
         let connections_iter = cluster_data.iter_working_connections()?;
 
@@ -1831,7 +1883,17 @@ impl Session {
 
         let local_version: Uuid = versions[0];
         let in_agreement = versions.into_iter().all(|v| v == local_version);
-        Ok(in_agreement.then_some(local_version))
+        Ok(in_agreement)
+    }
+
+    pub async fn fetch_schema_version(&self) -> Result<Uuid, QueryError> {
+        // We ignore custom Consistency that a retry policy could decide to put here, using the default instead.
+        self.schema_agreement_auxiliary(
+            |connection: Arc<Connection>, _: Consistency, _: &ExecutionProfileInner| async move {
+                connection.fetch_schema_version().await
+            },
+        )
+        .await
     }
 
     /// Retrieves the handle to execution profile that is used by this session
@@ -2030,6 +2092,13 @@ impl RequestSpan {
         }
     }
 
+    pub(crate) fn new_none() -> Self {
+        Self {
+            span: tracing::Span::none(),
+            speculative_executions: 0.into(),
+        }
+    }
+
     pub(crate) fn record_shard_id(&self, conn: &Connection) {
         if let Some(info) = conn.get_shard_info() {
             self.span.record("shard", info.shard);
diff --git a/scylla/src/transport/session_builder.rs b/scylla/src/transport/session_builder.rs
index c0467ec6..8804e10e 100644
--- a/scylla/src/transport/session_builder.rs
+++ b/scylla/src/transport/session_builder.rs
@@ -717,8 +717,8 @@ impl<K: SessionBuilderKind> GenericSessionBuilder<K> {
         self
     }
 
-    /// Sets the timeout for waiting for schema agreement.
-    /// By default, the timeout is 60 seconds.
+    /// Enables automatic wait for schema agreement and sets the timeout for it.
+    /// By default, it is enabled and the timeout is 60 seconds.
     ///
     /// # Example
     /// ```
@@ -726,19 +726,19 @@ impl<K: SessionBuilderKind> GenericSessionBuilder<K> {
     /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
     /// let session: Session = SessionBuilder::new()
     ///     .known_node("127.0.0.1:9042")
-    ///     .schema_agreement_timeout(std::time::Duration::from_secs(120))
+    ///     .auto_schema_agreement_timeout(std::time::Duration::from_secs(120))
     ///     .build()
     ///     .await?;
     /// # Ok(())
     /// # }
     /// ```
-    pub fn schema_agreement_timeout(mut self, timeout: Duration) -> Self {
-        self.config.schema_agreement_timeout = timeout;
+    pub fn auto_schema_agreement_timeout(mut self, timeout: Duration) -> Self {
+        self.config.auto_await_schema_agreement_timeout = Some(timeout);
         self
     }
 
-    /// Controls automatic waiting for schema agreement after a schema-altering
-    /// statement is sent. By default, it is enabled.
+    /// Disables automatic wait for schema agreement.
+    /// By default, it is enabled and the timeout is 60 seconds.
     ///
     /// # Example
     /// ```
@@ -746,14 +746,14 @@ impl<K: SessionBuilderKind> GenericSessionBuilder<K> {
     /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
     /// let session: Session = SessionBuilder::new()
     ///     .known_node("127.0.0.1:9042")
-    ///     .auto_await_schema_agreement(false)
+    ///     .no_auto_schema_agreement()
     ///     .build()
     ///     .await?;
     /// # Ok(())
     /// # }
     /// ```
-    pub fn auto_await_schema_agreement(mut self, enabled: bool) -> Self {
-        self.config.schema_agreement_automatic_waiting = enabled;
+    pub fn no_auto_schema_agreement(mut self) -> Self {
+        self.config.auto_await_schema_agreement_timeout = None;
         self
     }
 


Test patch:
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index 6bc9f0e4..5b9914f9 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -1105,16 +1105,26 @@ async fn assert_in_tracing_table(session: &Session, tracing_uuid: Uuid) {
     panic!("No rows for tracing with this session id!");
 }
 
+#[tokio::test]
+async fn test_fetch_schema_version() {
+    let session = create_new_session_builder().build().await.unwrap();
+    session.fetch_schema_version().await.unwrap();
+}
+
 #[tokio::test]
 async fn test_await_schema_agreement() {
     let session = create_new_session_builder().build().await.unwrap();
-    let _schema_version = session.await_schema_agreement().await.unwrap();
+    session.await_schema_agreement().await.unwrap();
 }
 
 #[tokio::test]
 async fn test_await_timed_schema_agreement() {
+    use std::time::Duration;
     let session = create_new_session_builder().build().await.unwrap();
-    session.await_schema_agreement().await.unwrap();
+    session
+        .await_timed_schema_agreement(Duration::from_millis(50))
+        .await
+        .unwrap();
 }
 
 #[tokio::test]