From c1dd8ec87dd455dea17a8c9738844f855ae16d39 Mon Sep 17 00:00:00 2001
From: samuel orji <awesomeorji@gmail.com>
Date: Sun, 8 Oct 2023 19:53:37 +0100
Subject: [PATCH] create a new query error called NoKnownNodeFoundError that is
 returned when there are no known nodes found during query plan and avoid
 ambiguity

---
 scylla-cql/src/errors.rs                      |  9 +++++
 .../src/transport/load_balancing/default.rs   |  1 +
 scylla/src/transport/session.rs               |  5 ++-
 scylla/src/transport/session_test.rs          | 34 +++++++++++++++++++
 4 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/scylla-cql/src/errors.rs b/scylla-cql/src/errors.rs
index 40587cfef6..cdb357bd1d 100644
--- a/scylla-cql/src/errors.rs
+++ b/scylla-cql/src/errors.rs
@@ -46,6 +46,10 @@ pub enum QueryError {
     #[error("Request timeout: {0}")]
     RequestTimeout(String),
 
+    /// No known node found to perform query
+    #[error("No known node found: {0}")]
+    NoKnownNodeFoundError(String),
+
     /// Address translation failed
     #[error("Address translation failed: {0}")]
     TranslationError(#[from] TranslationError),
@@ -400,6 +404,10 @@ pub enum NewSessionError {
     #[error("Client timeout: {0}")]
     RequestTimeout(String),
 
+    /// No known node found to perform query
+    #[error("No known node found: {0}")]
+    NoKnownNodeFoundError(String),
+
     /// Address translation failed
     #[error("Address translation failed: {0}")]
     TranslationError(#[from] TranslationError),
@@ -478,6 +486,7 @@ impl From<QueryError> for NewSessionError {
             QueryError::UnableToAllocStreamId => NewSessionError::UnableToAllocStreamId,
             QueryError::RequestTimeout(msg) => NewSessionError::RequestTimeout(msg),
             QueryError::TranslationError(e) => NewSessionError::TranslationError(e),
+            QueryError::NoKnownNodeFoundError(e) => NewSessionError::NoKnownNodeFoundError(e),
         }
     }
 }
diff --git a/scylla/src/transport/load_balancing/default.rs b/scylla/src/transport/load_balancing/default.rs
index 3c2097c9dc..f405b0af8c 100644
--- a/scylla/src/transport/load_balancing/default.rs
+++ b/scylla/src/transport/load_balancing/default.rs
@@ -2388,6 +2388,7 @@ mod latency_awareness {
                 | QueryError::IoError(_)
                 | QueryError::ProtocolError(_)
                 | QueryError::TimeoutError
+                | QueryError::NoKnownNodeFoundError(_)
                 | QueryError::RequestTimeout(_) => true,
             }
         }
diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
index 35ff25475f..c147bdac04 100644
--- a/scylla/src/transport/session.rs
+++ b/scylla/src/transport/session.rs
@@ -1653,7 +1653,10 @@ impl Session {
         QueryFut: Future<Output = Result<ResT, QueryError>>,
         ResT: AllowedRunQueryResTType,
     {
-        let mut last_error: Option<QueryError> = None;
+        // set default error as no known found as the query plan returns an empty iterator if there are no nodes in the plan
+        let mut last_error: Option<QueryError> = Some(QueryError::NoKnownNodeFoundError(
+            "Please confirm the supplied datacenters exists".to_string(),
+        ));
         let mut current_consistency: Consistency = context
             .consistency_set_on_statement
             .unwrap_or(execution_profile.consistency);
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index 805217053d..0e7844a84b 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -2,6 +2,7 @@ use crate as scylla;
 use crate::batch::{Batch, BatchStatement};
 use crate::frame::response::result::Row;
 use crate::frame::value::ValueList;
+use crate::load_balancing::DefaultPolicy;
 use crate::prepared_statement::PreparedStatement;
 use crate::query::Query;
 use crate::retry_policy::{QueryInfo, RetryDecision, RetryPolicy, RetrySession};
@@ -2857,3 +2858,36 @@ async fn test_manual_primary_key_computation() {
         .await;
     }
 }
+
+#[tokio::test]
+async fn test_non_existent_dc_return_correct_error() {
+    let ks = "iot";
+
+    let mut host = "127.0.0.1";
+    let mut dc = "non existent dc";
+
+    let default_policy = DefaultPolicy::builder()
+        .prefer_datacenter(dc.to_string())
+        .build();
+
+    let profile = ExecutionProfile::builder()
+        .load_balancing_policy(default_policy)
+        .build();
+
+    let handle = profile.into_handle();
+
+    let session: Session = SessionBuilder::new()
+        .known_node(host)
+        .default_execution_profile_handle(handle)
+        .build()
+        .await
+        .expect("cannot create session");
+
+    let ks_stmt = format!("CREATE KEYSPACE IF NOT EXISTS {} WITH replication = {{'class': 'NetworkTopologyStrategy', '{}': 1}}", ks, dc);
+    let query_result = session.query(ks_stmt, &[]).await;
+
+    assert_matches!(
+        query_result.unwrap_err(),
+        QueryError::NoKnownNodeFoundError(_)
+    )
+}