tank_tests/
metrics.rs

1use std::sync::LazyLock;
2use tank::{DataSet, Entity, Executor, cols, expr, stream::TryStreamExt};
3use time::{Date, macros::date};
4use tokio::sync::Mutex;
5
6#[derive(Entity)]
7#[tank(primary_key = (name, country, date, value))]
8pub struct Metric {
9    person: String,
10    country: String,
11    #[tank(clustering_key)]
12    date: Date,
13    name: String,
14    #[tank(clustering_key)]
15    value: f64,
16}
17static MUTEX: LazyLock<Mutex<()>> = LazyLock::new(|| Mutex::new(()));
18
19pub async fn metrics<E: Executor>(executor: &mut E) {
20    let _lock = MUTEX.lock().await;
21
22    // Setup
23    Metric::drop_table(executor, true, false)
24        .await
25        .expect("Failed to drop Metric table");
26    Metric::create_table(executor, true, true)
27        .await
28        .expect("Failed to create Metric table");
29
30    // Insert
31    let values = vec![
32        // Alice, IT (female)
33        Metric {
34            person: "alice".into(),
35            country: "IT".into(),
36            date: date!(2024 - 1 - 1),
37            name: "height_cm".into(),
38            value: 165.0,
39        },
40        Metric {
41            person: "alice".into(),
42            country: "IT".into(),
43            date: date!(2024 - 1 - 1),
44            name: "weight_kg".into(),
45            value: 62.5,
46        },
47        Metric {
48            person: "alice".into(),
49            country: "IT".into(),
50            date: date!(2024 - 6 - 1),
51            name: "weight_kg".into(),
52            value: 64.2,
53        },
54        Metric {
55            person: "alice".into(),
56            country: "IT".into(),
57            date: date!(2024 - 6 - 1),
58            name: "income_eur".into(),
59            value: 56000.0,
60        },
61        Metric {
62            person: "alice".into(),
63            country: "IT".into(),
64            date: date!(2025 - 4 - 1),
65            name: "income_eur".into(),
66            value: 68000.0,
67        },
68        // Bob, NL (male)
69        Metric {
70            person: "bob".into(),
71            country: "NL".into(),
72            date: date!(2024 - 1 - 1),
73            name: "height_cm".into(),
74            value: 188.0,
75        },
76        Metric {
77            person: "bob".into(),
78            country: "NL".into(),
79            date: date!(2024 - 1 - 1),
80            name: "weight_kg".into(),
81            value: 81.0,
82        },
83        Metric {
84            person: "bob".into(),
85            country: "NL".into(),
86            date: date!(2025 - 6 - 12),
87            name: "weight_kg".into(),
88            value: 82.5,
89        },
90        Metric {
91            person: "bob".into(),
92            country: "NL".into(),
93            date: date!(2024 - 12 - 1),
94            name: "income_eur".into(),
95            value: 88000.0,
96        },
97        Metric {
98            person: "bob".into(),
99            country: "NL".into(),
100            date: date!(2025 - 3 - 2),
101            name: "income_eur".into(),
102            value: 120000.0,
103        },
104        Metric {
105            person: "bob".into(),
106            country: "NL".into(),
107            date: date!(2026 - 2 - 2),
108            name: "income_eur".into(),
109            value: 130000.0,
110        },
111        // Clara, DE (female)
112        Metric {
113            person: "clara".into(),
114            country: "DE".into(),
115            date: date!(2024 - 2 - 1),
116            name: "height_cm".into(),
117            value: 170.0,
118        },
119        Metric {
120            person: "clara".into(),
121            country: "DE".into(),
122            date: date!(2024 - 2 - 1),
123            name: "weight_kg".into(),
124            value: 60.0,
125        },
126        Metric {
127            person: "clara".into(),
128            country: "DE".into(),
129            date: date!(2025 - 1 - 1),
130            name: "income_eur".into(),
131            value: 72000.0,
132        },
133        // David, UK (male)
134        Metric {
135            person: "david".into(),
136            country: "UK".into(),
137            date: date!(2024 - 3 - 10),
138            name: "height_cm".into(),
139            value: 182.0,
140        },
141        Metric {
142            person: "david".into(),
143            country: "UK".into(),
144            date: date!(2024 - 3 - 10),
145            name: "weight_kg".into(),
146            value: 86.0,
147        },
148        Metric {
149            person: "david".into(),
150            country: "UK".into(),
151            date: date!(2024 - 11 - 1),
152            name: "income_gbp".into(),
153            value: 65000.0,
154        },
155        Metric {
156            person: "david".into(),
157            country: "UK".into(),
158            date: date!(2025 - 11 - 1),
159            name: "income_gbp".into(),
160            value: 72000.0,
161        },
162        // Eva, ES (female)
163        Metric {
164            person: "eva".into(),
165            country: "ES".into(),
166            date: date!(2024 - 4 - 5),
167            name: "height_cm".into(),
168            value: 162.0,
169        },
170        Metric {
171            person: "eva".into(),
172            country: "ES".into(),
173            date: date!(2024 - 4 - 5),
174            name: "weight_kg".into(),
175            value: 58.0,
176        },
177        Metric {
178            person: "eva".into(),
179            country: "ES".into(),
180            date: date!(2024 - 12 - 1),
181            name: "income_eur".into(),
182            value: 42000.0,
183        },
184        Metric {
185            person: "eva".into(),
186            country: "ES".into(),
187            date: date!(2025 - 12 - 1),
188            name: "income_eur".into(),
189            value: 47000.0,
190        },
191        // Marco, IT (male)
192        Metric {
193            person: "marco".into(),
194            country: "IT".into(),
195            date: date!(2024 - 1 - 15),
196            name: "height_cm".into(),
197            value: 178.0,
198        },
199        Metric {
200            person: "marco".into(),
201            country: "IT".into(),
202            date: date!(2024 - 1 - 15),
203            name: "weight_kg".into(),
204            value: 78.0,
205        },
206        Metric {
207            person: "marco".into(),
208            country: "IT".into(),
209            date: date!(2024 - 10 - 1),
210            name: "weight_kg".into(),
211            value: 80.5,
212        },
213        Metric {
214            person: "marco".into(),
215            country: "IT".into(),
216            date: date!(2024 - 10 - 1),
217            name: "income_eur".into(),
218            value: 61000.0,
219        },
220        Metric {
221            person: "marco".into(),
222            country: "IT".into(),
223            date: date!(2025 - 10 - 1),
224            name: "income_eur".into(),
225            value: 72000.0,
226        },
227        // Sophie, UK (female)
228        Metric {
229            person: "sophie".into(),
230            country: "UK".into(),
231            date: date!(2024 - 2 - 20),
232            name: "height_cm".into(),
233            value: 168.0,
234        },
235        Metric {
236            person: "sophie".into(),
237            country: "UK".into(),
238            date: date!(2024 - 2 - 20),
239            name: "weight_kg".into(),
240            value: 61.0,
241        },
242        Metric {
243            person: "sophie".into(),
244            country: "UK".into(),
245            date: date!(2024 - 12 - 1),
246            name: "income_gbp".into(),
247            value: 52000.0,
248        },
249        Metric {
250            person: "sophie".into(),
251            country: "UK".into(),
252            date: date!(2025 - 12 - 1),
253            name: "income_gbp".into(),
254            value: 58000.0,
255        },
256        // Partition isolation sanity check
257        Metric {
258            person: "alice".into(),
259            country: "CA".into(),
260            date: date!(2024 - 6 - 1),
261            name: "income_usd".into(),
262            value: 69000.0,
263        },
264    ];
265
266    Metric::insert_many(executor, &values)
267        .await
268        .expect("Could not insert the entities");
269
270    #[derive(Entity)]
271    struct MetricValue {
272        pub value: f64,
273    }
274
275    let date = date!(2000 - 01 - 01);
276    let heights = Metric::table()
277        .select(
278            executor,
279            cols!(Metric::date DESC, Metric::value DESC),
280            expr!(
281                Metric::name == "height_cm"
282                    && Metric::country == "IT"
283                    && Metric::date >= #date
284            ),
285            None,
286        )
287        .map_ok(|v| {
288            MetricValue::from_row(v)
289                .expect("Could not read the value")
290                .value
291        })
292        .try_collect::<Vec<_>>()
293        .await
294        .expect("Coult not get the Italy height values");
295    assert_eq!(heights, [178.0, 165.0]);
296}