cassandra_protocol/query/
batch_query_builder.rs

1use crate::consistency::Consistency;
2use crate::error::{Error as CError, Result as CResult};
3use crate::frame::message_batch::{BatchQuery, BatchQuerySubj, BatchType, BodyReqBatch};
4use crate::query::{PreparedQuery, QueryValues};
5use crate::types::{CBytesShort, CInt, CLong};
6use derivative::Derivative;
7use derive_more::Constructor;
8use std::collections::HashMap;
9
10#[derive(Debug, Clone, PartialEq, Eq, Constructor)]
11pub struct QueryBatchPreparedStatement {
12    pub query: String,
13    pub keyspace: Option<String>,
14}
15
16#[derive(Debug, Clone, PartialEq, Eq, Constructor, Derivative)]
17pub struct QueryBatch {
18    pub request: BodyReqBatch,
19    #[derivative(Debug = "ignore")]
20    pub prepared_queries: HashMap<CBytesShort, QueryBatchPreparedStatement>,
21}
22
23#[derive(Debug)]
24pub struct BatchQueryBuilder {
25    batch_type: BatchType,
26    queries: Vec<BatchQuery>,
27    prepared_queries: HashMap<CBytesShort, QueryBatchPreparedStatement>,
28    consistency: Consistency,
29    serial_consistency: Option<Consistency>,
30    timestamp: Option<CLong>,
31    keyspace: Option<String>,
32    now_in_seconds: Option<CInt>,
33}
34
35impl Default for BatchQueryBuilder {
36    fn default() -> Self {
37        BatchQueryBuilder {
38            batch_type: BatchType::Logged,
39            queries: vec![],
40            prepared_queries: HashMap::new(),
41            consistency: Consistency::One,
42            serial_consistency: None,
43            timestamp: None,
44            keyspace: None,
45            now_in_seconds: None,
46        }
47    }
48}
49
50impl BatchQueryBuilder {
51    pub fn new() -> BatchQueryBuilder {
52        Default::default()
53    }
54
55    #[must_use]
56    pub fn with_batch_type(mut self, batch_type: BatchType) -> Self {
57        self.batch_type = batch_type;
58        self
59    }
60
61    /// Add a query (non-prepared one)
62    #[must_use]
63    pub fn add_query<T: Into<String>>(mut self, query: T, values: QueryValues) -> Self {
64        self.queries.push(BatchQuery {
65            subject: BatchQuerySubj::QueryString(query.into()),
66            values,
67        });
68        self
69    }
70
71    /// Add a query (prepared one)
72    #[must_use]
73    pub fn add_query_prepared(mut self, query: &PreparedQuery, values: QueryValues) -> Self {
74        self.queries.push(BatchQuery {
75            subject: BatchQuerySubj::PreparedId(query.id.clone()),
76            values,
77        });
78        self.prepared_queries.insert(
79            query.id.clone(),
80            QueryBatchPreparedStatement::new(query.query.clone(), query.keyspace.clone()),
81        );
82
83        self
84    }
85
86    #[must_use]
87    pub fn clear_queries(mut self) -> Self {
88        self.queries = vec![];
89        self
90    }
91
92    #[must_use]
93    pub fn with_consistency(mut self, consistency: Consistency) -> Self {
94        self.consistency = consistency;
95        self
96    }
97
98    #[must_use]
99    pub fn with_serial_consistency(mut self, serial_consistency: Consistency) -> Self {
100        self.serial_consistency = Some(serial_consistency);
101        self
102    }
103
104    #[must_use]
105    pub fn with_timestamp(mut self, timestamp: CLong) -> Self {
106        self.timestamp = Some(timestamp);
107        self
108    }
109
110    #[must_use]
111    pub fn with_keyspace(mut self, keyspace: String) -> Self {
112        self.keyspace = Some(keyspace);
113        self
114    }
115
116    #[must_use]
117    pub fn with_now_in_seconds(mut self, now_in_seconds: CInt) -> Self {
118        self.now_in_seconds = Some(now_in_seconds);
119        self
120    }
121
122    pub fn build(self) -> CResult<QueryBatch> {
123        let with_names_for_values = self.queries.iter().all(|q| q.values.has_names());
124
125        if !with_names_for_values {
126            let some_names_for_values = self.queries.iter().any(|q| q.values.has_names());
127
128            if some_names_for_values {
129                return Err(CError::General(String::from(
130                    "Inconsistent query values - mixed with and without names values",
131                )));
132            }
133        }
134
135        Ok(QueryBatch::new(
136            BodyReqBatch {
137                batch_type: self.batch_type,
138                queries: self.queries,
139                consistency: self.consistency,
140                serial_consistency: self.serial_consistency,
141                timestamp: self.timestamp,
142                keyspace: self.keyspace,
143                now_in_seconds: self.now_in_seconds,
144            },
145            self.prepared_queries,
146        ))
147    }
148}