pancake_db_client/
partition_helpers.rs

1use std::time::SystemTime;
2
3use pancake_db_idl::dml::partition_field_value::Value;
4
5/// Re-export for the purpose of [`make_partition`].
6pub use pancake_db_idl::dml::PartitionFieldValue;
7use prost_types::Timestamp;
8
9/// Trait used by [`make_partition`] to convert native types to Pancake IDL types.
10pub trait PartitionFieldValueConverter {
11  fn to_value(self) -> Value;
12}
13
14impl PartitionFieldValueConverter for i64 {
15  fn to_value(self) -> Value {
16    Value::Int64Val(self)
17  }
18}
19
20impl PartitionFieldValueConverter for bool {
21  fn to_value(self) -> Value {
22    Value::BoolVal(self)
23  }
24}
25
26impl PartitionFieldValueConverter for String {
27  fn to_value(self) -> Value {
28    Value::StringVal(self)
29  }
30}
31
32impl PartitionFieldValueConverter for SystemTime {
33  fn to_value(self) -> Value {
34    Value::TimestampVal(Timestamp::from(self))
35  }
36}
37
38/// Helper macro to support [`make_partition`].
39#[macro_export]
40macro_rules! make_partition_insert {
41  {$partition: expr;} => {};
42  {$partition: expr; $key:expr => $val:expr $(,$keys:expr => $vals:expr)* $(,)?} => {
43    let fv = $crate::partition_helpers::PartitionFieldValue {
44      value: Some($crate::partition_helpers::PartitionFieldValueConverter::to_value($val)),
45    };
46    $partition.insert($key.to_string(), fv);
47    $crate::make_partition_insert! { $partition; $($keys => $vals),* }
48  };
49}
50
51/// Outputs a partition, given native Rust key => value pairings.
52///
53/// Since instantiating protobuf-generated types is very verbose,
54/// this macro exists to make partitions
55/// (`HashMap<String, PartitionFieldValue>`) with ease:
56///
57/// ```
58/// use pancake_db_client::make_partition;
59/// use std::time::SystemTime;
60///
61/// let my_partition = make_partition! {
62///   "t" => SystemTime::now(),
63///   "action" => "click".to_string(),
64///   "is_final" => true,
65///   "int_bucket" => 7,
66/// };
67/// ```
68///
69/// Keys can be any type supporting `.to_string()`.
70/// Values can be `i64`s, `bool`s, `String`s, or `Timestamp`s.
71#[macro_export]
72macro_rules! make_partition {
73  {} => {
74    std::collections::HashMap::<String, $crate::partition_helpers::PartitionFieldValue>::new()
75  };
76  {$($keys:expr => $vals:expr),+ $(,)?} => {
77    {
78      let mut row = std::collections::HashMap::<String, $crate::partition_helpers::PartitionFieldValue>::new();
79      $crate::make_partition_insert! { row; $($keys => $vals),+ }
80      row
81    }
82  };
83}
84
85#[cfg(test)]
86mod tests {
87  use std::collections::HashMap;
88  use std::time::SystemTime;
89
90  use pancake_db_idl::dml::partition_field_value::Value;
91  use pancake_db_idl::dml::PartitionFieldValue;
92  use prost_types::Timestamp;
93
94  use crate::make_partition;
95
96  #[test]
97  fn test_partition_macro() {
98    let timestamp = SystemTime::now();
99    let p0 = make_partition! {};
100    let p1 = make_partition! { "i64" => 5_i64 };
101    let p2 = make_partition! {
102      "i64" => 5_i64,
103      "bool" => true,
104      "timestamp" => timestamp.clone(),
105      "string" => "asdf".to_string(),
106    };
107
108    assert!(p0.is_empty());
109
110    assert_eq!(p1.len(), 1);
111
112    assert_eq!(p2.len(), 4);
113    fn assert_val_eq(partition: &HashMap<String, PartitionFieldValue>, key: &str, value: Value) {
114      assert_eq!(partition[key].clone(), PartitionFieldValue {
115        value: Some(value),
116      });
117    }
118    assert_val_eq(&p2, "i64", Value::Int64Val(5));
119    assert_val_eq(&p2, "bool", Value::BoolVal(true));
120    assert_val_eq(&p2, "timestamp", Value::TimestampVal(Timestamp::from(timestamp.clone())));
121    assert_val_eq(&p2, "string", Value::StringVal("asdf".to_string()));
122  }
123}
124
125#[cfg(test)]
126mod tests_no_imports {
127  use crate::make_partition;
128
129  #[test]
130  fn test_partition_macro() {
131    println!("{:?}", make_partition! {});
132    println!("{:?}", make_partition! { "a" => 5_i64 });
133  }
134}
pancake_db_client/partition_helpers.rs

pancake_db_client/
partition_helpers.rs
