imessage_database/util/
query_context.rs1use std::collections::BTreeSet;
5
6use chrono::prelude::*;
7
8use crate::{
9 error::query_context::QueryContextError,
10 util::dates::{TIMESTAMP_FACTOR, get_offset},
11};
12
13#[derive(Debug, Default, PartialEq, Eq)]
14pub struct QueryContext {
16 pub start: Option<i64>,
18 pub end: Option<i64>,
20 pub selected_handle_ids: Option<BTreeSet<i32>>,
22 pub selected_chat_ids: Option<BTreeSet<i32>>,
24}
25
26impl QueryContext {
27 pub fn set_start(&mut self, start: &str) -> Result<(), QueryContextError> {
37 let timestamp = QueryContext::sanitize_date(start)
38 .ok_or(QueryContextError::InvalidDate(start.to_string()))?;
39 self.start = Some(timestamp);
40 Ok(())
41 }
42
43 pub fn set_end(&mut self, end: &str) -> Result<(), QueryContextError> {
53 let timestamp = QueryContext::sanitize_date(end)
54 .ok_or(QueryContextError::InvalidDate(end.to_string()))?;
55 self.end = Some(timestamp);
56 Ok(())
57 }
58
59 pub fn set_selected_handle_ids(&mut self, selected_handle_ids: BTreeSet<i32>) {
71 self.selected_handle_ids = (!selected_handle_ids.is_empty()).then_some(selected_handle_ids);
72 }
73
74 pub fn set_selected_chat_ids(&mut self, selected_chat_ids: BTreeSet<i32>) {
86 self.selected_chat_ids = (!selected_chat_ids.is_empty()).then_some(selected_chat_ids);
87 }
88
89 fn sanitize_date(date: &str) -> Option<i64> {
91 if date.len() < 9 {
92 return None;
93 }
94
95 let year = date.get(0..4)?.parse::<i32>().ok()?;
96
97 if !date.get(4..5)?.eq("-") {
98 return None;
99 }
100
101 let month = date.get(5..7)?.parse::<u32>().ok()?;
102 if month > 12 {
103 return None;
104 }
105
106 if !date.get(7..8)?.eq("-") {
107 return None;
108 }
109
110 let day = date.get(8..)?.parse::<u32>().ok()?;
111 if day > 31 {
112 return None;
113 }
114
115 let local = Local.with_ymd_and_hms(year, month, day, 0, 0, 0).single()?;
116 let stamp = local.timestamp_nanos_opt().unwrap_or(0);
117
118 Some(stamp - (get_offset() * TIMESTAMP_FACTOR))
119 }
120
121 pub fn has_filters(&self) -> bool {
134 self.start.is_some()
135 || self.end.is_some()
136 || self.selected_chat_ids.is_some()
137 || self.selected_handle_ids.is_some()
138 }
139}
140
141#[cfg(test)]
142mod use_tests {
143 use chrono::prelude::*;
144
145 use crate::util::{
146 dates::{TIMESTAMP_FACTOR, format, get_offset},
147 query_context::QueryContext,
148 };
149
150 #[test]
151 fn can_create() {
152 let context = QueryContext::default();
153 assert!(context.start.is_none());
154 assert!(context.end.is_none());
155 assert!(!context.has_filters());
156 }
157
158 #[test]
159 fn can_create_start() {
160 let mut context = QueryContext::default();
161 context.set_start("2020-01-01").unwrap();
162
163 let from_timestamp = DateTime::from_timestamp(
164 (context.start.unwrap() / TIMESTAMP_FACTOR) + get_offset(),
165 0,
166 )
167 .unwrap()
168 .naive_utc();
169 let local = Local.from_utc_datetime(&from_timestamp);
170
171 assert_eq!(format(&Ok(local)), "Jan 01, 2020 12:00:00 AM");
172 assert!(context.start.is_some());
173 assert!(context.end.is_none());
174 assert!(context.has_filters());
175 }
176
177 #[test]
178 fn can_create_end() {
179 let mut context = QueryContext::default();
180 context.set_end("2020-01-01").unwrap();
181
182 let from_timestamp =
183 DateTime::from_timestamp((context.end.unwrap() / TIMESTAMP_FACTOR) + get_offset(), 0)
184 .unwrap()
185 .naive_utc();
186 let local = Local.from_utc_datetime(&from_timestamp);
187
188 assert_eq!(format(&Ok(local)), "Jan 01, 2020 12:00:00 AM");
189 assert!(context.start.is_none());
190 assert!(context.end.is_some());
191 assert!(context.has_filters());
192 }
193
194 #[test]
195 fn can_create_both() {
196 let mut context = QueryContext::default();
197 context.set_start("2020-01-01").unwrap();
198 context.set_end("2020-02-02").unwrap();
199
200 let from_timestamp = DateTime::from_timestamp(
201 (context.start.unwrap() / TIMESTAMP_FACTOR) + get_offset(),
202 0,
203 )
204 .unwrap()
205 .naive_utc();
206 let local_start = Local.from_utc_datetime(&from_timestamp);
207
208 let from_timestamp =
209 DateTime::from_timestamp((context.end.unwrap() / TIMESTAMP_FACTOR) + get_offset(), 0)
210 .unwrap()
211 .naive_utc();
212 let local_end = Local.from_utc_datetime(&from_timestamp);
213
214 assert_eq!(format(&Ok(local_start)), "Jan 01, 2020 12:00:00 AM");
215 assert_eq!(format(&Ok(local_end)), "Feb 02, 2020 12:00:00 AM");
216 assert!(context.start.is_some());
217 assert!(context.end.is_some());
218 assert!(context.has_filters());
219 }
220}
221
222#[cfg(test)]
223mod id_tests {
224 use std::collections::BTreeSet;
225
226 use crate::util::query_context::QueryContext;
227
228 #[test]
229 fn test_can_set_selected_chat_ids() {
230 let mut qc = QueryContext::default();
231 qc.set_selected_chat_ids(BTreeSet::from([1, 2, 3]));
232
233 assert_eq!(qc.selected_chat_ids, Some(BTreeSet::from([1, 2, 3])));
234 assert!(qc.has_filters());
235 }
236
237 #[test]
238 fn test_can_set_selected_chat_ids_empty() {
239 let mut qc = QueryContext::default();
240 qc.set_selected_chat_ids(BTreeSet::new());
241
242 assert_eq!(qc.selected_chat_ids, None);
243 assert!(!qc.has_filters());
244 }
245
246 #[test]
247 fn test_can_overwrite_selected_chat_ids_empty() {
248 let mut qc = QueryContext::default();
249 qc.set_selected_chat_ids(BTreeSet::from([1, 2, 3]));
250 qc.set_selected_chat_ids(BTreeSet::new());
251
252 assert_eq!(qc.selected_chat_ids, None);
253 assert!(!qc.has_filters());
254 }
255
256 #[test]
257 fn test_can_set_selected_handle_ids() {
258 let mut qc = QueryContext::default();
259 qc.set_selected_handle_ids(BTreeSet::from([1, 2, 3]));
260
261 assert_eq!(qc.selected_handle_ids, Some(BTreeSet::from([1, 2, 3])));
262 assert!(qc.has_filters());
263 }
264
265 #[test]
266 fn test_can_set_selected_handle_ids_empty() {
267 let mut qc = QueryContext::default();
268 qc.set_selected_handle_ids(BTreeSet::new());
269
270 assert_eq!(qc.selected_handle_ids, None);
271 assert!(!qc.has_filters());
272 }
273
274 #[test]
275 fn test_can_overwrite_selected_handle_ids_empty() {
276 let mut qc = QueryContext::default();
277 qc.set_selected_handle_ids(BTreeSet::from([1, 2, 3]));
278 qc.set_selected_handle_ids(BTreeSet::new());
279
280 assert_eq!(qc.selected_handle_ids, None);
281 assert!(!qc.has_filters());
282 }
283}
284
285#[cfg(test)]
286mod sanitize_tests {
287 use crate::util::query_context::QueryContext;
288
289 #[test]
290 fn can_sanitize_good() {
291 let res = QueryContext::sanitize_date("2020-01-01");
292 assert!(res.is_some());
293 }
294
295 #[test]
296 fn can_reject_bad_short() {
297 let res = QueryContext::sanitize_date("1-1-20");
298 assert!(res.is_none());
299 }
300
301 #[test]
302 fn can_reject_bad_order() {
303 let res = QueryContext::sanitize_date("01-01-2020");
304 assert!(res.is_none());
305 }
306
307 #[test]
308 fn can_reject_bad_month() {
309 let res = QueryContext::sanitize_date("2020-31-01");
310 assert!(res.is_none());
311 }
312
313 #[test]
314 fn can_reject_bad_day() {
315 let res = QueryContext::sanitize_date("2020-01-32");
316 assert!(res.is_none());
317 }
318
319 #[test]
320 fn can_reject_bad_data() {
321 let res = QueryContext::sanitize_date("2020-AB-CD");
322 assert!(res.is_none());
323 }
324
325 #[test]
326 fn can_reject_wrong_hyphen() {
327 let res = QueryContext::sanitize_date("2020–01–01");
328 assert!(res.is_none());
329 }
330}