stargate_grpc/
query.rs

1//! Utilities for building queries.
2
3use crate::into_value::IntoValue;
4use crate::proto::{
5    Batch, BatchParameters, BatchQuery, Consistency, Query, QueryParameters, Value, Values,
6};
7
8impl From<Vec<Value>> for Values {
9    fn from(v: Vec<Value>) -> Self {
10        Values {
11            value_names: vec![],
12            values: v,
13        }
14    }
15}
16
17/// Builds a [`Query`].
18/// Sets the CQL string, binds values and sets query execution parameters.
19///
20/// # Example
21/// ```
22/// use stargate_grpc::{Query, Consistency};
23///
24/// let query = Query::builder()
25///     .keyspace("ks")
26///     .consistency(Consistency::LocalQuorum)
27///     .query("SELECT * FROM table WHERE year = :year and month = :month")
28///     .bind_name("year", 2021)
29///     .bind_name("month", "October")
30///     .build();
31/// ```
32///
33/// A single `QueryBuilder` can be used to create more than one query, because it is cloneable.
34/// You can set the default query parameters at the beginning, and then clone it multiple times
35/// to set a different query string or different query arguments each time:
36///
37/// # Example
38/// ```
39/// use stargate_grpc::{Query, Consistency};
40///
41/// let query_defaults = Query::builder()
42///     .keyspace("ks")
43///     .consistency(Consistency::LocalQuorum);
44///
45/// let query1 = query_defaults.clone().query("SELECT * FROM table1").build();
46/// let query2 = query_defaults.clone().query("SELECT * FROM table2").build();
47/// ```
48///
49#[derive(Default, Clone)]
50pub struct QueryBuilder {
51    cql: Option<String>,
52    values: ValuesBuilder,
53    parameters: QueryParameters,
54}
55
56impl QueryBuilder {
57    /// Creates a new `QueryBuilder` with default parameters and no query.
58    ///
59    /// You have to call [`query`](QueryBuilder::query) to set the CQL query string
60    /// before calling [`build`](QueryBuilder::build).
61    pub fn new() -> Self {
62        Default::default()
63    }
64
65    /// Sets the CQL query string.
66    pub fn query(mut self, cql: &str) -> Self {
67        self.cql = Some(cql.to_string());
68        self
69    }
70
71    /// Sets all values at once, from a vector or a value that can
72    /// be converted to a vector, e.g. a tuple.
73    ///
74    /// # Example
75    /// ```
76    /// use stargate_grpc::{Query, Value};
77    ///
78    /// let cql = "SELECT * FROM table WHERE year = ? and month = ?";
79    ///
80    /// let query1 = Query::builder()
81    ///     .query(cql)
82    ///     .bind((2021, "October"))
83    ///     .build();
84    ///
85    /// let query2 = Query::builder()
86    ///     .query(cql)
87    ///     .bind(vec![Value::bigint(2021), Value::string("October")])
88    ///     .build();
89    ///
90    /// assert_eq!(query1.values, query2.values);
91    /// ```
92    ///
93    /// # Panics
94    /// Will panic if it is called after a call to [`bind_name`](QueryBuilder::bind_name)
95    pub fn bind<I: Into<Values>>(mut self, values: I) -> Self {
96        self.values.bind(values);
97        self
98    }
99
100    /// Sets a value at a given index.
101    ///
102    /// If the internal vector of values is too small, it is automatically resized to
103    /// so that the `index` is valid, and any previously
104    /// unset values are filled with [`Value::unset`].
105    ///
106    /// # Example
107    /// ```
108    /// use stargate_grpc::Query;
109    ///
110    /// let query = Query::builder()
111    ///     .query("SELECT * FROM table WHERE year = ? and month = ?")
112    ///     .bind_ith(0, 2021)
113    ///     .bind_ith(1, "October")
114    ///     .build();
115    /// ```
116    /// # Panics
117    /// Will panic if it is called after a call to [`bind_name`](QueryBuilder::bind_name)
118    pub fn bind_ith<T: Into<Value>>(mut self, index: usize, value: T) -> Self {
119        self.values.bind_ith(index, value);
120        self
121    }
122
123    /// Binds a name to a value.
124    ///
125    /// # Example
126    /// ```
127    /// use stargate_grpc::Query;
128    ///
129    /// let query = Query::builder()
130    ///     .query("SELECT * FROM table WHERE year = :year and month = :month")
131    ///     .bind_name("year", 2021)
132    ///     .bind_name("month", "October")
133    ///     .build();
134    /// ```
135    ///
136    /// # Panics
137    /// Will panic if mixed with calls to [`bind`](QueryBuilder::bind)
138    /// or [`bind_ith`](QueryBuilder::bind_ith).
139    pub fn bind_name<T: Into<Value>>(mut self, name: &str, value: T) -> Self {
140        self.values.bind_name(name, value);
141        self
142    }
143
144    /// Sets the keyspace the query will apply to.
145    ///
146    /// See [`QueryParameters::keyspace`].
147    pub fn keyspace(mut self, keyspace: &str) -> Self {
148        self.parameters.keyspace = Some(keyspace.to_string());
149        self
150    }
151
152    /// Sets the consistency level of the query.
153    /// # Example
154    /// ```
155    /// use stargate_grpc::{Consistency, Query};
156    ///
157    /// let query = Query::builder()
158    ///     .query("SELECT * FROM table")
159    ///     .consistency(Consistency::One);
160    /// ```
161    /// See [`QueryParameters::consistency`].
162    pub fn consistency(mut self, consistency: Consistency) -> Self {
163        self.parameters.consistency = Some(crate::proto::ConsistencyValue {
164            value: consistency.into(),
165        });
166        self
167    }
168
169    /// Sets the serial consistency level (if the query is a lightweight transaction).
170    ///
171    /// See [`QueryParameters::serial_consistency`].
172    pub fn serial_consistency(mut self, consistency: Consistency) -> Self {
173        self.parameters.serial_consistency = Some(crate::proto::ConsistencyValue {
174            value: consistency.into(),
175        });
176        self
177    }
178
179    /// Sets the maximum number of rows that will be returned in the response.
180    ///
181    /// See [`QueryParameters::page_size`].
182    pub fn page_size(mut self, size: i32) -> Self {
183        self.parameters.page_size = Some(size);
184        self
185    }
186
187    /// Sets a paging state that indicates where to resume iteration in the result set.
188    ///
189    /// See [`QueryParameters::paging_state`].
190    pub fn paging_state(mut self, paging_state: Vec<u8>) -> Self {
191        self.parameters.paging_state = Some(paging_state);
192        self
193    }
194
195    /// Sets whether the server should collect tracing information about the execution of the query.
196    ///
197    /// See [`QueryParameters::tracing`].
198    pub fn tracing(mut self, tracing: bool) -> Self {
199        self.parameters.tracing = tracing;
200        self
201    }
202
203    /// Sets the query timestamp (in microseconds).
204    ///
205    /// See [`QueryParameters::timestamp`].
206    pub fn timestamp(mut self, timestamp: i64) -> Self {
207        self.parameters.timestamp = Some(timestamp);
208        self
209    }
210
211    /// Sets all parameters of the query at once.
212    ///
213    /// Overwrites any parameters that were set before.
214    pub fn parameters(self, parameters: QueryParameters) -> Self {
215        QueryBuilder { parameters, ..self }
216    }
217
218    /// Builds the query that can be passed to
219    /// [`StargateClient::execute_query`](crate::StargateClient::execute_query).
220    ///
221    /// # Panics
222    /// Will panic if the query string was not set.
223    pub fn build(mut self) -> Query {
224        Query {
225            cql: self.cql.expect("cql string"),
226            values: self.values.build(),
227            parameters: Some(self.parameters),
228        }
229    }
230}
231
232impl Query {
233    /// Returns a fresh builder for building a query
234    pub fn builder() -> QueryBuilder {
235        QueryBuilder::new()
236    }
237}
238
239/// Builds a batch of queries.
240///
241/// # Example
242/// ```
243/// use stargate_grpc::{Batch, Consistency};
244///
245/// let batch = Batch::builder()
246///     .keyspace("example")
247///     .consistency(Consistency::LocalQuorum)
248///     .query("INSERT INTO users (id, login, email) VALUES (?, ?, ?)")
249///     .bind((0, "admin", "admin@example.net"))
250///     .query("INSERT INTO users_by_login (id, login) VALUES (?, ?)")
251///     .bind((0, "admin"))
252///     .build();
253/// ```
254#[derive(Default, Clone)]
255pub struct BatchBuilder {
256    cql: Option<String>,
257    values: ValuesBuilder,
258    parameters: BatchParameters,
259    built_queries: Vec<BatchQuery>,
260}
261
262impl BatchBuilder {
263    /// Creates a new `BatchBuilder` with no queries in it.
264    pub fn new() -> BatchBuilder {
265        Default::default()
266    }
267
268    /// Adds a CQL query to the batch.
269    ///
270    /// If the query has arguments, set their values with
271    /// one of the `bind` functions.
272    pub fn query(mut self, cql: &str) -> Self {
273        self.finalize_query();
274        self.cql = Some(cql.to_string());
275        self
276    }
277
278    /// Binds all arguments of the lately added query at once,
279    /// from a vector or a value that can be converted to a vector, e.g. a tuple.
280    ///
281    /// # Panics
282    /// Will panic if it is called after a call to [`bind_name`](BatchBuilder::bind_name)
283    pub fn bind<I: Into<Values>>(mut self, values: I) -> Self {
284        self.values.bind(values);
285        self
286    }
287
288    /// Binds an argument of the recently added query at a given index.
289    ///
290    /// This function can be called multiple times, to bind several arguments.
291    /// If the internal vector of values is too small, it is automatically resized to
292    /// so that the `index` is valid, and any previously
293    /// unset values are filled with [`Value::unset`].
294    pub fn bind_ith<T: Into<Value>>(mut self, index: usize, value: T) -> Self {
295        self.values.bind_ith(index, value);
296        self
297    }
298
299    /// Binds a name to a value.
300    ///
301    /// This function can be called multiple times, to bind several arguments.
302    ///
303    /// # Panics
304    /// Will panic if mixed with calls to [`bind`](BatchBuilder::bind)
305    /// or [`bind_ith`](BatchBuilder::bind_ith).
306    pub fn bind_name<T: Into<Value>>(mut self, name: &str, value: T) -> Self {
307        self.values.bind_name(name, value);
308        self
309    }
310
311    /// Sets the keyspace every query in the batch will apply to.
312    ///
313    /// See [`BatchParameters::keyspace`].
314    pub fn keyspace(mut self, keyspace: &str) -> Self {
315        self.parameters.keyspace = Some(keyspace.to_string());
316        self
317    }
318
319    /// Sets the consistency level of all queries in the batch.
320    ///
321    /// See [`BatchParameters::consistency`].
322    pub fn consistency(mut self, consistency: Consistency) -> Self {
323        self.parameters.consistency = Some(crate::proto::ConsistencyValue {
324            value: consistency.into(),
325        });
326        self
327    }
328
329    /// Sets whether the server should collect tracing information about the execution of the batch.
330    ///
331    /// See [`BatchParameters::tracing`].
332    pub fn tracing(mut self, tracing: bool) -> Self {
333        self.parameters.tracing = tracing;
334        self
335    }
336
337    /// Sets the query timestamp (in microseconds).
338    ///
339    /// See [`BatchParameters::timestamp`].
340    pub fn timestamp(mut self, timestamp: i64) -> Self {
341        self.parameters.timestamp = Some(timestamp);
342        self
343    }
344
345    /// Sets the serial consistency level (if the query is a lightweight transaction).
346    ///
347    /// See [`BatchParameters::serial_consistency`].
348    pub fn serial_consistency(mut self, consistency: Consistency) -> Self {
349        self.parameters.serial_consistency = Some(crate::proto::ConsistencyValue {
350            value: consistency.into(),
351        });
352        self
353    }
354
355    /// Sets all parameters of the batch at once.
356    ///
357    /// Overwrites any parameters that were set before.
358    pub fn parameters(mut self, parameters: BatchParameters) -> Self {
359        self.parameters = parameters;
360        self
361    }
362
363    /// Finalizes building and returns the `Batch` that can be passed to
364    /// [`StargateClient::execute_batch`](crate::StargateClient::execute_batch).
365    pub fn build(mut self) -> Batch {
366        self.finalize_query();
367        Batch {
368            r#type: 0,
369            queries: self.built_queries,
370            parameters: Some(self.parameters),
371        }
372    }
373
374    fn finalize_query(&mut self) {
375        if let Some(cql) = self.cql.take() {
376            self.built_queries.push(BatchQuery {
377                cql,
378                values: self.values.build(),
379            });
380        }
381    }
382}
383
384impl Batch {
385    /// Returns a fresh builder for building a batch of queries
386    pub fn builder() -> BatchBuilder {
387        BatchBuilder::new()
388    }
389}
390
391/// The logic of building the query argument values,
392/// shared between [`QueryBuilder`] and [`BatchBuilder`]
393#[derive(Default, Clone)]
394struct ValuesBuilder {
395    values: Vec<Value>,
396    value_names: Vec<String>,
397}
398
399impl ValuesBuilder {
400    pub fn bind<I: Into<Values>>(&mut self, values: I) {
401        if !self.value_names.is_empty() {
402            panic!("Mixing named with non-named values is not allowed")
403        }
404        let values = values.into();
405        self.values.extend(values.values);
406        self.value_names.extend(values.value_names)
407    }
408
409    pub fn bind_ith<T: Into<Value>>(&mut self, index: usize, value: T) {
410        if !self.value_names.is_empty() {
411            panic!("Mixing named with non-named values is not allowed")
412        }
413        if index >= self.values.len() {
414            self.values.resize(index + 1, Value::unset());
415        }
416        self.values[index] = value.into_value();
417    }
418
419    pub fn bind_name<T: Into<Value>>(&mut self, name: &str, value: T) {
420        if self.values.len() != self.value_names.len() {
421            panic!("Mixing named with non-named values is not allowed")
422        }
423        self.value_names.push(name.to_string());
424        self.values.push(value.into_value());
425    }
426
427    /// If there were any values bound with one of the `bind_` calls, moves them to the
428    /// return `Values` object. If no values were bound, returns `None`.
429    /// After returning, `self` is left in a clean, empty state (the value vectors are cleared).
430    /// This method can be called multiple times.
431    pub fn build(&mut self) -> Option<Values> {
432        if self.values.is_empty() {
433            None
434        } else {
435            Some(Values {
436                values: self.values.drain(0..).collect(),
437                value_names: self.value_names.drain(0..).collect(),
438            })
439        }
440    }
441}
442
443#[cfg(test)]
444mod test {
445    use crate::proto::Values;
446    use crate::query::ValuesBuilder;
447    use crate::Value;
448
449    #[test]
450    fn bind_a_single_item_tuple() {
451        let mut builder = ValuesBuilder::default();
452        builder.bind((1,));
453        let values = builder.build();
454        assert_eq!(
455            values,
456            Some(Values {
457                values: vec![Value::int(1)],
458                value_names: vec![]
459            })
460        )
461    }
462
463    #[test]
464    fn bind_a_tuple() {
465        let mut builder = ValuesBuilder::default();
466        builder.bind((1, 2.0, "foo"));
467        let values = builder.build();
468        assert_eq!(
469            values,
470            Some(Values {
471                values: vec![Value::int(1), Value::double(2.0), Value::string("foo")],
472                value_names: vec![]
473            })
474        )
475    }
476
477    #[test]
478    fn bind_ith() {
479        let mut builder = ValuesBuilder::default();
480        builder.bind_ith(0, 1);
481        builder.bind_ith(1, 2.0);
482        builder.bind_ith(2, "foo");
483        let values = builder.build();
484        assert_eq!(
485            values,
486            Some(Values {
487                values: vec![Value::int(1), Value::double(2.0), Value::string("foo")],
488                value_names: vec![]
489            })
490        )
491    }
492
493    #[test]
494    fn bind_name() {
495        let mut builder = ValuesBuilder::default();
496        builder.bind_name("a", 1);
497        builder.bind_name("b", 2.0);
498        builder.bind_name("c", "foo");
499        let values = builder.build();
500        assert_eq!(
501            values,
502            Some(Values {
503                values: vec![Value::int(1), Value::double(2.0), Value::string("foo")],
504                value_names: vec!["a".to_string(), "b".to_string(), "c".to_string()]
505            })
506        );
507    }
508}