From 486cc1209e4e02625a9c3fe75dee3491015d21e3 Mon Sep 17 00:00:00 2001
From: muzarski <mikolaj.uzarski@scylladb.com>
Date: Fri, 5 Jan 2024 13:16:09 +0100
Subject: [PATCH] prepared: add result_metadata field

Added a `result_metadata` to `PreparedStatementSharedData` struct so
the driver caches the result metadata provided by the server after
statement preparation.
---
 scylla-cql/src/frame/response/result.rs    |  2 +-
 scylla/src/statement/prepared_statement.rs | 15 ++++++++++++++-
 scylla/src/transport/caching_session.rs    |  5 ++++-
 scylla/src/transport/connection.rs         |  1 +
 4 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/scylla-cql/src/frame/response/result.rs b/scylla-cql/src/frame/response/result.rs
index 745b925c00..ab5a588fba 100644
--- a/scylla-cql/src/frame/response/result.rs
+++ b/scylla-cql/src/frame/response/result.rs
@@ -386,7 +386,7 @@ pub struct ColumnSpec {
     pub typ: ColumnType,
 }
 
-#[derive(Debug, Default)]
+#[derive(Debug, Clone, Default)]
 pub struct ResultMetadata {
     col_count: usize,
     pub paging_state: Option<Bytes>,
diff --git a/scylla/src/statement/prepared_statement.rs b/scylla/src/statement/prepared_statement.rs
index b9e4894872..3d56e43d1c 100644
--- a/scylla/src/statement/prepared_statement.rs
+++ b/scylla/src/statement/prepared_statement.rs
@@ -10,7 +10,7 @@ use std::time::Duration;
 use thiserror::Error;
 use uuid::Uuid;
 
-use scylla_cql::frame::response::result::{ColumnSpec, PartitionKeyIndex};
+use scylla_cql::frame::response::result::{ColumnSpec, PartitionKeyIndex, ResultMetadata};
 
 use super::StatementConfig;
 use crate::frame::response::result::PreparedMetadata;
@@ -37,6 +37,7 @@ pub struct PreparedStatement {
 #[derive(Debug)]
 struct PreparedStatementSharedData {
     metadata: PreparedMetadata,
+    result_metadata: ResultMetadata,
     statement: String,
 }
 
@@ -59,6 +60,7 @@ impl PreparedStatement {
         id: Bytes,
         is_lwt: bool,
         metadata: PreparedMetadata,
+        result_metadata: ResultMetadata,
         statement: String,
         page_size: Option<i32>,
         config: StatementConfig,
@@ -67,6 +69,7 @@ impl PreparedStatement {
             id,
             shared: Arc::new(PreparedStatementSharedData {
                 metadata,
+                result_metadata,
                 statement,
             }),
             prepare_tracing_ids: Vec::new(),
@@ -335,6 +338,16 @@ impl PreparedStatement {
         &self.shared.metadata.pk_indexes
     }
 
+    /// Access metadata about the result of prepared statement returned by the database
+    pub(crate) fn get_result_metadata(&self) -> &ResultMetadata {
+        &self.shared.result_metadata
+    }
+
+    /// Access column specifications of the result set returned after the execution of this statement
+    pub fn get_result_set_col_specs(&self) -> &[ColumnSpec] {
+        &self.shared.result_metadata.col_specs
+    }
+
     /// Get the name of the partitioner used for this statement.
     pub(crate) fn get_partitioner_name(&self) -> &PartitionerName {
         &self.partitioner_name
diff --git a/scylla/src/transport/caching_session.rs b/scylla/src/transport/caching_session.rs
index 9af2d08dbe..0bdc2ad138 100644
--- a/scylla/src/transport/caching_session.rs
+++ b/scylla/src/transport/caching_session.rs
@@ -8,7 +8,7 @@ use crate::{QueryResult, Session};
 use bytes::Bytes;
 use dashmap::DashMap;
 use futures::future::try_join_all;
-use scylla_cql::frame::response::result::PreparedMetadata;
+use scylla_cql::frame::response::result::{PreparedMetadata, ResultMetadata};
 use scylla_cql::types::serialize::batch::BatchValues;
 use scylla_cql::types::serialize::row::SerializeRow;
 use std::collections::hash_map::RandomState;
@@ -23,6 +23,7 @@ struct RawPreparedStatementData {
     id: Bytes,
     is_confirmed_lwt: bool,
     metadata: PreparedMetadata,
+    result_metadata: ResultMetadata,
     partitioner_name: PartitionerName,
 }
 
@@ -168,6 +169,7 @@ where
                 raw.id.clone(),
                 raw.is_confirmed_lwt,
                 raw.metadata.clone(),
+                raw.result_metadata.clone(),
                 query.contents,
                 page_size,
                 query.config,
@@ -195,6 +197,7 @@ where
                 id: prepared.get_id().clone(),
                 is_confirmed_lwt: prepared.is_confirmed_lwt(),
                 metadata: prepared.get_variable_metadata().clone(),
+                result_metadata: prepared.get_result_metadata().clone(),
                 partitioner_name: prepared.get_partitioner_name().clone(),
             };
             self.cache.insert(query_contents, raw);
diff --git a/scylla/src/transport/connection.rs b/scylla/src/transport/connection.rs
index ff4aa4da08..edaa1aad8e 100644
--- a/scylla/src/transport/connection.rs
+++ b/scylla/src/transport/connection.rs
@@ -552,6 +552,7 @@ impl Connection {
                     .protocol_features
                     .prepared_flags_contain_lwt_mark(p.prepared_metadata.flags as u32),
                 p.prepared_metadata,
+                p.result_metadata,
                 query.contents.clone(),
                 query.get_page_size(),
                 query.config.clone(),