Merge pull request #1445 from Lorak-mmk/prepared-mutable-metadata

Mutable metadata in PreparedStatement

Author: Lorak-mmk

scylla/src/network/connection.rs

@@ -930,14 +930,14 @@ impl Connection {
 
         let cached_metadata = prepared_statement
             .get_use_cached_result_metadata()
-            .then(|| prepared_statement.get_result_metadata());
+            .then(|| prepared_statement.get_current_result_metadata());
 
         let query_response = self
             .send_request(
                 &execute_frame,
                 true,
                 prepared_statement.config.tracing,
-                cached_metadata,
+                cached_metadata.as_ref(),
             )
             .await?;
 
@@ -967,7 +967,7 @@ impl Connection {
                         &execute_frame,
                         true,
                         prepared_statement.config.tracing,
-                        cached_metadata,
+                        cached_metadata.as_ref(),
                     )
                     .await?;
 

scylla/src/statement/prepared.rs

@@ -1,6 +1,7 @@
 //! Defines the [`PreparedStatement`] type, which represents a statement
 //! that has been prepared in advance on the server.
 
+use arc_swap::{ArcSwap, Guard};
 use bytes::{Bytes, BytesMut};
 use scylla_cql::frame::response::result::{
     ColumnSpec, PartitionKeyIndex, ResultMetadata, TableSpec,
@@ -187,6 +188,7 @@ struct PreparedStatementSharedData {
     id: Bytes,
     metadata: PreparedMetadata,
     initial_result_metadata: Arc<ResultMetadata<'static>>,
+    current_result_metadata: ArcSwap<ResultMetadata<'static>>,
     statement: String,
     is_confirmed_lwt: bool,
 }
@@ -203,6 +205,18 @@ impl Clone for PreparedStatement {
     }
 }
 
+/// Stores a snapshot of current result metadata column specs.
+pub struct ColumnSpecsGuard {
+    result: Guard<Arc<ResultMetadata<'static>>>,
+}
+
+impl ColumnSpecsGuard {
+    /// Retrieves current result metadata column specs.
+    pub fn get(&self) -> ColumnSpecs<'_, 'static> {
+        ColumnSpecs::new(self.result.col_specs())
+    }
+}
+
 impl PreparedStatement {
     fn new(
         id: Bytes,
@@ -217,7 +231,8 @@ impl PreparedStatement {
             shared: Arc::new(PreparedStatementSharedData {
                 id,
                 metadata,
-                initial_result_metadata: result_metadata,
+                initial_result_metadata: Arc::clone(&result_metadata),
+                current_result_metadata: ArcSwap::from(result_metadata),
                 statement,
                 is_confirmed_lwt: is_lwt,
             }),
@@ -507,15 +522,42 @@ impl PreparedStatement {
     }
 
     /// Access metadata about the result of prepared statement returned by the database
-    pub(crate) fn get_result_metadata(&self) -> &Arc<ResultMetadata<'static>> {
-        &self.shared.initial_result_metadata
+    pub(crate) fn get_current_result_metadata(&self) -> Arc<ResultMetadata<'static>> {
+        self.shared.current_result_metadata.load_full()
     }
 
-    /// Access column specifications of the result set returned after the execution of this statement
+    /// Update metadata about the result of prepared statement.
+    // Will be used when we implement support for metadata id extension.
+    #[allow(dead_code)]
+    pub(crate) fn update_current_result_metadata(
+        &self,
+        new_metadata: Arc<ResultMetadata<'static>>,
+    ) {
+        self.shared.current_result_metadata.store(new_metadata);
+    }
+
+    /// Access column specifications of the result set returned after the preparation of this statement
+    ///
+    /// In 1.4.0, result metadata became mutable to allow us to update it when server
+    /// sends a new one (which happens in CQLv5, or CQLv4 with Scylla's metadata id extension). This method can't
+    /// be changed to support it because of Copy bound on ColumnSpecs. This method now uses metadata initially sent
+    /// by the server, which may be different than the one currently used. Please use get_current_result_set_col_specs instead."
+    // TODO(2.0): Remove this
+    #[deprecated(
+        since = "1.4.0",
+        note = "This method may return outdated metadata. Use get_current_result_set_col_specs() instead."
+    )]
     pub fn get_result_set_col_specs(&self) -> ColumnSpecs<'_, 'static> {
         ColumnSpecs::new(self.shared.initial_result_metadata.col_specs())
     }
 
+    /// Access column specifications of the result set returned after the execution of this statement
+    pub fn get_current_result_set_col_specs(&self) -> ColumnSpecsGuard {
+        ColumnSpecsGuard {
+            result: self.shared.current_result_metadata.load(),
+        }
+    }
+
     /// Get the name of the partitioner used for this statement.
     pub fn get_partitioner_name(&self) -> &PartitionerName {
         &self.partitioner_name

scylla/tests/integration/statements/prepared.rs

@@ -474,7 +474,8 @@ async fn test_prepared_statement_col_specs() {
     ];
     assert_eq!(variable_col_specs, expected_variable_col_specs);
 
-    let result_set_col_specs = prepared.get_result_set_col_specs().as_slice();
+    let col_specs_guard = prepared.get_current_result_set_col_specs();
+    let result_set_col_specs = col_specs_guard.get().as_slice();
     let expected_result_set_col_specs = &[
         spec("k1", ColumnType::Native(NativeType::Int)),
         spec("k2", ColumnType::Native(NativeType::Varint)),