Skip to main content

smriti/db/
document_repo.rs

1//! Document repository
2//!
3//! Queries and updates for document categorization and OCR metadata.
4
5use rusqlite::{params, params_from_iter, Connection, Result as SqliteResult};
6
7use crate::models::Photo;
8
9pub struct DocumentRepo<'a> {
10    conn: &'a Connection,
11}
12
13impl<'a> DocumentRepo<'a> {
14    pub fn new(conn: &'a Connection) -> Self {
15        Self { conn }
16    }
17
18    pub fn get_non_photo_documents(&self, limit: i64, offset: i64) -> SqliteResult<Vec<Photo>> {
19        let mut stmt = self.conn.prepare(
20            r#"
21            SELECT
22                id, file_path, file_name, file_hash, file_size,
23                date_taken, date_taken_source,
24                gps_latitude, gps_longitude,
25                location_city, location_country,
26                camera_make, camera_model,
27                iso, aperture, shutter_speed, focal_length,
28                lens_model, flash, gps_altitude,
29                width, height, orientation,
30                media_type, duration_ms, video_codec, audio_codec,
31                frame_rate, bitrate, has_audio,
32                thumbnail_path, faces_processed,
33                content_category, ocr_text, ocr_processed, ocr_confidence,
34                is_favorite,
35                is_trashed, trashed_at,
36                indexed_at, updated_at
37            FROM photos
38            WHERE is_trashed = FALSE AND content_category != 'photo'
39            ORDER BY date_taken DESC
40            LIMIT ?1 OFFSET ?2
41            "#,
42        )?;
43
44        let rows = stmt.query_map(params![limit, offset], crate::db::photo_repo::row_to_photo)?;
45        let mut out = Vec::new();
46        for row in rows {
47            out.push(row?);
48        }
49        Ok(out)
50    }
51
52    pub fn get_documents_by_category(
53        &self,
54        category: &str,
55        limit: i64,
56        offset: i64,
57    ) -> SqliteResult<Vec<Photo>> {
58        let mut stmt = self.conn.prepare(
59            r#"
60            SELECT
61                id, file_path, file_name, file_hash, file_size,
62                date_taken, date_taken_source,
63                gps_latitude, gps_longitude,
64                location_city, location_country,
65                camera_make, camera_model,
66                iso, aperture, shutter_speed, focal_length,
67                lens_model, flash, gps_altitude,
68                width, height, orientation,
69                media_type, duration_ms, video_codec, audio_codec,
70                frame_rate, bitrate, has_audio,
71                thumbnail_path, faces_processed,
72                content_category, ocr_text, ocr_processed, ocr_confidence,
73                is_favorite,
74                is_trashed, trashed_at,
75                indexed_at, updated_at
76            FROM photos
77            WHERE is_trashed = FALSE AND content_category = ?1
78            ORDER BY date_taken DESC
79            LIMIT ?2 OFFSET ?3
80            "#,
81        )?;
82
83        let rows = stmt.query_map(
84            params![category, limit, offset],
85            crate::db::photo_repo::row_to_photo,
86        )?;
87        let mut out = Vec::new();
88        for row in rows {
89            out.push(row?);
90        }
91        Ok(out)
92    }
93
94    pub fn get_documents_by_categories(
95        &self,
96        categories: &[String],
97        limit: i64,
98        offset: i64,
99    ) -> SqliteResult<Vec<Photo>> {
100        if categories.is_empty() {
101            return Ok(Vec::new());
102        }
103        let placeholders = std::iter::repeat_n("?", categories.len())
104            .collect::<Vec<_>>()
105            .join(", ");
106        let sql = format!(
107            r#"
108            SELECT
109                id, file_path, file_name, file_hash, file_size,
110                date_taken, date_taken_source,
111                gps_latitude, gps_longitude,
112                location_city, location_country,
113                camera_make, camera_model,
114                iso, aperture, shutter_speed, focal_length,
115                lens_model, flash, gps_altitude,
116                width, height, orientation,
117                media_type, duration_ms, video_codec, audio_codec,
118                frame_rate, bitrate, has_audio,
119                thumbnail_path, faces_processed,
120                content_category, ocr_text, ocr_processed, ocr_confidence,
121                is_favorite,
122                is_trashed, trashed_at,
123                indexed_at, updated_at
124            FROM photos
125            WHERE is_trashed = FALSE AND content_category IN ({placeholders})
126            ORDER BY date_taken DESC
127            LIMIT ? OFFSET ?
128            "#
129        );
130        let params = categories.iter().map(|s| s as &dyn rusqlite::ToSql).chain([
131            &limit as &dyn rusqlite::ToSql,
132            &offset as &dyn rusqlite::ToSql,
133        ]);
134        let mut stmt = self.conn.prepare(&sql)?;
135        let rows = stmt.query_map(
136            params_from_iter(params),
137            crate::db::photo_repo::row_to_photo,
138        )?;
139        let mut out = Vec::new();
140        for row in rows {
141            out.push(row?);
142        }
143        Ok(out)
144    }
145
146    pub fn search_documents_fts(
147        &self,
148        query: &str,
149        limit: i64,
150        offset: i64,
151    ) -> SqliteResult<Vec<Photo>> {
152        let Some(query) = fts_literal_query(query) else {
153            return Ok(Vec::new());
154        };
155        let mut stmt = self.conn.prepare(
156            r#"
157            SELECT
158                p.id, p.file_path, p.file_name, p.file_hash, p.file_size,
159                p.date_taken, p.date_taken_source,
160                p.gps_latitude, p.gps_longitude,
161                p.location_city, p.location_country,
162                p.camera_make, p.camera_model,
163                p.iso, p.aperture, p.shutter_speed, p.focal_length,
164                p.lens_model, p.flash, p.gps_altitude,
165                p.width, p.height, p.orientation,
166                p.media_type, p.duration_ms, p.video_codec, p.audio_codec,
167                p.frame_rate, p.bitrate, p.has_audio,
168                p.thumbnail_path, p.faces_processed,
169                p.content_category, p.ocr_text, p.ocr_processed, p.ocr_confidence,
170                p.is_favorite,
171                p.is_trashed, p.trashed_at,
172                p.indexed_at, p.updated_at
173            FROM photos p
174            JOIN photos_fts fts ON fts.rowid = p.id
175            WHERE p.is_trashed = FALSE
176                AND p.content_category != 'photo'
177                AND photos_fts MATCH ?1
178            ORDER BY p.date_taken DESC
179            LIMIT ?2 OFFSET ?3
180            "#,
181        )?;
182
183        let rows = stmt.query_map(
184            params![query, limit, offset],
185            crate::db::photo_repo::row_to_photo,
186        )?;
187        let mut out = Vec::new();
188        for row in rows {
189            out.push(row?);
190        }
191        Ok(out)
192    }
193
194    pub fn update_content_category(&self, photo_id: i64, category: &str) -> SqliteResult<usize> {
195        self.conn.execute(
196            "UPDATE photos SET content_category = ?1, updated_at = CURRENT_TIMESTAMP WHERE id = ?2",
197            params![category, photo_id],
198        )
199    }
200
201    pub fn update_ocr_metadata(
202        &self,
203        photo_id: i64,
204        text: Option<&str>,
205        confidence: Option<f32>,
206        processed: bool,
207    ) -> SqliteResult<()> {
208        self.conn.execute(
209            r#"
210            UPDATE photos
211            SET ocr_text = ?1,
212                ocr_confidence = ?2,
213                ocr_processed = ?3,
214                updated_at = CURRENT_TIMESTAMP
215            WHERE id = ?4
216            "#,
217            params![text, confidence, processed, photo_id],
218        )?;
219        Ok(())
220    }
221
222    pub fn get_unprocessed_for_document_analysis(
223        &self,
224        limit: i64,
225    ) -> SqliteResult<Vec<(i64, String, i32)>> {
226        let mut stmt = self.conn.prepare(
227            r#"
228            SELECT id, file_path, COALESCE(orientation, 1)
229            FROM photos
230            WHERE is_trashed = FALSE AND ocr_processed = FALSE
231            ORDER BY date_taken DESC
232            LIMIT ?1
233            "#,
234        )?;
235
236        let rows = stmt.query_map(params![limit], |row| {
237            Ok((
238                row.get::<_, i64>(0)?,
239                row.get::<_, String>(1)?,
240                row.get::<_, i32>(2)?,
241            ))
242        })?;
243
244        let mut out = Vec::new();
245        for row in rows {
246            out.push(row?);
247        }
248        Ok(out)
249    }
250}
251
252fn fts_literal_query(query: &str) -> Option<String> {
253    let terms: Vec<String> = query
254        .split_whitespace()
255        .filter_map(|term| {
256            let cleaned = term.trim_matches(|ch: char| {
257                !ch.is_alphanumeric() && ch != '\'' && ch != '-' && ch != '_'
258            });
259            if cleaned.is_empty() {
260                None
261            } else {
262                Some(format!("\"{}\"", cleaned.replace('"', "\"\"")))
263            }
264        })
265        .collect();
266    if terms.is_empty() {
267        None
268    } else {
269        Some(terms.join(" "))
270    }
271}
272
273#[cfg(test)]
274mod tests {
275    use super::{fts_literal_query, DocumentRepo};
276    use crate::db::create_schema;
277    use rusqlite::Connection;
278
279    #[test]
280    fn fts_literal_query_quotes_user_terms() {
281        assert_eq!(
282            fts_literal_query("invoice: (delhi) \"trip\""),
283            Some("\"invoice\" \"delhi\" \"trip\"".to_string())
284        );
285    }
286
287    #[test]
288    fn fts_literal_query_ignores_operator_only_input() {
289        assert_eq!(fts_literal_query("() OR *"), Some("\"OR\"".to_string()));
290        assert_eq!(fts_literal_query("() *"), None);
291    }
292
293    #[test]
294    fn category_list_pages_across_categories_in_one_order() {
295        let conn = Connection::open_in_memory().unwrap();
296        create_schema(&conn).unwrap();
297        conn.execute(
298            "INSERT INTO photos
299                (id, file_path, file_name, file_hash, file_size, date_taken, content_category, is_trashed)
300             VALUES
301                (1, 'a.jpg', 'a.jpg', 'hash-a', 10, '2026-01-01T00:00:00Z', 'document', 0),
302                (2, 'b.jpg', 'b.jpg', 'hash-b', 10, '2026-01-03T00:00:00Z', 'receipt', 0),
303                (3, 'c.jpg', 'c.jpg', 'hash-c', 10, '2026-01-02T00:00:00Z', 'document', 0)",
304            [],
305        )
306        .unwrap();
307
308        let repo = DocumentRepo::new(&conn);
309        let categories = vec!["document".to_string(), "receipt".to_string()];
310        let page = repo.get_documents_by_categories(&categories, 2, 1).unwrap();
311
312        assert_eq!(page.iter().map(|p| p.id).collect::<Vec<_>>(), vec![3, 1]);
313    }
314}