Repository: scylladb/scylla-rust-driver
Pull request #276 (2021-07-23 13:56:18)

Add getting column specification by name
This series is a first baby step towards solving #271 - which is about enabling users to get column values from result rows by name instead of by index.

## 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.


psarna: (some extra unrelated commits were initially pushed by mistake, fixed)
psarna: > Looks good, but we should expose column specification information in the `RowIterator` struct, too.

You're right, column spec should probably become part of `ReceivedPage`, which gets propagated from the iterator worker to the iterator. I'd strongly suggest splitting this to a follow-up, it has two advantages:
 - non-iterator queries can start using the new way of getting proper indexes by name and I'm not sure when I'll be sending a new patchset
 - somebody else might pick it up and deliver it sooner than me (:

psarna: v2:
 - `s/get_spec_by_name/get_column_spec/g`

Patch:
diff --git a/examples/Cargo.toml b/examples/Cargo.toml
index ab64c485..7aab1fb8 100644
--- a/examples/Cargo.toml
+++ b/examples/Cargo.toml
@@ -75,3 +75,7 @@ path = "clone.rs"
 [[example]]
 name = "speculative-execution"
 path = "speculative-execution.rs"
+
+[[example]]
+name = "get_by_name"
+path = "get_by_name.rs"
diff --git a/examples/get_by_name.rs b/examples/get_by_name.rs
new file mode 100644
index 00000000..d02af56c
--- /dev/null
+++ b/examples/get_by_name.rs
@@ -0,0 +1,54 @@
+use anyhow::{anyhow, Result};
+use scylla::transport::session::Session;
+use scylla::SessionBuilder;
+use std::env;
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    tracing_subscriber::fmt::init();
+    let uri = env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1".to_string());
+
+    println!("Connecting to {} ...", uri);
+
+    let session: Session = SessionBuilder::new().known_node(uri).build().await?;
+
+    session.query("CREATE KEYSPACE IF NOT EXISTS ks WITH REPLICATION = {'class' : 'SimpleStrategy', 'replication_factor' : 1}", &[]).await?;
+
+    session
+        .query(
+            "CREATE TABLE IF NOT EXISTS ks.hello (pk int, ck int, value text, primary key (pk, ck))",
+            &[],
+        )
+        .await?;
+
+    session
+        .query(
+            "INSERT INTO ks.hello (pk, ck, value) VALUES (?, ?, ?)",
+            (3, 4, "def"),
+        )
+        .await?;
+
+    session
+        .query(
+            "INSERT INTO ks.hello (pk, ck, value) VALUES (1, 2, 'abc')",
+            &[],
+        )
+        .await?;
+
+    let query_result = session
+        .query("SELECT pk, ck, value FROM ks.hello", &[])
+        .await?;
+    let (ck_idx, _) = query_result
+        .get_column_spec("ck")
+        .ok_or_else(|| anyhow!("No ck column found"))?;
+    let (value_idx, _) = query_result
+        .get_column_spec("value")
+        .ok_or_else(|| anyhow!("No value column found"))?;
+    println!("ck           |  value");
+    println!("---------------------");
+    for row in query_result.rows.ok_or_else(|| anyhow!("no rows found"))? {
+        println!("{:?} | {:?}", row.columns[ck_idx], row.columns[value_idx]);
+    }
+
+    Ok(())
+}
diff --git a/scylla/src/frame/response/result.rs b/scylla/src/frame/response/result.rs
index c31f2e9c..a78ccdc4 100644
--- a/scylla/src/frame/response/result.rs
+++ b/scylla/src/frame/response/result.rs
@@ -293,7 +293,7 @@ impl CqlValue {
 #[derive(Debug, Clone)]
 pub struct ColumnSpec {
     pub table_spec: TableSpec,
-    name: String,
+    pub name: String,
     typ: ColumnType,
 }
 
@@ -301,7 +301,7 @@ pub struct ColumnSpec {
 pub struct ResultMetadata {
     col_count: usize,
     pub paging_state: Option<Bytes>,
-    col_specs: Vec<ColumnSpec>,
+    pub col_specs: Vec<ColumnSpec>,
 }
 
 #[derive(Debug, Clone)]
diff --git a/scylla/src/transport/connection.rs b/scylla/src/transport/connection.rs
index 1f62dca7..6eb32d73 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},
-    response::{event::Event, result, Response, ResponseOpcode},
+    response::{event::Event, result, result::ColumnSpec, Response, ResponseOpcode},
     server_event_type::EventType,
     value::{BatchValues, ValueList},
     FrameParams, SerializedRequest,
@@ -90,6 +90,18 @@ pub struct QueryResult {
     pub tracing_id: Option<Uuid>,
     /// Paging state returned from the server
     pub paging_state: Option<Bytes>,
+    /// Column specification returned from the server
+    pub col_specs: Vec<ColumnSpec>,
+}
+
+impl QueryResult {
+    // Returns a column specification for a column with given name, or None if not found
+    pub fn get_column_spec<'a>(&'a self, name: &str) -> Option<(usize, &'a ColumnSpec)> {
+        self.col_specs
+            .iter()
+            .enumerate()
+            .find(|(_id, spec)| spec.name == name)
+    }
 }
 
 /// Result of Session::batch(). Contains no rows, only some useful information.
@@ -102,10 +114,14 @@ pub struct BatchResult {
 
 impl QueryResponse {
     pub fn into_query_result(self) -> Result<QueryResult, QueryError> {
-        let (rows, paging_state) = match self.response {
+        let (rows, paging_state, col_specs) = match self.response {
             Response::Error(err) => return Err(err.into()),
-            Response::Result(result::Result::Rows(rs)) => (Some(rs.rows), rs.metadata.paging_state),
-            Response::Result(_) => (None, None),
+            Response::Result(result::Result::Rows(rs)) => (
+                Some(rs.rows),
+                rs.metadata.paging_state,
+                rs.metadata.col_specs,
+            ),
+            Response::Result(_) => (None, None, vec![]),
             _ => {
                 return Err(QueryError::ProtocolError(
                     "Unexpected server response, expected Result or Error",
@@ -118,6 +134,7 @@ impl QueryResponse {
             warnings: self.warnings,
             tracing_id: self.tracing_id,
             paging_state,
+            col_specs,
         })
     }
 }


Test patch:
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index 0d73a2ac..cb9b1715 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -44,19 +44,24 @@ async fn test_unprepared_statement() {
         .await
         .unwrap();
 
-    let rs = session
+    let query_result = session
         .query("SELECT a, b, c FROM ks.t", &[])
         .await
-        .unwrap()
-        .rows
         .unwrap();
 
+    let (a_idx, _) = query_result.get_column_spec("a").unwrap();
+    let (b_idx, _) = query_result.get_column_spec("b").unwrap();
+    let (c_idx, _) = query_result.get_column_spec("c").unwrap();
+    assert!(query_result.get_column_spec("d").is_none());
+
+    let rs = query_result.rows.unwrap();
+
     let mut results: Vec<(i32, i32, &String)> = rs
         .iter()
         .map(|r| {
-            let a = r.columns[0].as_ref().unwrap().as_int().unwrap();
-            let b = r.columns[1].as_ref().unwrap().as_int().unwrap();
-            let c = r.columns[2].as_ref().unwrap().as_text().unwrap();
+            let a = r.columns[a_idx].as_ref().unwrap().as_int().unwrap();
+            let b = r.columns[b_idx].as_ref().unwrap().as_int().unwrap();
+            let c = r.columns[c_idx].as_ref().unwrap().as_text().unwrap();
             (a, b, c)
         })
         .collect();