Repository: scylladb/scylla-rust-driver
Pull request #223 (2021-04-19 11:55:44)

Schema agreement
## 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.
- [x] 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.
- [x] I added appropriate `Fixes:` annotations to PR description.

Fixes: https://github.com/scylladb/scylla-rust-driver/issues/83

Created waiting for schema agreement mechanism for `Session` and setting the interval of that check in `SessionBuilder`.
User can now check local schema version and await for schema agreement with timeout or with endless timeout.
New methods: `Session::fetch_schema_version`, `Session::await_schema_agreement`, `Session::await_timed_schema_agreement`, `SessionBuilder::schema_agreement_interval`.

I'm not really sure what kind of tests should I provide to test that feature, so I'm asking now if adding line that awaits schema agreement in existing test is satisfactory? @psarna // solved

psarna: @Kejmer we don't have any environment for automated distributed tests (yet), so yes, it's satisfactory. But, please also test manually that this code works on a multi-node cluster as well. If I remember correctly, I sent the instructions to run a local multi-node cluster with ccm to @havaker once, so he might be able to give some more guidance (https://github.com/scylladb/scylla/wiki/Using-CCM)  
Kejmer: We'll let you know when we test it on multi-node cluster
Jasperav: > 
> 
> @Kejmer we don't have any environment for automated distributed tests (yet), so yes, it's satisfactory. But, please also test manually that this code works on a multi-node cluster as well. If I remember correctly, I sent the instructions to run a local multi-node cluster with ccm to @havaker once, so he might be able to give some more guidance (https://github.com/scylladb/scylla/wiki/Using-CCM)

@psarna we can add a cluster to the CI like this: https://github.com/Jasperav/docker-compose-check/blob/master/docker-compose.yml. Sadly Scylla sometimes fails to startup with docker compose :( https://github.com/scylladb/scylla/issues/8458
psarna: Yeah, eventually we'd like to have a multi-node cluster CI, but it shouldn't block this series. @Kejmer please rebase on top of master, since there are some conflicts (probably due to events getting merged)

Patch:
diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
index 902709ae..941c3181 100644
--- a/.github/workflows/rust.yml
+++ b/.github/workflows/rust.yml
@@ -22,13 +22,13 @@ jobs:
           - ${{ github.workspace }}:/workspace
     steps:
     - uses: actions/checkout@v2
+    - name: Run tests
+      # test threads must be one because else database tests will run in parallel and will result in flaky tests
+      run: cargo test --verbose -- --test-threads=1
     - name: Format check
       run: cargo fmt --verbose --all -- --check
     - name: Clippy check
       run: cargo clippy --verbose --examples --tests -- -D warnings
     - name: Build
       run: cargo build --verbose --examples
-    - name: Run tests
-      # test threads must be one because else database tests will run in parallel and will result in flaky tests
-      run: cargo test --verbose -- --test-threads=1
       
\ No newline at end of file
diff --git a/book/src/SUMMARY.md b/book/src/SUMMARY.md
index b866ae2d..2502d7c2 100644
--- a/book/src/SUMMARY.md
+++ b/book/src/SUMMARY.md
@@ -20,6 +20,7 @@
     - [Batch statement](queries/batch.md)
     - [Paged query](queries/paged.md)
     - [USE keyspace](queries/usekeyspace.md)
+    - [Schema agreement](queries/schema_agreement.md)
 
 - [Data Types](data-types/data-types.md)
     - [Bool, Tinyint, Smallint, Int, Bigint, Float, Double](data-types/primitive.md)
diff --git a/book/src/queries/schema_agreement.md b/book/src/queries/schema_agreement.md
new file mode 100644
index 00000000..412d416c
--- /dev/null
+++ b/book/src/queries/schema_agreement.md
@@ -0,0 +1,90 @@
+# Schema agreement
+
+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.
+
+### Checking schema version
+`Session::fetch_schema_version` returns an `Uuid` of local node's schema version. 
+
+```rust
+# extern crate scylla;
+# use scylla::Session;
+# use std::error::Error;
+# async fn check_only_compiles(session: &Session) -> Result<(), Box<dyn Error>> {
+println!("Local schema version is: {}", session.fetch_schema_version().await?);
+# Ok(())
+# }
+```
+
+### Awaiting schema agreement
+
+`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;
+# use scylla::Session;
+# use std::error::Error;
+# async fn check_only_compiles(session: &Session) -> Result<(), Box<dyn Error>> {
+session.await_schema_agreement().await?;
+# Ok(())
+# }
+```
+
+### Awaiting with timeout
+We can also set timeout in miliseconds 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 miliseconds but it can be changed with `SessionBuilder::schema_agreement_interval`.
+
+
+```rust
+# extern crate scylla;
+# use scylla::SessionBuilder;
+# use std::error::Error;
+# use std::time::Duration;
+# async fn check_only_compiles() -> Result<(), Box<dyn Error>> {
+SessionBuilder::new()
+    .known_node("127.0.0.1:9042")
+    .schema_agreement_interval(Duration::from_secs(1))
+    .build()
+    .await?;
+# Ok(())
+# }
+```
+
+### 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.
+
+
+```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? { 
+    println!("SCHEMA AGREED");
+} else {
+    println!("SCHEMA IS NOT IN AGREEMENT");
+}
+# Ok(())
+# }
+```
+
+
diff --git a/examples/Cargo.toml b/examples/Cargo.toml
index accb9a53..ed1f2bc8 100644
--- a/examples/Cargo.toml
+++ b/examples/Cargo.toml
@@ -60,3 +60,7 @@ path = "cql-time-types.rs"
 name = "tracing"
 path = "tracing.rs"
 
+[[example]]
+name = "schema_agreement"
+path = "schema_agreement.rs"
+
diff --git a/examples/schema_agreement.rs b/examples/schema_agreement.rs
new file mode 100644
index 00000000..96dfb3f9
--- /dev/null
+++ b/examples/schema_agreement.rs
@@ -0,0 +1,78 @@
+use anyhow::Result;
+use scylla::transport::session::{IntoTypedRows, Session};
+use scylla::SessionBuilder;
+use std::env;
+use std::time::Duration;
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    // Create connection
+    let uri = env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
+
+    println!("Connecting to {} ...", uri);
+
+    let session: Session = SessionBuilder::new()
+        .known_node(uri)
+        .schema_agreement_interval(Duration::from_secs(1)) // check every second for schema agreement if not agreed first check
+        .build()
+        .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' : 'SimpleStrategy', 'replication_factor' : 1}", &[]).await?;
+
+    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))",
+            &[],
+        )
+        .await?;
+
+    session.await_schema_agreement().await?;
+    session
+        .query("INSERT INTO ks.t (a, b, c) VALUES (?, ?, ?)", (3, 4, "def"))
+        .await?;
+
+    session.await_schema_agreement().await?;
+    session
+        .query("INSERT INTO ks.t (a, b, c) VALUES (1, 2, 'abc')", &[])
+        .await?;
+
+    let prepared = session
+        .prepare("INSERT INTO ks.t (a, b, c) VALUES (?, 7, ?)")
+        .await?;
+    session
+        .execute(&prepared, (42_i32, "I'm prepared!"))
+        .await?;
+    session
+        .execute(&prepared, (43_i32, "I'm prepared 2!"))
+        .await?;
+    session
+        .execute(&prepared, (44_i32, "I'm prepared 3!"))
+        .await?;
+
+    // Rows can be parsed as tuples
+    if let Some(rows) = session.query("SELECT a, b, c FROM ks.t", &[]).await?.rows {
+        for row in rows.into_typed::<(i32, i32, String)>() {
+            let (a, b, c) = row?;
+            println!("a, b, c: {}, {}, {}", a, b, c);
+        }
+    }
+    println!("Ok.");
+
+    let schema_version = session.fetch_schema_version().await?;
+    println!("Schema version: {}", schema_version);
+
+    Ok(())
+}
diff --git a/scylla/src/transport/connection.rs b/scylla/src/transport/connection.rs
index a5690f5f..f30c5499 100644
--- a/scylla/src/transport/connection.rs
+++ b/scylla/src/transport/connection.rs
@@ -30,7 +30,7 @@ use crate::batch::{Batch, BatchStatement};
 use crate::frame::{
     self,
     request::{self, batch, execute, query, register, Request, RequestOpcode},
-    response::{event::Event, result, Response, ResponseOpcode},
+    response::{event::Event, result, result::Row, Response, ResponseOpcode},
     server_event_type::EventType,
     value::{BatchValues, ValueList},
     FrameParams, RequestBodyWithExtensions,
@@ -38,6 +38,7 @@ use crate::frame::{
 use crate::query::Query;
 use crate::routing::ShardInfo;
 use crate::statement::prepared_statement::PreparedStatement;
+use crate::transport::session::IntoTypedRows;
 use crate::transport::Authenticator;
 use crate::transport::Authenticator::{
     AllowAllAuthenticator, CassandraAllowAllAuthenticator, CassandraPasswordAuthenticator,
@@ -45,6 +46,10 @@ use crate::transport::Authenticator::{
 };
 use crate::transport::Compression;
 
+// Queries for schema agreement
+const LOCAL_VERSION: &str = "SELECT schema_version FROM system.local WHERE key='local'";
+const PEERS_VERSION: &str = "SELECT schema_version FROM system.peers";
+
 pub struct Connection {
     submit_channel: mpsc::Sender<Task>,
     _worker_handle: RemoteHandle<()>,
@@ -445,6 +450,41 @@ impl Connection {
         }
     }
 
+    pub async fn check_schema_agreement(&self) -> Result<bool, QueryError> {
+        let peers_rows = self.fetch_peers_schema_version();
+        let local_version = self.fetch_schema_version();
+        let (peers_rows, local_version) = tokio::try_join!(peers_rows, local_version)?;
+        for row in peers_rows {
+            let (version,) = row
+                .into_typed::<(Uuid,)>()
+                .map_err(|_| QueryError::ProtocolError("Row is not uuid type as it should be"))?;
+            if local_version != version {
+                return Ok(false);
+            }
+        }
+        Ok(true)
+    }
+
+    async fn fetch_peers_schema_version(&self) -> Result<Vec<Row>, QueryError> {
+        self.query_single_page(PEERS_VERSION, &[])
+            .await?
+            .rows
+            .ok_or(QueryError::ProtocolError("Peers query returned not rows"))
+    }
+
+    pub async fn fetch_schema_version(&self) -> Result<Uuid, QueryError> {
+        let (version_id,): (Uuid,) = self
+            .query_single_page(LOCAL_VERSION, &[])
+            .await?
+            .rows
+            .ok_or(QueryError::ProtocolError("Version query returned not rows"))?
+            .into_typed::<(Uuid,)>()
+            .next()
+            .ok_or(QueryError::ProtocolError("Admin table returned empty rows"))?
+            .map_err(|_| QueryError::ProtocolError("Row is not uuid type as it should be"))?;
+        Ok(version_id)
+    }
+
     async fn send_request<R: Request>(
         &self,
         request: &R,
diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
index c0f60e3e..a524cd11 100644
--- a/scylla/src/transport/session.rs
+++ b/scylla/src/transport/session.rs
@@ -2,7 +2,9 @@ use futures::future::join_all;
 use std::future::Future;
 use std::net::SocketAddr;
 use std::sync::Arc;
+use std::time::Duration;
 use tokio::net::lookup_host;
+use tokio::time::timeout;
 use tracing::warn;
 use uuid::Uuid;
 
@@ -35,6 +37,7 @@ pub struct Session {
     cluster: Cluster,
     load_balancer: Arc<dyn LoadBalancingPolicy>,
     retry_policy: Box<dyn RetryPolicy + Send + Sync>,
+    schema_agreement_interval: Duration,
 
     metrics: Arc<Metrics>,
 }
@@ -67,6 +70,8 @@ pub struct SessionConfig {
 
     pub auth_username: Option<String>,
     pub auth_password: Option<String>,
+
+    pub schema_agreement_interval: Duration,
     /*
     These configuration options will be added in the future:
 
@@ -100,6 +105,7 @@ impl SessionConfig {
             known_nodes: Vec::new(),
             compression: None,
             tcp_nodelay: false,
+            schema_agreement_interval: Duration::from_millis(200),
             load_balancing: Arc::new(TokenAwarePolicy::new(Box::new(RoundRobinPolicy::new()))),
             used_keyspace: None,
             keyspace_case_sensitive: false,
@@ -259,6 +265,7 @@ impl Session {
             cluster,
             load_balancer: config.load_balancing,
             retry_policy: config.retry_policy,
+            schema_agreement_interval: config.schema_agreement_interval,
             metrics,
         };
 
@@ -733,6 +740,54 @@ impl Session {
 
         return Err(last_error);
     }
+
+    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_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)))
+    }
+
+    async fn schema_agreement_auxilary<ResT, QueryFut>(
+        &self,
+        do_query: impl Fn(Arc<Connection>) -> QueryFut,
+    ) -> Result<ResT, QueryError>
+    where
+        QueryFut: Future<Output = Result<ResT, QueryError>>,
+    {
+        self.run_query(
+            Statement::default(),
+            true,
+            Consistency::One,
+            self.retry_policy.new_session(),
+            |node: Arc<Node>| async move { node.random_connection().await },
+            do_query,
+        )
+        .await
+    }
+
+    pub async fn check_schema_agreement(&self) -> Result<bool, QueryError> {
+        self.schema_agreement_auxilary(|connection: Arc<Connection>| async move {
+            connection.check_schema_agreement().await
+        })
+        .await
+    }
+
+    pub async fn fetch_schema_version(&self) -> Result<Uuid, QueryError> {
+        self.schema_agreement_auxilary(|connection: Arc<Connection>| async move {
+            connection.fetch_schema_version().await
+        })
+        .await
+    }
 }
 
 /// Checks if a query sets a keyspace
diff --git a/scylla/src/transport/session_builder.rs b/scylla/src/transport/session_builder.rs
index 8aaa1b66..8656ea23 100644
--- a/scylla/src/transport/session_builder.rs
+++ b/scylla/src/transport/session_builder.rs
@@ -5,6 +5,7 @@ use super::Compression;
 use crate::transport::retry_policy::RetryPolicy;
 use std::net::SocketAddr;
 use std::sync::Arc;
+use std::time::Duration;
 
 #[cfg(feature = "ssl")]
 use openssl::ssl::SslContext;
@@ -205,6 +206,27 @@ impl SessionBuilder {
         self
     }
 
+    /// Set the delay for schema agreement check. How often driver should ask if schema is in agreement
+    /// The default is 200 miliseconds.
+    ///
+    /// # Example
+    /// ```
+    /// # use scylla::{Session, SessionBuilder};
+    /// # use std::time::Duration;
+    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
+    /// let session: Session = SessionBuilder::new()
+    ///     .known_node("127.0.0.1:9042")
+    ///     .schema_agreement_interval(Duration::from_secs(5))
+    ///     .build()
+    ///     .await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn schema_agreement_interval(mut self, timeout: Duration) -> Self {
+        self.config.schema_agreement_interval = timeout;
+        self
+    }
+
     /// Set the load balancing policy
     /// The default is Token-aware Round-robin.
     ///


Test patch:
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index 5077eecd..de9a727a 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -845,3 +845,28 @@ async fn assert_in_tracing_table(session: &Session, tracing_uuid: Uuid) {
     // If all retries failed panic with an error
     panic!("No rows for tracing with this session id!");
 }
+
+#[tokio::test]
+async fn test_fetch_schema_version() {
+    let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
+    let session = SessionBuilder::new().known_node(uri).build().await.unwrap();
+    session.fetch_schema_version().await.unwrap();
+}
+
+#[tokio::test]
+async fn test_await_schema_agreement() {
+    let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
+    let session = SessionBuilder::new().known_node(uri).build().await.unwrap();
+    session.await_schema_agreement().await.unwrap();
+}
+
+#[tokio::test]
+async fn test_await_timed_schema_agreement() {
+    use std::time::Duration;
+    let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
+    let session = SessionBuilder::new().known_node(uri).build().await.unwrap();
+    session
+        .await_timed_schema_agreement(Duration::from_millis(50))
+        .await
+        .unwrap();
+}