Merge pull request scylladb#1249 from muzarski/refactor-column-specs

Refactor `ColumnSpecs` and adjust `PreparedStatement` API

Author: Lorak-mmk

scylla/src/client/pager.rs

@@ -646,7 +646,7 @@ impl QueryPager {
     pub fn type_check<'frame, 'metadata, RowT: DeserializeRow<'frame, 'metadata>>(
         &self,
     ) -> Result<(), TypeCheckError> {
-        RowT::type_check(self.column_specs().inner())
+        RowT::type_check(self.column_specs().as_slice())
     }
 
     /// Casts the iterator to a given row type, enabling [Stream]'ed operations
@@ -955,7 +955,7 @@ impl QueryPager {
 
     /// Returns specification of row columns
     #[inline]
-    pub fn column_specs(&self) -> ColumnSpecs<'_> {
+    pub fn column_specs(&self) -> ColumnSpecs<'_, '_> {
         ColumnSpecs::new(self.current_page.metadata().col_specs())
     }
 

scylla/src/response/query_result.rs

@@ -8,87 +8,23 @@ use scylla_cql::deserialize::row::DeserializeRow;
 use scylla_cql::deserialize::{DeserializationError, TypeCheckError};
 use scylla_cql::frame::frame_errors::ResultMetadataAndRowsCountParseError;
 use scylla_cql::frame::response::result::{
-    ColumnSpec, ColumnType, DeserializedMetadataAndRawRows, RawMetadataAndRawRows, TableSpec,
+    ColumnSpec, DeserializedMetadataAndRawRows, RawMetadataAndRawRows,
 };
 
-/// A view over specification of a table in the database.
-#[derive(Debug, Clone, Copy)]
-#[cfg_attr(test, derive(PartialEq, Eq))]
-pub struct TableSpecView<'res> {
-    table_name: &'res str,
-    ks_name: &'res str,
-}
-
-impl<'res> TableSpecView<'res> {
-    pub(crate) fn new_from_table_spec(spec: &'res TableSpec) -> Self {
-        Self {
-            table_name: spec.table_name(),
-            ks_name: spec.ks_name(),
-        }
-    }
-
-    /// The name of the table.
-    #[inline]
-    pub fn table_name(&self) -> &'res str {
-        self.table_name
-    }
-
-    /// The name of the keyspace the table resides in.
-    #[inline]
-    pub fn ks_name(&self) -> &'res str {
-        self.ks_name
-    }
-}
-
-/// A view over specification of a column returned by the database.
-#[derive(Debug, Clone, Copy)]
-#[cfg_attr(test, derive(PartialEq, Eq))]
-pub struct ColumnSpecView<'res> {
-    table_spec: TableSpecView<'res>,
-    name: &'res str,
-    typ: &'res ColumnType<'res>,
-}
-
-impl<'res> ColumnSpecView<'res> {
-    pub(crate) fn new_from_column_spec(spec: &'res ColumnSpec) -> Self {
-        Self {
-            table_spec: TableSpecView::new_from_table_spec(spec.table_spec()),
-            name: spec.name(),
-            typ: spec.typ(),
-        }
-    }
-
-    /// Returns a view over specification of the table the column is part of.
-    #[inline]
-    pub fn table_spec(&self) -> TableSpecView<'res> {
-        self.table_spec
-    }
-
-    /// The column's name.
-    #[inline]
-    pub fn name(&self) -> &'res str {
-        self.name
-    }
-
-    /// The column's CQL type.
-    #[inline]
-    pub fn typ(&self) -> &'res ColumnType {
-        self.typ
-    }
-}
-
 /// A view over specification of columns returned by the database.
 #[derive(Debug, Clone, Copy)]
-pub struct ColumnSpecs<'res> {
-    specs: &'res [ColumnSpec<'res>],
+pub struct ColumnSpecs<'slice, 'spec> {
+    specs: &'slice [ColumnSpec<'spec>],
 }
 
-impl<'res> ColumnSpecs<'res> {
-    pub(crate) fn new(specs: &'res [ColumnSpec<'res>]) -> Self {
+impl<'slice, 'spec> ColumnSpecs<'slice, 'spec> {
+    /// Creates new [`ColumnSpecs`] wrapper from a slice.
+    pub fn new(specs: &'slice [ColumnSpec<'spec>]) -> Self {
         Self { specs }
     }
 
-    pub(crate) fn inner(&self) -> &'res [ColumnSpec<'res>] {
+    /// Returns a slice of col specs encompassed by this struct.
+    pub fn as_slice(&self) -> &'slice [ColumnSpec<'spec>] {
         self.specs
     }
 
@@ -101,25 +37,24 @@ impl<'res> ColumnSpecs<'res> {
 
     /// Returns specification of k-th column returned from the database.
     #[inline]
-    pub fn get_by_index(&self, k: usize) -> Option<ColumnSpecView<'res>> {
-        self.specs.get(k).map(ColumnSpecView::new_from_column_spec)
+    pub fn get_by_index(&self, k: usize) -> Option<&'slice ColumnSpec<'spec>> {
+        self.specs.get(k)
     }
 
     /// Returns specification of the column with given name returned from the database.
     #[inline]
-    pub fn get_by_name(&self, name: &str) -> Option<(usize, ColumnSpecView<'res>)> {
+    pub fn get_by_name(&self, name: &str) -> Option<(usize, &'slice ColumnSpec<'spec>)> {
         self.specs
             .iter()
             .enumerate()
             .find(|(_idx, spec)| spec.name() == name)
-            .map(|(idx, spec)| (idx, ColumnSpecView::new_from_column_spec(spec)))
     }
 
     /// Returns iterator over specification of columns returned from the database,
     /// ordered by column order in the response.
     #[inline]
-    pub fn iter(&self) -> impl Iterator<Item = ColumnSpecView<'res>> {
-        self.specs.iter().map(ColumnSpecView::new_from_column_spec)
+    pub fn iter(&self) -> impl Iterator<Item = &'slice ColumnSpec<'spec>> {
+        self.specs.iter()
     }
 }
 
@@ -316,7 +251,7 @@ impl QueryRowsResult {
 
     /// Returns column specifications.
     #[inline]
-    pub fn column_specs(&self) -> ColumnSpecs {
+    pub fn column_specs(&self) -> ColumnSpecs<'_, '_> {
         ColumnSpecs::new(self.raw_rows_with_metadata.metadata().col_specs())
     }
 
@@ -487,8 +422,8 @@ mod tests {
     use assert_matches::assert_matches;
     use bytes::{Bytes, BytesMut};
     use itertools::Itertools as _;
-    use scylla_cql::frame::response::result::NativeType;
-    use scylla_cql::frame::response::result::ResultMetadata;
+    use scylla_cql::frame::response::result::{ColumnType, ResultMetadata};
+    use scylla_cql::frame::response::result::{NativeType, TableSpec};
     use scylla_cql::frame::types;
 
     use super::*;
@@ -589,9 +524,7 @@ mod tests {
                 // By index
                 {
                     for (i, expected_col_spec) in column_spec_infinite_iter().enumerate().take(n) {
-                        let expected_view =
-                            ColumnSpecView::new_from_column_spec(&expected_col_spec);
-                        assert_eq!(column_specs.get_by_index(i), Some(expected_view));
+                        assert_eq!(column_specs.get_by_index(i), Some(&expected_col_spec));
                     }
 
                     assert_matches!(column_specs.get_by_index(n), None);
@@ -602,9 +535,10 @@ mod tests {
                     for (idx, expected_col_spec) in column_spec_infinite_iter().enumerate().take(n)
                     {
                         let name = expected_col_spec.name();
-                        let expected_view =
-                            ColumnSpecView::new_from_column_spec(&expected_col_spec);
-                        assert_eq!(column_specs.get_by_name(name), Some((idx, expected_view)));
+                        assert_eq!(
+                            column_specs.get_by_name(name),
+                            Some((idx, &expected_col_spec))
+                        );
                     }
 
                     assert_matches!(column_specs.get_by_name("ala ma kota"), None);
@@ -615,9 +549,7 @@ mod tests {
                     for (got_view, expected_col_spec) in
                         column_specs.iter().zip(column_spec_infinite_iter())
                     {
-                        let expected_view =
-                            ColumnSpecView::new_from_column_spec(&expected_col_spec);
-                        assert_eq!(got_view, expected_view);
+                        assert_eq!(got_view, &expected_col_spec);
                     }
                 }
             }

scylla/src/statement/prepared_statement.rs

@@ -19,6 +19,7 @@ use crate::frame::response::result::PreparedMetadata;
 use crate::frame::types::{Consistency, SerialConsistency};
 use crate::observability::history::HistoryListener;
 use crate::policies::retry::RetryPolicy;
+use crate::response::query_result::ColumnSpecs;
 use crate::routing::partitioner::{Partitioner, PartitionerHasher, PartitionerName};
 use crate::routing::Token;
 
@@ -398,8 +399,8 @@ impl PreparedStatement {
     }
 
     /// Access column specifications of the bind variables of this statement
-    pub fn get_variable_col_specs(&self) -> &[ColumnSpec<'static>] {
-        &self.shared.metadata.col_specs
+    pub fn get_variable_col_specs(&self) -> ColumnSpecs<'_, 'static> {
+        ColumnSpecs::new(&self.shared.metadata.col_specs)
     }
 
     /// Access info about partition key indexes of the bind variables of this statement
@@ -413,8 +414,8 @@ impl PreparedStatement {
     }
 
     /// Access column specifications of the result set returned after the execution of this statement
-    pub fn get_result_set_col_specs(&self) -> &[ColumnSpec<'static>] {
-        self.shared.result_metadata.col_specs()
+    pub fn get_result_set_col_specs(&self) -> ColumnSpecs<'_, 'static> {
+        ColumnSpecs::new(self.shared.result_metadata.col_specs())
     }
 
     /// Get the name of the partitioner used for this statement.

scylla/tests/integration/main.rs

@@ -17,5 +17,6 @@ mod shards;
 mod silent_prepare_batch;
 mod silent_prepare_query;
 mod skip_metadata_optimization;
+mod statement;
 mod tablets;
 pub(crate) mod utils;

scylla/tests/integration/statement.rs

@@ -0,0 +1,57 @@
+use scylla::cluster::metadata::{ColumnType, NativeType};
+use scylla::frame::response::result::{ColumnSpec, TableSpec};
+
+use crate::utils::{create_new_session_builder, setup_tracing, unique_keyspace_name, PerformDDL};
+
+#[tokio::test]
+async fn test_prepared_statement_col_specs() {
+    setup_tracing();
+    let session = create_new_session_builder().build().await.unwrap();
+
+    let ks = unique_keyspace_name();
+    session
+        .ddl(format!(
+            "CREATE KEYSPACE IF NOT EXISTS {} WITH REPLICATION = 
+            {{'class' : 'NetworkTopologyStrategy', 'replication_factor' : 3}}",
+            ks
+        ))
+        .await
+        .unwrap();
+    session.use_keyspace(&ks, false).await.unwrap();
+
+    session
+        .ddl(
+            "CREATE TABLE t (k1 int, k2 varint, c1 timestamp,
+            a tinyint, b text, c smallint, PRIMARY KEY ((k1, k2), c1))",
+        )
+        .await
+        .unwrap();
+
+    let spec = |name: &'static str, typ: ColumnType<'static>| -> ColumnSpec<'_> {
+        ColumnSpec::borrowed(name, typ, TableSpec::borrowed(&ks, "t"))
+    };
+
+    let prepared = session
+        .prepare("SELECT * FROM t WHERE k1 = ? AND k2 = ? AND c1 > ?")
+        .await
+        .unwrap();
+
+    let variable_col_specs = prepared.get_variable_col_specs().as_slice();
+    let expected_variable_col_specs = &[
+        spec("k1", ColumnType::Native(NativeType::Int)),
+        spec("k2", ColumnType::Native(NativeType::Varint)),
+        spec("c1", ColumnType::Native(NativeType::Timestamp)),
+    ];
+    assert_eq!(variable_col_specs, expected_variable_col_specs);
+
+    let result_set_col_specs = prepared.get_result_set_col_specs().as_slice();
+    let expected_result_set_col_specs = &[
+        spec("k1", ColumnType::Native(NativeType::Int)),
+        spec("k2", ColumnType::Native(NativeType::Varint)),
+        spec("c1", ColumnType::Native(NativeType::Timestamp)),
+        spec("a", ColumnType::Native(NativeType::TinyInt)),
+        spec("b", ColumnType::Native(NativeType::Text)),
+        spec("c", ColumnType::Native(NativeType::SmallInt)),
+    ];
+    assert_eq!(result_set_col_specs, expected_result_set_col_specs);
+}