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