examples: introduce example for enforcing request coordinator

Author: muzarski

examples/Cargo.toml

@@ -98,6 +98,10 @@ path = "get_by_name.rs"
 name = "value_list"
 path = "value_list.rs"
 
+[[example]]
+name = "enforce_coordinator"
+path = "enforce_coordinator.rs"
+
 [[example]]
 name = "custom_load_balancing_policy"
 path = "custom_load_balancing_policy.rs"

examples/enforce_coordinator.rs

@@ -0,0 +1,94 @@
+//! This example show how to enforce the target the request is sent to.
+
+use std::net::{IpAddr, SocketAddr};
+use std::sync::Arc;
+
+use anyhow::Result;
+use scylla::client::session::Session;
+use scylla::client::session_builder::SessionBuilder;
+use scylla::cluster::Node;
+use scylla::policies::load_balancing::{NodeIdentifier, SingleTargetLoadBalancingPolicy};
+use scylla::statement::prepared::PreparedStatement;
+
+/// Executes "SELECT host_id, rpc_address FROM system.local" query with `node` as the enforced target.
+/// Checks whether the result matches the expected values (i.e. ones stored in peers metadata).
+async fn query_system_local_and_verify(
+    session: &Session,
+    node: &Arc<Node>,
+    query_local: &PreparedStatement,
+) {
+    let (actual_host_id, actual_node_ip) = session
+        .execute_unpaged(query_local, ())
+        .await
+        .unwrap()
+        .into_rows_result()
+        .unwrap()
+        .single_row::<(uuid::Uuid, IpAddr)>()
+        .unwrap();
+
+    println!(
+        "queried host_id: {}; queried node_ip: {}",
+        actual_host_id, actual_node_ip
+    );
+    assert_eq!(node.host_id, actual_host_id);
+    assert_eq!(node.address.ip(), actual_node_ip);
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
+
+    let session: Session = SessionBuilder::new().known_node(uri).build().await?;
+
+    let state = session.get_cluster_state();
+    let node = state
+        .get_nodes_info()
+        .first()
+        .ok_or_else(|| anyhow::anyhow!("No nodes in metadata!"))?;
+
+    let expected_host_id = node.host_id;
+    let expected_node_ip = node.address.ip();
+
+    let mut query_local = session
+        .prepare("SELECT host_id, rpc_address FROM system.local where key='local'")
+        .await?;
+
+    // Enforce the node using `Arc<Node>`.
+    {
+        let node_identifier = NodeIdentifier::Node(Arc::clone(node));
+        println!("Enforcing target using {:?}...", node_identifier);
+        query_local.set_load_balancing_policy(Some(SingleTargetLoadBalancingPolicy::new(
+            node_identifier,
+            None,
+        )));
+
+        query_system_local_and_verify(&session, node, &query_local).await;
+    }
+
+    // Enforce the node using host_id.
+    {
+        let node_identifier = NodeIdentifier::HostId(expected_host_id);
+        println!("Enforcing target using {:?}...", node_identifier);
+        query_local.set_load_balancing_policy(Some(SingleTargetLoadBalancingPolicy::new(
+            node_identifier,
+            None,
+        )));
+
+        query_system_local_and_verify(&session, node, &query_local).await;
+    }
+
+    // Enforce the node using **untranslated** node address.
+    {
+        let node_identifier =
+            NodeIdentifier::NodeAddress(SocketAddr::new(expected_node_ip, node.address.port()));
+        println!("Enforcing target using {:?}...", node_identifier);
+        query_local.set_load_balancing_policy(Some(SingleTargetLoadBalancingPolicy::new(
+            node_identifier,
+            None,
+        )));
+
+        query_system_local_and_verify(&session, node, &query_local).await;
+    }
+
+    Ok(())
+}