Skip to main content

imessage_database/util/
query_context.rs

1/*!
2 SQL filter state for message queries.
3*/
4use 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)]
14/// Filters applied to message count and stream queries.
15pub struct QueryContext {
16    /// Inclusive lower-bound timestamp for message dates.
17    pub start: Option<i64>,
18    /// Inclusive upper-bound timestamp for message dates.
19    pub end: Option<i64>,
20    /// Handle row IDs included in the query.
21    pub selected_handle_ids: Option<BTreeSet<i32>>,
22    /// Chat row IDs included in the query.
23    pub selected_chat_ids: Option<BTreeSet<i32>>,
24}
25
26impl QueryContext {
27    /// Set the inclusive lower-bound date from `YYYY-MM-DD`.
28    ///
29    /// # Example:
30    ///
31    /// ```
32    /// use imessage_database::util::query_context::QueryContext;
33    ///
34    /// let mut context = QueryContext::default();
35    /// context.set_start("2023-01-01");
36    /// ```
37    pub fn set_start(&mut self, start: &str) -> Result<(), QueryContextError> {
38        let timestamp = QueryContext::sanitize_date(start)
39            .ok_or(QueryContextError::InvalidDate(start.to_string()))?;
40        self.start = Some(timestamp);
41        Ok(())
42    }
43
44    /// Set the inclusive upper-bound date from `YYYY-MM-DD`.
45    ///
46    /// # Example:
47    ///
48    /// ```
49    /// use imessage_database::util::query_context::QueryContext;
50    ///
51    /// let mut context = QueryContext::default();
52    /// context.set_end("2023-01-01");
53    /// ```
54    pub fn set_end(&mut self, end: &str) -> Result<(), QueryContextError> {
55        let timestamp = QueryContext::sanitize_date(end)
56            .ok_or(QueryContextError::InvalidDate(end.to_string()))?;
57        self.end = Some(timestamp);
58        Ok(())
59    }
60
61    /// Set the handle row IDs included in the query.
62    ///
63    /// # Example:
64    ///
65    /// ```
66    /// use std::collections::BTreeSet;
67    /// use imessage_database::util::query_context::QueryContext;
68    ///
69    /// let mut context = QueryContext::default();
70    /// context.set_selected_handle_ids(BTreeSet::from([1, 2, 3]));
71    /// ```
72    pub fn set_selected_handle_ids(&mut self, selected_handle_ids: BTreeSet<i32>) {
73        self.selected_handle_ids = (!selected_handle_ids.is_empty()).then_some(selected_handle_ids);
74    }
75
76    /// Set the chat row IDs included in the query.
77    ///
78    /// # Example:
79    ///
80    /// ```
81    /// use std::collections::BTreeSet;
82    /// use imessage_database::util::query_context::QueryContext;
83    ///
84    /// let mut context = QueryContext::default();
85    /// context.set_selected_chat_ids(BTreeSet::from([1, 2, 3]));
86    /// ```
87    pub fn set_selected_chat_ids(&mut self, selected_chat_ids: BTreeSet<i32>) {
88        self.selected_chat_ids = (!selected_chat_ids.is_empty()).then_some(selected_chat_ids);
89    }
90
91    /// Parse `YYYY-MM-DD` as a Messages timestamp.
92    fn sanitize_date(date: &str) -> Option<i64> {
93        if date.len() < 9 {
94            return None;
95        }
96
97        let year = date.get(0..4)?.parse::<i32>().ok()?;
98
99        if !date.get(4..5)?.eq("-") {
100            return None;
101        }
102
103        let month = date.get(5..7)?.parse::<u32>().ok()?;
104        if month > 12 {
105            return None;
106        }
107
108        if !date.get(7..8)?.eq("-") {
109            return None;
110        }
111
112        let day = date.get(8..)?.parse::<u32>().ok()?;
113        if day > 31 {
114            return None;
115        }
116
117        let local = Local.with_ymd_and_hms(year, month, day, 0, 0, 0).single()?;
118        let stamp = local.timestamp_nanos_opt().unwrap_or(0);
119
120        Some(stamp - (get_offset() * TIMESTAMP_FACTOR))
121    }
122
123    /// `true` when any filter is set.
124    ///
125    /// # Example:
126    ///
127    /// ```
128    /// use imessage_database::util::query_context::QueryContext;
129    ///
130    /// let mut context = QueryContext::default();
131    /// assert!(!context.has_filters());
132    /// context.set_start("2023-01-01");
133    /// assert!(context.has_filters());
134    /// ```
135    #[must_use]
136    pub fn has_filters(&self) -> bool {
137        self.start.is_some()
138            || self.end.is_some()
139            || self.selected_chat_ids.is_some()
140            || self.selected_handle_ids.is_some()
141    }
142}
143
144#[cfg(test)]
145mod use_tests {
146    use chrono::prelude::*;
147
148    use crate::util::{
149        dates::{TIMESTAMP_FACTOR, format, get_offset},
150        query_context::QueryContext,
151    };
152
153    #[test]
154    fn can_create() {
155        let context = QueryContext::default();
156        assert!(context.start.is_none());
157        assert!(context.end.is_none());
158        assert!(!context.has_filters());
159    }
160
161    #[test]
162    fn can_create_start() {
163        let mut context = QueryContext::default();
164        context.set_start("2020-01-01").unwrap();
165
166        let from_timestamp = DateTime::from_timestamp(
167            (context.start.unwrap() / TIMESTAMP_FACTOR) + get_offset(),
168            0,
169        )
170        .unwrap()
171        .naive_utc();
172        let local = Local.from_utc_datetime(&from_timestamp);
173
174        assert_eq!(format(&local), "Jan 01, 2020 12:00:00 AM");
175        assert!(context.start.is_some());
176        assert!(context.end.is_none());
177        assert!(context.has_filters());
178    }
179
180    #[test]
181    fn can_create_end() {
182        let mut context = QueryContext::default();
183        context.set_end("2020-01-01").unwrap();
184
185        let from_timestamp =
186            DateTime::from_timestamp((context.end.unwrap() / TIMESTAMP_FACTOR) + get_offset(), 0)
187                .unwrap()
188                .naive_utc();
189        let local = Local.from_utc_datetime(&from_timestamp);
190
191        assert_eq!(format(&local), "Jan 01, 2020 12:00:00 AM");
192        assert!(context.start.is_none());
193        assert!(context.end.is_some());
194        assert!(context.has_filters());
195    }
196
197    #[test]
198    fn can_create_both() {
199        let mut context = QueryContext::default();
200        context.set_start("2020-01-01").unwrap();
201        context.set_end("2020-02-02").unwrap();
202
203        let from_timestamp = DateTime::from_timestamp(
204            (context.start.unwrap() / TIMESTAMP_FACTOR) + get_offset(),
205            0,
206        )
207        .unwrap()
208        .naive_utc();
209        let local_start = Local.from_utc_datetime(&from_timestamp);
210
211        let from_timestamp =
212            DateTime::from_timestamp((context.end.unwrap() / TIMESTAMP_FACTOR) + get_offset(), 0)
213                .unwrap()
214                .naive_utc();
215        let local_end = Local.from_utc_datetime(&from_timestamp);
216
217        assert_eq!(format(&local_start), "Jan 01, 2020 12:00:00 AM");
218        assert_eq!(format(&local_end), "Feb 02, 2020 12:00:00 AM");
219        assert!(context.start.is_some());
220        assert!(context.end.is_some());
221        assert!(context.has_filters());
222    }
223}
224
225#[cfg(test)]
226mod id_tests {
227    use std::collections::BTreeSet;
228
229    use crate::util::query_context::QueryContext;
230
231    #[test]
232    fn test_can_set_selected_chat_ids() {
233        let mut qc = QueryContext::default();
234        qc.set_selected_chat_ids(BTreeSet::from([1, 2, 3]));
235
236        assert_eq!(qc.selected_chat_ids, Some(BTreeSet::from([1, 2, 3])));
237        assert!(qc.has_filters());
238    }
239
240    #[test]
241    fn test_can_set_selected_chat_ids_empty() {
242        let mut qc = QueryContext::default();
243        qc.set_selected_chat_ids(BTreeSet::new());
244
245        assert_eq!(qc.selected_chat_ids, None);
246        assert!(!qc.has_filters());
247    }
248
249    #[test]
250    fn test_can_overwrite_selected_chat_ids_empty() {
251        let mut qc = QueryContext::default();
252        qc.set_selected_chat_ids(BTreeSet::from([1, 2, 3]));
253        qc.set_selected_chat_ids(BTreeSet::new());
254
255        assert_eq!(qc.selected_chat_ids, None);
256        assert!(!qc.has_filters());
257    }
258
259    #[test]
260    fn test_can_set_selected_handle_ids() {
261        let mut qc = QueryContext::default();
262        qc.set_selected_handle_ids(BTreeSet::from([1, 2, 3]));
263
264        assert_eq!(qc.selected_handle_ids, Some(BTreeSet::from([1, 2, 3])));
265        assert!(qc.has_filters());
266    }
267
268    #[test]
269    fn test_can_set_selected_handle_ids_empty() {
270        let mut qc = QueryContext::default();
271        qc.set_selected_handle_ids(BTreeSet::new());
272
273        assert_eq!(qc.selected_handle_ids, None);
274        assert!(!qc.has_filters());
275    }
276
277    #[test]
278    fn test_can_overwrite_selected_handle_ids_empty() {
279        let mut qc = QueryContext::default();
280        qc.set_selected_handle_ids(BTreeSet::from([1, 2, 3]));
281        qc.set_selected_handle_ids(BTreeSet::new());
282
283        assert_eq!(qc.selected_handle_ids, None);
284        assert!(!qc.has_filters());
285    }
286}
287
288#[cfg(test)]
289mod sanitize_tests {
290    use crate::util::query_context::QueryContext;
291
292    #[test]
293    fn can_sanitize_good() {
294        let res = QueryContext::sanitize_date("2020-01-01");
295        assert!(res.is_some());
296    }
297
298    #[test]
299    fn can_reject_bad_short() {
300        let res = QueryContext::sanitize_date("1-1-20");
301        assert!(res.is_none());
302    }
303
304    #[test]
305    fn can_reject_bad_order() {
306        let res = QueryContext::sanitize_date("01-01-2020");
307        assert!(res.is_none());
308    }
309
310    #[test]
311    fn can_reject_bad_month() {
312        let res = QueryContext::sanitize_date("2020-31-01");
313        assert!(res.is_none());
314    }
315
316    #[test]
317    fn can_reject_bad_day() {
318        let res = QueryContext::sanitize_date("2020-01-32");
319        assert!(res.is_none());
320    }
321
322    #[test]
323    fn can_reject_bad_data() {
324        let res = QueryContext::sanitize_date("2020-AB-CD");
325        assert!(res.is_none());
326    }
327
328    #[test]
329    fn can_reject_wrong_hyphen() {
330        let res = QueryContext::sanitize_date("2020–01–01");
331        assert!(res.is_none());
332    }
333}