typedb_driver/
given.rs

1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements.  See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership.  The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License.  You may obtain a copy of the License at
9 *
10 *   http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied.  See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19use std::{collections::HashMap, sync::Arc};
20
21use chrono::{DateTime, NaiveDate, NaiveDateTime};
22
23use crate::{
24    Result,
25    concept::{
26        Attribute, Concept, Entity, Relation, Value,
27        value::{Decimal, Duration, TimeZone},
28    },
29    error::{Error, QueryError},
30};
31
32#[derive(Debug, Clone)]
33pub enum GivenRowEntry {
34    Empty,
35    Entity(Entity),
36    Relation(Relation),
37    Attribute(Attribute),
38    Value(Value),
39}
40
41macro_rules! impl_from_for_given_row_entry {
42    ($($t:ty),*) => {
43        $(impl From<$t> for GivenRowEntry {
44            fn from(value: $t) -> Self { Self::Value(Value::from(value)) }
45        })*
46    };
47}
48
49impl_from_for_given_row_entry!(bool, i64, f64, Decimal, String, NaiveDate, NaiveDateTime, DateTime<TimeZone>, Duration);
50
51impl From<&str> for GivenRowEntry {
52    fn from(value: &str) -> Self {
53        Self::Value(Value::from(value))
54    }
55}
56
57impl TryFrom<Concept> for GivenRowEntry {
58    type Error = Error;
59
60    fn try_from(concept: Concept) -> Result<Self> {
61        match concept {
62            Concept::Entity(e) => Ok(Self::Entity(e)),
63            Concept::Relation(r) => Ok(Self::Relation(r)),
64            Concept::Attribute(a) => Ok(Self::Attribute(a)),
65            Concept::Value(v) => Ok(Self::Value(v)),
66
67            Concept::EntityType(_) | Concept::RelationType(_) | Concept::RoleType(_) | Concept::AttributeType(_) => {
68                Err(Error::Query(QueryError::InvalidTypeToGivenRow))
69            }
70        }
71    }
72}
73
74impl From<Entity> for GivenRowEntry {
75    fn from(value: Entity) -> Self {
76        Self::Entity(value)
77    }
78}
79
80impl From<Relation> for GivenRowEntry {
81    fn from(value: Relation) -> Self {
82        Self::Relation(value)
83    }
84}
85
86impl From<Attribute> for GivenRowEntry {
87    fn from(value: Attribute) -> Self {
88        Self::Attribute(value)
89    }
90}
91
92impl From<Value> for GivenRowEntry {
93    fn from(value: Value) -> Self {
94        Self::Value(value)
95    }
96}
97
98/// Rows of data to be used as input to a TypeQL query using a `given` stage.
99///
100/// # Examples
101///
102/// ```rust
103/// let query = "given $n: string, $a: integer; insert $p isa person, has name == $n, has age == $a;";
104/// let mut rows = GivenRows::new(vec!["n".to_owned(), "a".to_owned()], 1);
105/// rows.push_row(vec![GivenRowEntry::Value(Value::String("Alice".to_owned())), GivenRowEntry::Value(Value::Integer(28))]);
106/// rows.push_row(vec!["Bob".into(), 26.into()]);
107/// transaction.query_with_rows(query, rows)
108/// ```
109#[derive(Debug, Clone)]
110pub struct GivenRows {
111    pub header: Arc<GivenRowsHeader>,
112    pub(crate) rows: Vec<Vec<GivenRowEntry>>,
113}
114
115impl GivenRows {
116    pub fn new(variables: Vec<String>, row_count_hint: usize) -> Self {
117        Self::new_with_headers(Arc::new(GivenRowsHeader::new(variables)), row_count_hint)
118    }
119
120    pub fn new_with_headers(header: Arc<GivenRowsHeader>, row_count_hint: usize) -> Self {
121        let rows = Vec::with_capacity(row_count_hint);
122        Self { header, rows }
123    }
124
125    pub fn into_parts(self) -> (Arc<GivenRowsHeader>, Vec<Vec<GivenRowEntry>>) {
126        let Self { header, rows } = self;
127        (header, rows)
128    }
129
130    pub fn push(&mut self, row: GivenRow) -> Result {
131        self.push_row(row.row)
132    }
133
134    pub fn push_row(&mut self, row: Vec<GivenRowEntry>) -> Result {
135        if row.len() == self.header.width() {
136            self.rows.push(row);
137            Ok(())
138        } else {
139            Err(Error::Query(QueryError::GivenRowsSizeMismatch { actual: row.len(), expected: self.header.width() }))
140        }
141    }
142
143    pub fn push_map(&mut self, values: impl IntoIterator<Item = (String, GivenRowEntry)>) -> Result {
144        let mut row = GivenRow::new(self.header.clone());
145        values.into_iter().try_for_each(|(var, entry)| row.set(var, entry))?;
146        self.push(row)
147    }
148}
149
150#[derive(Debug, Clone)]
151pub struct GivenRowsHeader {
152    pub(crate) variables: Vec<String>,
153    pub(crate) index: HashMap<String, usize>,
154}
155
156impl GivenRowsHeader {
157    pub fn new(variables: Vec<String>) -> Self {
158        let index = variables.iter().cloned().enumerate().map(|(i, v)| (v, i)).collect();
159        Self { variables, index }
160    }
161
162    pub fn width(&self) -> usize {
163        self.variables.len()
164    }
165}
166
167/// Helper for building a single row to add to a <code>GivenRows</code> instance.
168#[derive(Debug, Clone)]
169pub struct GivenRow {
170    header: Arc<GivenRowsHeader>,
171    row: Vec<GivenRowEntry>,
172}
173
174impl GivenRow {
175    pub fn new(header: Arc<GivenRowsHeader>) -> GivenRow {
176        let mut row = Vec::with_capacity(header.width());
177        row.resize(header.width(), GivenRowEntry::Empty);
178        GivenRow { header, row }
179    }
180
181    pub fn width(&self) -> usize {
182        self.header.width()
183    }
184
185    pub fn set(&mut self, variable: String, entry: GivenRowEntry) -> Result {
186        let index =
187            self.header.index.get(&variable).ok_or(Error::Query(QueryError::GivenRowUnknownVariable { variable }))?;
188        self.set_at(*index, entry)
189    }
190
191    pub fn set_at(&mut self, index: usize, entry: GivenRowEntry) -> Result {
192        if index < self.row.len() {
193            self.row[index] = entry;
194            Ok(())
195        } else {
196            let width = self.header.width();
197            Err(Error::Query(QueryError::GivenRowIndexOutOfBounds { index, width }))
198        }
199    }
200}
typedb_driver/given.rs

typedb_driver/
given.rs
