Allow to prioritize nodes based on their datacenter

src/config.rs

@@ -181,6 +181,10 @@ pub struct ConnectionConf {
     #[clap(long("ssl-key"), value_name = "PATH")]
     pub ssl_key_file: Option<PathBuf>,
 
+    /// Datacenter name
+    #[clap(long("datacenter"), required = false, default_value = "")]
+    pub datacenter: String,
+
     /// Default CQL query consistency level
     #[clap(long("consistency"), required = false, default_value = "LOCAL_QUORUM")]
     pub consistency: Consistency,

src/context.rs

@@ -25,6 +25,7 @@ use rune::{Any, Value};
 use rust_embed::RustEmbed;
 use scylla::_macro_internal::ColumnType;
 use scylla::frame::response::result::CqlValue;
+use scylla::load_balancing::DefaultPolicy;
 use scylla::prepared_statement::PreparedStatement;
 use scylla::transport::errors::{DbError, NewSessionError, QueryError};
 use scylla::transport::session::PoolSize;
@@ -57,8 +58,14 @@ fn ssl_context(conf: &&ConnectionConf) -> Result<Option<SslContext>, CassError>
 
 /// Configures connection to Cassandra.
 pub async fn connect(conf: &ConnectionConf) -> Result<Context, CassError> {
+    let mut policy_builder = DefaultPolicy::builder().token_aware(true);
+    let dc = &conf.datacenter;
+    if !dc.is_empty() {
+        policy_builder = policy_builder.prefer_datacenter(dc.to_owned()).permit_dc_failover(true);
+    }
     let profile = ExecutionProfile::builder()
         .consistency(conf.consistency.scylla_consistency())
+        .load_balancing_policy(policy_builder.build())
         .request_timeout(Some(Duration::from_secs(conf.request_timeout.get() as u64)))
         .build();
 

src/report.rs

@@ -510,6 +510,10 @@ impl<'a> Display for RunConfigCmp<'a> {
         }
 
         let lines: Vec<Box<dyn Display>> = vec![
+            self.line("Datacenter", "", |conf| {conf.connection.datacenter.clone()}),
+            self.line("Consistency", "", |conf| {
+                conf.connection.consistency.scylla_consistency().to_string()
+            }),
             self.line("Threads", "", |conf| Quantity::from(conf.threads)),
             self.line("Connections", "", |conf| {
                 Quantity::from(conf.connection.count)