session: automatically wait for schema agreement (#460)

Now, the session automatically waits for schema agreement after
executing a statement which modifies the schema. The rationale for this
change is that, previously, it was the user's responsibility to wait for
schema agreement and it was easy to forget about this check.

The automatic wait for schema agreement is never a bug, although it can
slow down some programs which do a lot of schema changes in a short
succession. Those applications can opt-out of this feature by disabling
the feature in the configuration and can still wait for schema agreement
manually.

The session will wait for schema agreement for a configured amount of
time, and then - if the agreement is not reached - it will return
a TimeoutError.

Fixes: #459

Author: piodul

scylla/src/transport/connection.rs

@@ -158,6 +158,13 @@ impl QueryResponse {
         }
     }
 
+    pub fn as_schema_change(&self) -> Option<&result::SchemaChange> {
+        match &self.response {
+            Response::Result(result::Result::SchemaChange(sc)) => Some(sc),
+            _ => None,
+        }
+    }
+
     pub fn into_query_result(self) -> Result<QueryResult, QueryError> {
         let (rows, paging_state, col_specs) = match self.response {
             Response::Error(err) => return Err(err.into()),

scylla/src/transport/session.rs

@@ -11,7 +11,7 @@ use std::sync::Arc;
 use std::time::Duration;
 use tokio::net::lookup_host;
 use tokio::time::timeout;
-use tracing::{debug, trace, trace_span, Instrument};
+use tracing::{debug, error, trace, trace_span, Instrument};
 use uuid::Uuid;
 
 use super::connection::QueryResponse;
@@ -62,6 +62,7 @@ pub struct Session {
     speculative_execution_policy: Option<Arc<dyn SpeculativeExecutionPolicy>>,
     metrics: Arc<Metrics>,
     default_consistency: Consistency,
+    auto_await_schema_agreement_timeout: Option<Duration>,
 }
 
 /// Configuration options for [`Session`].
@@ -113,6 +114,10 @@ pub struct SessionConfig {
 
     /// Interval of sending keepalive requests
     pub keepalive_interval: Option<Duration>,
+
+    /// 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>,
 }
 
 /// Describes database server known on Session startup.
@@ -154,6 +159,7 @@ impl SessionConfig {
             default_consistency: Consistency::LocalQuorum,
             fetch_schema_metadata: true,
             keepalive_interval: None,
+            auto_await_schema_agreement_timeout: Some(std::time::Duration::from_secs(60)),
         }
     }
 
@@ -344,6 +350,7 @@ impl Session {
             speculative_execution_policy: config.speculative_execution_policy,
             metrics: Arc::new(Metrics::new()),
             default_consistency: config.default_consistency,
+            auto_await_schema_agreement_timeout: config.auto_await_schema_agreement_timeout,
         };
 
         if let Some(keyspace_name) = config.used_keyspace {
@@ -446,6 +453,8 @@ impl Session {
             .instrument(span)
             .await?;
         self.handle_set_keyspace_response(&response).await?;
+        self.handle_auto_await_schema_agreement(&query.contents, &response)
+            .await?;
 
         response.into_query_result()
     }
@@ -466,6 +475,28 @@ impl Session {
         Ok(())
     }
 
+    async fn handle_auto_await_schema_agreement(
+        &self,
+        contents: &str,
+        response: &QueryResponse,
+    ) -> Result<(), QueryError> {
+        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);
+            }
+        }
+
+        Ok(())
+    }
+
     /// Run a simple query with paging\
     /// This method will query all pages of the result\
     ///
@@ -713,6 +744,8 @@ impl Session {
             .instrument(span)
             .await?;
         self.handle_set_keyspace_response(&response).await?;
+        self.handle_auto_await_schema_agreement(prepared.get_statement(), &response)
+            .await?;
 
         response.into_query_result()
     }

scylla/src/transport/session_builder.rs

@@ -500,6 +500,46 @@ impl SessionBuilder {
         self.config.keepalive_interval = Some(interval);
         self
     }
+
+    /// Enables automatic wait for schema agreement and sets the timeout for it.
+    /// By default, it is enabled and the timeout is 60 seconds.
+    ///
+    /// # Example
+    /// ```
+    /// # use scylla::{Session, SessionBuilder};
+    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
+    /// let session: Session = SessionBuilder::new()
+    ///     .known_node("127.0.0.1:9042")
+    ///     .auto_schema_agreement_timeout(std::time::Duration::from_secs(120))
+    ///     .build()
+    ///     .await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn auto_schema_agreement_timeout(mut self, timeout: Duration) -> Self {
+        self.config.auto_await_schema_agreement_timeout = Some(timeout);
+        self
+    }
+
+    /// Disables automatic wait for schema agreement.
+    /// By default, it is enabled and the timeout is 60 seconds.
+    ///
+    /// # Example
+    /// ```
+    /// # use scylla::{Session, SessionBuilder};
+    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
+    /// let session: Session = SessionBuilder::new()
+    ///     .known_node("127.0.0.1:9042")
+    ///     .no_auto_schema_agreement()
+    ///     .build()
+    ///     .await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn no_auto_schema_agreement(mut self) -> Self {
+        self.config.auto_await_schema_agreement_timeout = None;
+        self
+    }
 }
 
 /// Creates a [`SessionBuilder`] with default configuration, same as [`SessionBuilder::new`]