mentat/
query_builder.rs

1// Copyright 2018 Mozilla
2//
3// Licensed under the Apache License, Version 2.0 (the "License"); you may not use
4// this file except in compliance with the License. You may obtain a copy of the
5// License at http://www.apache.org/licenses/LICENSE-2.0
6// Unless required by applicable law or agreed to in writing, software distributed
7// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
8// CONDITIONS OF ANY KIND, either express or implied. See the License for the
9// specific language governing permissions and limitations under the License.
10
11#![macro_use]
12use std::collections::BTreeMap;
13
14pub use core_traits::{Binding, Entid, TypedValue, ValueType};
15
16use mentat_core::{DateTime, Keyword, Utc};
17
18use super::{HasSchema, QueryInputs, QueryOutput, Queryable, RelResult, Store, Variable};
19
20use public_traits::errors::{MentatError, Result};
21
22pub struct QueryBuilder<'a> {
23    query: String,
24    values: BTreeMap<Variable, TypedValue>,
25    types: BTreeMap<Variable, ValueType>,
26    store: &'a mut Store,
27}
28
29impl<'a> QueryBuilder<'a> {
30    pub fn new<T>(store: &mut Store, query: T) -> QueryBuilder<'_>
31    where
32        T: Into<String>,
33    {
34        QueryBuilder {
35            query: query.into(),
36            values: BTreeMap::new(),
37            types: BTreeMap::new(),
38            store,
39        }
40    }
41
42    pub fn bind_value<T>(&mut self, var: &str, value: T) -> &mut Self
43    where
44        T: Into<TypedValue>,
45    {
46        self.values
47            .insert(Variable::from_valid_name(var), value.into());
48        self
49    }
50
51    pub fn bind_ref_from_kw(&mut self, var: &str, value: Keyword) -> Result<&mut Self> {
52        let entid = self
53            .store
54            .conn()
55            .current_schema()
56            .get_entid(&value)
57            .ok_or_else(|| MentatError::UnknownAttribute(value.to_string()))?;
58        self.values.insert(
59            Variable::from_valid_name(var),
60            TypedValue::Ref(entid.into()),
61        );
62        Ok(self)
63    }
64
65    pub fn bind_ref<T>(&mut self, var: &str, value: T) -> &mut Self
66    where
67        T: Into<Entid>,
68    {
69        self.values.insert(
70            Variable::from_valid_name(var),
71            TypedValue::Ref(value.into()),
72        );
73        self
74    }
75
76    pub fn bind_long(&mut self, var: &str, value: i64) -> &mut Self {
77        self.values
78            .insert(Variable::from_valid_name(var), TypedValue::Long(value));
79        self
80    }
81
82    pub fn bind_instant(&mut self, var: &str, value: i64) -> &mut Self {
83        self.values
84            .insert(Variable::from_valid_name(var), TypedValue::instant(value));
85
86        self
87    }
88
89    pub fn bind_date_time(&mut self, var: &str, value: DateTime<Utc>) -> &mut Self {
90        self.values
91            .insert(Variable::from_valid_name(var), TypedValue::Instant(value));
92        self
93    }
94
95    pub fn bind_type(&mut self, var: &str, value_type: ValueType) -> &mut Self {
96        self.types
97            .insert(Variable::from_valid_name(var), value_type);
98        self
99    }
100
101    pub fn execute(&mut self) -> Result<QueryOutput> {
102        let values = ::std::mem::take(&mut self.values);
103        let types = ::std::mem::take(&mut self.types);
104        let query_inputs = QueryInputs::new(types, values)?;
105        let read = self.store.begin_read()?;
106        read.q_once(&self.query, query_inputs)
107    }
108
109    pub fn execute_scalar(&mut self) -> Result<Option<Binding>> {
110        let results = self.execute()?;
111        results.into_scalar().map_err(|e| e.into())
112    }
113
114    pub fn execute_coll(&mut self) -> Result<Vec<Binding>> {
115        let results = self.execute()?;
116        results.into_coll().map_err(|e| e.into())
117    }
118
119    pub fn execute_tuple(&mut self) -> Result<Option<Vec<Binding>>> {
120        let results = self.execute()?;
121        results.into_tuple().map_err(|e| e.into())
122    }
123
124    pub fn execute_rel(&mut self) -> Result<RelResult<Binding>> {
125        let results = self.execute()?;
126        results.into_rel().map_err(|e| e.into())
127    }
128}
129
130#[cfg(test)]
131mod test {
132    use super::{QueryBuilder, Store, TypedValue};
133
134    #[test]
135    fn test_scalar_query() {
136        let mut store = Store::open("").expect("store connection");
137        store
138            .transact(
139                r#"[
140            [:db/add "s" :db/ident :foo/boolean]
141            [:db/add "s" :db/valueType :db.type/boolean]
142            [:db/add "s" :db/cardinality :db.cardinality/one]
143        ]"#,
144            )
145            .expect("successful transaction");
146
147        let report = store
148            .transact(
149                r#"[
150            [:db/add "u" :foo/boolean true]
151            [:db/add "p" :foo/boolean false]
152        ]"#,
153            )
154            .expect("successful transaction");
155
156        let yes = *report.tempids.get("u").expect("found it");
157
158        let entid = QueryBuilder::new(
159            &mut store,
160            r#"[:find ?x .
161                                                      :in ?v
162                                                      :where [?x :foo/boolean ?v]]"#,
163        )
164        .bind_value("?v", true)
165        .execute_scalar()
166        .expect("ScalarResult")
167        .and_then(|t| t.into_entid());
168        assert_eq!(entid, Some(yes));
169    }
170
171    #[test]
172    fn test_coll_query() {
173        let mut store = Store::open("").expect("store connection");
174        store
175            .transact(
176                r#"[
177            [:db/add "s" :db/ident :foo/boolean]
178            [:db/add "s" :db/valueType :db.type/boolean]
179            [:db/add "s" :db/cardinality :db.cardinality/one]
180            [:db/add "t" :db/ident :foo/long]
181            [:db/add "t" :db/valueType :db.type/long]
182            [:db/add "t" :db/cardinality :db.cardinality/one]
183        ]"#,
184            )
185            .expect("successful transaction");
186
187        let report = store
188            .transact(
189                r#"[
190            [:db/add "l" :foo/boolean true]
191            [:db/add "l" :foo/long 25]
192            [:db/add "m" :foo/boolean false]
193            [:db/add "m" :foo/long 26]
194            [:db/add "n" :foo/boolean true]
195            [:db/add "n" :foo/long 27]
196            [:db/add "p" :foo/boolean false]
197            [:db/add "p" :foo/long 24]
198            [:db/add "u" :foo/boolean true]
199            [:db/add "u" :foo/long 23]
200        ]"#,
201            )
202            .expect("successful transaction");
203
204        let u_yes = *report.tempids.get("u").expect("found it");
205        let l_yes = *report.tempids.get("l").expect("found it");
206        let n_yes = *report.tempids.get("n").expect("found it");
207
208        let entids: Vec<i64> = QueryBuilder::new(
209            &mut store,
210            r#"[:find [?x ...]
211                                                                 :in ?v
212                                                                 :where [?x :foo/boolean ?v]]"#,
213        )
214        .bind_value("?v", true)
215        .execute_coll()
216        .expect("CollResult")
217        .into_iter()
218        .map(|v| v.into_entid().expect("val"))
219        .collect();
220
221        assert_eq!(entids, vec![l_yes, n_yes, u_yes]);
222    }
223
224    #[test]
225    fn test_coll_query_by_row() {
226        let mut store = Store::open("").expect("store connection");
227        store
228            .transact(
229                r#"[
230            [:db/add "s" :db/ident :foo/boolean]
231            [:db/add "s" :db/valueType :db.type/boolean]
232            [:db/add "s" :db/cardinality :db.cardinality/one]
233            [:db/add "t" :db/ident :foo/long]
234            [:db/add "t" :db/valueType :db.type/long]
235            [:db/add "t" :db/cardinality :db.cardinality/one]
236        ]"#,
237            )
238            .expect("successful transaction");
239
240        let report = store
241            .transact(
242                r#"[
243            [:db/add "l" :foo/boolean true]
244            [:db/add "l" :foo/long 25]
245            [:db/add "m" :foo/boolean false]
246            [:db/add "m" :foo/long 26]
247            [:db/add "n" :foo/boolean true]
248            [:db/add "n" :foo/long 27]
249            [:db/add "p" :foo/boolean false]
250            [:db/add "p" :foo/long 24]
251            [:db/add "u" :foo/boolean true]
252            [:db/add "u" :foo/long 23]
253        ]"#,
254            )
255            .expect("successful transaction");
256
257        let n_yes = *report.tempids.get("n").expect("found it");
258
259        let results = QueryBuilder::new(
260            &mut store,
261            r#"[:find [?x ...]
262                                                        :in ?v
263                                                        :where [?x :foo/boolean ?v]]"#,
264        )
265        .bind_value("?v", true)
266        .execute_coll()
267        .expect("CollResult");
268        let entid = results
269            .get(1)
270            .and_then(|t| t.to_owned().into_entid())
271            .expect("entid");
272
273        assert_eq!(entid, n_yes);
274    }
275
276    #[test]
277    fn test_tuple_query_result_by_column() {
278        let mut store = Store::open("").expect("store connection");
279        store
280            .transact(
281                r#"[
282            [:db/add "s" :db/ident :foo/boolean]
283            [:db/add "s" :db/valueType :db.type/boolean]
284            [:db/add "s" :db/cardinality :db.cardinality/one]
285            [:db/add "t" :db/ident :foo/long]
286            [:db/add "t" :db/valueType :db.type/long]
287            [:db/add "t" :db/cardinality :db.cardinality/one]
288        ]"#,
289            )
290            .expect("successful transaction");
291
292        let report = store
293            .transact(
294                r#"[
295            [:db/add "l" :foo/boolean true]
296            [:db/add "l" :foo/long 25]
297            [:db/add "m" :foo/boolean false]
298            [:db/add "m" :foo/long 26]
299            [:db/add "n" :foo/boolean true]
300            [:db/add "n" :foo/long 27]
301            [:db/add "p" :foo/boolean false]
302            [:db/add "p" :foo/long 24]
303            [:db/add "u" :foo/boolean true]
304            [:db/add "u" :foo/long 23]
305        ]"#,
306            )
307            .expect("successful transaction");
308
309        let n_yes = *report.tempids.get("n").expect("found it");
310
311        let results = QueryBuilder::new(
312            &mut store,
313            r#"[:find [?x, ?i]
314                                                        :in ?v ?i
315                                                        :where [?x :foo/boolean ?v]
316                                                               [?x :foo/long ?i]]"#,
317        )
318        .bind_value("?v", true)
319        .bind_long("?i", 27)
320        .execute_tuple()
321        .expect("TupleResult")
322        .expect("Vec<TypedValue>");
323        let entid = results
324            .first()
325            .and_then(|t| t.to_owned().into_entid())
326            .expect("entid");
327        let long_val = results
328            .get(1)
329            .and_then(|t| t.to_owned().into_long())
330            .expect("long");
331
332        assert_eq!(entid, n_yes);
333        assert_eq!(long_val, 27);
334    }
335
336    #[test]
337    fn test_tuple_query_result_by_iter() {
338        let mut store = Store::open("").expect("store connection");
339        store
340            .transact(
341                r#"[
342            [:db/add "s" :db/ident :foo/boolean]
343            [:db/add "s" :db/valueType :db.type/boolean]
344            [:db/add "s" :db/cardinality :db.cardinality/one]
345            [:db/add "t" :db/ident :foo/long]
346            [:db/add "t" :db/valueType :db.type/long]
347            [:db/add "t" :db/cardinality :db.cardinality/one]
348        ]"#,
349            )
350            .expect("successful transaction");
351
352        let report = store
353            .transact(
354                r#"[
355            [:db/add "l" :foo/boolean true]
356            [:db/add "l" :foo/long 25]
357            [:db/add "m" :foo/boolean false]
358            [:db/add "m" :foo/long 26]
359            [:db/add "n" :foo/boolean true]
360            [:db/add "n" :foo/long 27]
361            [:db/add "p" :foo/boolean false]
362            [:db/add "p" :foo/long 24]
363            [:db/add "u" :foo/boolean true]
364            [:db/add "u" :foo/long 23]
365        ]"#,
366            )
367            .expect("successful transaction");
368
369        let n_yes = *report.tempids.get("n").expect("found it");
370
371        let results: Vec<_> = QueryBuilder::new(
372            &mut store,
373            r#"[:find [?x, ?i]
374                                                                :in ?v ?i
375                                                                :where [?x :foo/boolean ?v]
376                                                                       [?x :foo/long ?i]]"#,
377        )
378        .bind_value("?v", true)
379        .bind_long("?i", 27)
380        .execute_tuple()
381        .expect("TupleResult")
382        .unwrap_or_default();
383        let entid = TypedValue::Ref(n_yes).into();
384        let long_val = TypedValue::Long(27).into();
385
386        assert_eq!(results, vec![entid, long_val]);
387    }
388
389    #[test]
390    fn test_rel_query_result() {
391        let mut store = Store::open("").expect("store connection");
392        store
393            .transact(
394                r#"[
395            [:db/add "s" :db/ident :foo/boolean]
396            [:db/add "s" :db/valueType :db.type/boolean]
397            [:db/add "s" :db/cardinality :db.cardinality/one]
398            [:db/add "t" :db/ident :foo/long]
399            [:db/add "t" :db/valueType :db.type/long]
400            [:db/add "t" :db/cardinality :db.cardinality/one]
401        ]"#,
402            )
403            .expect("successful transaction");
404
405        let report = store
406            .transact(
407                r#"[
408            [:db/add "l" :foo/boolean true]
409            [:db/add "l" :foo/long 25]
410            [:db/add "m" :foo/boolean false]
411            [:db/add "m" :foo/long 26]
412            [:db/add "n" :foo/boolean true]
413            [:db/add "n" :foo/long 27]
414        ]"#,
415            )
416            .expect("successful transaction");
417
418        let l_yes = *report.tempids.get("l").expect("found it");
419        let m_yes = *report.tempids.get("m").expect("found it");
420        let n_yes = *report.tempids.get("n").expect("found it");
421
422        #[derive(Debug, PartialEq)]
423        struct Res {
424            entid: i64,
425            boolean: bool,
426            long_val: i64,
427        }
428
429        let mut results: Vec<Res> = QueryBuilder::new(
430            &mut store,
431            r#"[:find ?x ?v ?i
432                                                                      :where [?x :foo/boolean ?v]
433                                                                             [?x :foo/long ?i]]"#,
434        )
435        .execute_rel()
436        .expect("RelResult")
437        .into_iter()
438        .map(|row| Res {
439            entid: row
440                .first()
441                .and_then(|t| t.to_owned().into_entid())
442                .expect("entid"),
443            boolean: row
444                .get(1)
445                .and_then(|t| t.to_owned().into_boolean())
446                .expect("boolean"),
447            long_val: row
448                .get(2)
449                .and_then(|t| t.to_owned().into_long())
450                .expect("long"),
451        })
452        .collect();
453
454        let res1 = results.pop().expect("res");
455        assert_eq!(
456            res1,
457            Res {
458                entid: n_yes,
459                boolean: true,
460                long_val: 27
461            }
462        );
463        let res2 = results.pop().expect("res");
464        assert_eq!(
465            res2,
466            Res {
467                entid: m_yes,
468                boolean: false,
469                long_val: 26
470            }
471        );
472        let res3 = results.pop().expect("res");
473        assert_eq!(
474            res3,
475            Res {
476                entid: l_yes,
477                boolean: true,
478                long_val: 25
479            }
480        );
481        assert_eq!(results.pop(), None);
482    }
483
484    #[test]
485    fn test_bind_ref() {
486        let mut store = Store::open("").expect("store connection");
487        store
488            .transact(
489                r#"[
490            [:db/add "s" :db/ident :foo/boolean]
491            [:db/add "s" :db/valueType :db.type/boolean]
492            [:db/add "s" :db/cardinality :db.cardinality/one]
493            [:db/add "t" :db/ident :foo/long]
494            [:db/add "t" :db/valueType :db.type/long]
495            [:db/add "t" :db/cardinality :db.cardinality/one]
496        ]"#,
497            )
498            .expect("successful transaction");
499
500        let report = store
501            .transact(
502                r#"[
503            [:db/add "l" :foo/boolean true]
504            [:db/add "l" :foo/long 25]
505            [:db/add "m" :foo/boolean false]
506            [:db/add "m" :foo/long 26]
507            [:db/add "n" :foo/boolean true]
508            [:db/add "n" :foo/long 27]
509        ]"#,
510            )
511            .expect("successful transaction");
512
513        let l_yes = *report.tempids.get("l").expect("found it");
514
515        let results = QueryBuilder::new(
516            &mut store,
517            r#"[:find [?v ?i]
518                                                        :in ?x
519                                                        :where [?x :foo/boolean ?v]
520                                                               [?x :foo/long ?i]]"#,
521        )
522        .bind_ref("?x", l_yes)
523        .execute_tuple()
524        .expect("TupleResult")
525        .unwrap_or_default();
526        assert!(results
527            .first()
528            .and_then(|t| t.to_owned().into_boolean())
529            .expect("boolean"));
530        assert_eq!(
531            results
532                .get(1)
533                .and_then(|t| t.to_owned().into_long())
534                .expect("long"),
535            25
536        );
537    }
538}
mentat/query_builder.rs

mentat/
query_builder.rs
