gcloud_spanner/
statement.rs

1use std::collections::{BTreeMap, HashMap};
2
3use base64::prelude::*;
4use prost_types::value::Kind;
5use prost_types::value::Kind::StringValue;
6use prost_types::{value, ListValue, Struct, Value};
7use time::format_description::well_known::Rfc3339;
8use time::macros::format_description;
9use time::{Date, OffsetDateTime};
10
11use google_cloud_googleapis::spanner::v1::struct_type::Field;
12use google_cloud_googleapis::spanner::v1::{StructType, Type, TypeAnnotationCode, TypeCode};
13
14use crate::bigdecimal::BigDecimal;
15use crate::value::CommitTimestamp;
16
17/// A Statement is a SQL query with named parameters.
18///
19/// A parameter placeholder consists of '@' followed by the parameter name.
20/// The parameter name is an identifier which must conform to the naming
21/// requirements in <https://cloud.google.com/spanner/docs/lexical#identifiers>.
22/// Parameters may appear anywhere that a literal value is expected. The same
23/// parameter name may be used more than once.  It is an error to execute a
24/// statement with unbound parameters. On the other hand, it is allowable to
25/// bind parameter names that are not used.
26///
27/// See the documentation of the Row type for how Go types are mapped to Cloud
28/// Spanner types.
29#[derive(Clone)]
30pub struct Statement {
31    pub(crate) sql: String,
32    pub(crate) params: BTreeMap<String, Value>,
33    pub(crate) param_types: HashMap<String, Type>,
34}
35
36impl Statement {
37    /// new returns a Statement with the given SQL and an empty Params map.
38    pub fn new<T: Into<String>>(sql: T) -> Self {
39        Statement {
40            sql: sql.into(),
41            params: Default::default(),
42            param_types: Default::default(),
43        }
44    }
45
46    /// add_params add the bind parameter.
47    /// Implement the ToKind trait to use non-predefined types.
48    pub fn add_param<T>(&mut self, name: &str, value: &T)
49    where
50        T: ToKind,
51    {
52        self.param_types.insert(name.to_string(), T::get_type());
53        self.params.insert(
54            name.to_string(),
55            Value {
56                kind: Some(value.to_kind()),
57            },
58        );
59    }
60}
61
62pub fn single_type<T>(code: T) -> Type
63where
64    T: Into<i32>,
65{
66    Type {
67        code: code.into(),
68        array_element_type: None,
69        struct_type: None,
70        //TODO support PG Numeric
71        type_annotation: TypeAnnotationCode::Unspecified.into(),
72        proto_type_fqn: "".to_string(),
73    }
74}
75
76pub trait ToKind {
77    fn to_kind(&self) -> value::Kind;
78    fn get_type() -> Type
79    where
80        Self: Sized;
81}
82
83pub type Kinds = Vec<(&'static str, Kind)>;
84pub type Types = Vec<(&'static str, Type)>;
85
86pub trait ToStruct {
87    fn to_kinds(&self) -> Kinds;
88    fn get_types() -> Types
89    where
90        Self: Sized;
91}
92
93impl<T> ToStruct for &T
94where
95    T: ToStruct,
96{
97    fn to_kinds(&self) -> Kinds {
98        (*self).to_kinds()
99    }
100
101    fn get_types() -> Types
102    where
103        Self: Sized,
104    {
105        T::get_types()
106    }
107}
108
109impl ToKind for String {
110    fn to_kind(&self) -> Kind {
111        StringValue(self.clone())
112    }
113    fn get_type() -> Type {
114        single_type(TypeCode::String)
115    }
116}
117
118impl ToKind for &str {
119    fn to_kind(&self) -> Kind {
120        StringValue(self.to_string())
121    }
122    fn get_type() -> Type {
123        single_type(TypeCode::String)
124    }
125}
126
127impl ToKind for i64 {
128    fn to_kind(&self) -> Kind {
129        self.to_string().to_kind()
130    }
131    fn get_type() -> Type {
132        single_type(TypeCode::Int64)
133    }
134}
135
136impl ToKind for f64 {
137    fn to_kind(&self) -> Kind {
138        value::Kind::NumberValue(*self)
139    }
140    fn get_type() -> Type {
141        single_type(TypeCode::Float64)
142    }
143}
144
145impl ToKind for bool {
146    fn to_kind(&self) -> Kind {
147        value::Kind::BoolValue(*self)
148    }
149    fn get_type() -> Type {
150        single_type(TypeCode::Bool)
151    }
152}
153
154impl ToKind for Date {
155    fn to_kind(&self) -> Kind {
156        self.format(format_description!("[year]-[month]-[day]"))
157            .unwrap()
158            .to_kind()
159    }
160    fn get_type() -> Type {
161        single_type(TypeCode::Date)
162    }
163}
164
165impl ToKind for OffsetDateTime {
166    fn to_kind(&self) -> Kind {
167        self.format(&Rfc3339).unwrap().to_kind()
168    }
169    fn get_type() -> Type {
170        single_type(TypeCode::Timestamp)
171    }
172}
173
174impl ToKind for CommitTimestamp {
175    fn to_kind(&self) -> Kind {
176        "spanner.commit_timestamp()".to_kind()
177    }
178    fn get_type() -> Type {
179        single_type(TypeCode::Timestamp)
180    }
181}
182
183impl ToKind for &[u8] {
184    fn to_kind(&self) -> Kind {
185        BASE64_STANDARD.encode(self).to_kind()
186    }
187    fn get_type() -> Type {
188        single_type(TypeCode::Bytes)
189    }
190}
191
192impl ToKind for Vec<u8> {
193    fn to_kind(&self) -> Kind {
194        BASE64_STANDARD.encode(self).to_kind()
195    }
196    fn get_type() -> Type {
197        single_type(TypeCode::Bytes)
198    }
199}
200
201impl ToKind for BigDecimal {
202    fn to_kind(&self) -> Kind {
203        self.to_string().to_kind()
204    }
205    fn get_type() -> Type {
206        single_type(TypeCode::Numeric)
207    }
208}
209
210impl ToKind for ::prost_types::Timestamp {
211    fn to_kind(&self) -> Kind {
212        // The protobuf timestamp type should be formatted in RFC3339
213        // See here for more details: https://docs.rs/prost-types/latest/prost_types/struct.Timestamp.html
214        let rfc3339 = format!("{self}");
215        rfc3339.to_kind()
216    }
217
218    fn get_type() -> Type
219    where
220        Self: Sized,
221    {
222        single_type(TypeCode::Timestamp)
223    }
224}
225
226impl<T> ToKind for T
227where
228    T: ToStruct,
229{
230    fn to_kind(&self) -> Kind {
231        let mut fields = BTreeMap::<String, Value>::default();
232        self.to_kinds().into_iter().for_each(|e| {
233            fields.insert(e.0.into(), Value { kind: Some(e.1) });
234        });
235        Kind::StructValue(Struct { fields })
236    }
237    fn get_type() -> Type {
238        Type {
239            code: TypeCode::Struct.into(),
240            array_element_type: None,
241            type_annotation: TypeAnnotationCode::Unspecified.into(),
242            struct_type: Some(StructType {
243                fields: T::get_types()
244                    .into_iter()
245                    .map(|t| Field {
246                        name: t.0.into(),
247                        r#type: Some(t.1),
248                    })
249                    .collect(),
250            }),
251            proto_type_fqn: "".to_string(),
252        }
253    }
254}
255
256impl<T> ToKind for Option<T>
257where
258    T: ToKind,
259{
260    fn to_kind(&self) -> Kind {
261        match self {
262            Some(vv) => vv.to_kind(),
263            None => value::Kind::NullValue(prost_types::NullValue::NullValue.into()),
264        }
265    }
266    fn get_type() -> Type {
267        T::get_type()
268    }
269}
270
271impl<T> ToKind for Vec<T>
272where
273    T: ToKind,
274{
275    #[inline]
276    fn to_kind(&self) -> Kind {
277        self.as_slice().to_kind()
278    }
279
280    #[inline]
281    fn get_type() -> Type {
282        <&[T] as ToKind>::get_type()
283    }
284}
285
286impl<T> ToKind for &[T]
287where
288    T: ToKind,
289{
290    fn to_kind(&self) -> Kind {
291        value::Kind::ListValue(ListValue {
292            values: self
293                .iter()
294                .map(|x| Value {
295                    kind: Some(x.to_kind()),
296                })
297                .collect(),
298        })
299    }
300
301    fn get_type() -> Type {
302        Type {
303            code: TypeCode::Array.into(),
304            array_element_type: Some(Box::new(T::get_type())),
305            struct_type: None,
306            type_annotation: TypeAnnotationCode::Unspecified.into(),
307            proto_type_fqn: "".to_string(),
308        }
309    }
310}
311
312#[cfg(test)]
313mod test {
314    use crate::statement::ToKind;
315    use prost_types::value::Kind;
316    use time::OffsetDateTime;
317
318    // Test that prost's to_kind implementation works as expected.
319    #[test]
320    fn prost_timestamp_to_kind_works() {
321        let ts = ::prost_types::Timestamp::date_time(2024, 1, 1, 12, 15, 36).unwrap();
322        let expected = String::from("2024-01-01T12:15:36Z");
323        // Make sure the formatting of prost_types::Timestamp hasn't changed
324        assert_eq!(format!("{ts:}"), expected);
325        let kind = ts.to_kind();
326        matches!(kind, Kind::StringValue(s) if s == expected);
327
328        // Prost's Timestamp type and OffsetDateTime should have the same representation in spanner
329        assert_eq!(prost_types::Timestamp::get_type(), OffsetDateTime::get_type());
330    }
331}