imessage_database/util/
query_context.rs1use 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)]
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<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 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 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 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 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 #[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 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}