Skip to main content

imessage_database/tables/
table.rs

1/*!
2 Table traits, database connection helpers, and shared table constants.
3
4 # Streaming API
5
6 The streaming API processes each row through a callback without collecting the
7 table into a `Vec`.
8
9 ```no_run
10 use imessage_database::{
11    error::table::TableError,
12    tables::{
13        table::{get_connection, Table},
14        messages::Message,
15    },
16    util::dirs::default_db_path
17 };
18
19 let db_path = default_db_path();
20 let db = get_connection(&db_path).unwrap();
21
22 Message::stream(&db, |message_result| {
23     match message_result {
24         Ok(message) => println!("Message: {:#?}", message),
25         Err(e) => eprintln!("Error: {:?}", e),
26     }
27    Ok::<(), TableError>(())
28 }).unwrap();
29 ```
30
31 The callback may return any error type that implements `From<TableError>`.
32*/
33
34use std::{collections::HashMap, fs::metadata, path::Path};
35
36use rusqlite::{
37    CachedStatement, Connection, Error, OpenFlags, Params, Result, Row, Statement, blob::Blob,
38};
39
40use crate::error::table::{TableConnectError, TableError};
41
42// MARK: Traits
43/// Database table model that can deserialize itself from SQLite rows.
44pub trait Table: Sized {
45    /// Deserialize a single row into `Self`. Returns [`rusqlite::Result`]
46    /// for direct use inside `rusqlite::query_map` / `query_row`
47    /// callbacks. For high-level iteration, prefer [`Table::rows`] or
48    /// [`Table::row`].
49    fn from_row(row: &Row) -> Result<Self>;
50
51    /// Prepare the table's default `SELECT *` statement.
52    fn get(db: &'_ Connection) -> Result<CachedStatement<'_>, TableError>;
53
54    /// Iterate over rows produced by `stmt`. The default implementation
55    /// deserializes each through [`from_row`](Self::from_row). Row-fetch and
56    /// row-deserialization failures are surfaced uniformly as [`TableError`].
57    /// Accepts both [`rusqlite::Statement`] and [`rusqlite::CachedStatement`]
58    /// (the latter via deref coercion).
59    ///
60    /// Use this when the caller owns a custom prepared statement (with
61    /// filters, joins, or bound parameters). For a full-table scan against
62    /// the default `SELECT *` with a callback API, see [`Table::stream`].
63    ///
64    /// Implementations may override this method to resolve column layout once
65    /// and reuse it for every row. Resolve through the first [`Row`] when the
66    /// schema may change concurrently: `sqlite3_step` may recompile the
67    /// statement after preparation.
68    fn rows<'stmt, P: Params>(
69        stmt: &'stmt mut Statement<'_>,
70        params: P,
71    ) -> Result<impl Iterator<Item = Result<Self, TableError>> + 'stmt, TableError>
72    where
73        Self: 'stmt,
74    {
75        let mapped = stmt.query_map(params, |row| Ok(Self::from_row(row)))?;
76        Ok(mapped.map(flatten_row))
77    }
78
79    /// Fetch exactly one row from `stmt`. Returns
80    /// [`TableError::QueryError`] if the row is missing or fails to
81    /// deserialize. Accepts both [`rusqlite::Statement`] and
82    /// [`rusqlite::CachedStatement`] (the latter via deref coercion).
83    ///
84    /// Implementations may override this method to resolve the stepped row's
85    /// column layout instead of decoding by name.
86    fn row<P: Params>(stmt: &mut Statement<'_>, params: P) -> Result<Self, TableError> {
87        flatten_row(stmt.query_row(params, |row| Ok(Self::from_row(row))))
88    }
89
90    /// Process every row from the table's default `SELECT *` query using a
91    /// callback. Builds and discards the prepared statement internally, so
92    /// the caller never sees it.
93    ///
94    /// Use this for full-table scans where the callback style fits. For
95    /// custom statements (filters, joins, bound parameters), prepare the
96    /// statement yourself and iterate via [`Table::rows`]. See the
97    /// [`message`](crate::tables::messages::message) module docs for an
98    /// example.
99    ///
100    /// # Example
101    ///
102    /// ```no_run
103    /// use imessage_database::{
104    ///    error::table::TableError,
105    ///    tables::{
106    ///        table::{get_connection, Table},
107    ///        handle::Handle,
108    ///    },
109    ///    util::dirs::default_db_path
110    /// };
111    ///
112    /// let db_path = default_db_path();
113    /// let db = get_connection(&db_path).unwrap();
114    ///
115    /// // Stream the Handle table, processing each row with a callback
116    /// Handle::stream(&db, |handle_result| {
117    ///     match handle_result {
118    ///         Ok(handle) => println!("Handle: {}", handle.id),
119    ///         Err(e) => eprintln!("Error: {:?}", e),
120    ///     }
121    ///     Ok::<(), TableError>(())
122    /// }).unwrap();
123    /// ```
124    fn stream<F, E>(db: &Connection, callback: F) -> Result<(), E>
125    where
126        E: From<TableError>,
127        F: FnMut(Result<Self, TableError>) -> Result<(), E>,
128    {
129        stream_table_callback::<Self, F, E>(db, callback)
130    }
131
132    /// Open a `BLOB` column for the supplied `rowid`.
133    fn get_blob<'a>(
134        &self,
135        db: &'a Connection,
136        table: &str,
137        column: &str,
138        rowid: i64,
139    ) -> Option<Blob<'a>> {
140        db.blob_open(rusqlite::MAIN_DB, table, column, rowid, true)
141            .ok()
142    }
143
144    /// Return whether a `BLOB` column is non-null for the supplied `rowid`.
145    fn has_blob(&self, db: &Connection, table: &str, column: &str, rowid: i64) -> bool {
146        let sql = std::format!(
147            "SELECT ({column} IS NOT NULL) AS not_null
148         FROM {table}
149         WHERE rowid = ?1",
150        );
151
152        // This returns 1 for true, 0 for false.
153        db.query_row(&sql, [rowid], |row| row.get(0))
154            .ok()
155            .is_some_and(|v: i32| v != 0)
156    }
157}
158
159/// Flatten the doubly-nested result produced by `rusqlite::query_map` /
160/// `query_row` callbacks into a single [`TableError`]. The outer layer
161/// represents row-fetch failures, the inner layer represents row-deserialize
162/// failures from [`Table::from_row`].
163pub(crate) fn flatten_row<T>(item: Result<Result<T, Error>, Error>) -> Result<T, TableError> {
164    match item {
165        Ok(Ok(row)) => Ok(row),
166        Err(why) | Ok(Err(why)) => Err(TableError::QueryError(why)),
167    }
168}
169
170fn stream_table_callback<T, F, E>(db: &Connection, mut callback: F) -> Result<(), E>
171where
172    T: Table + Sized,
173    E: From<TableError>,
174    F: FnMut(Result<T, TableError>) -> Result<(), E>,
175{
176    let mut stmt = T::get(db).map_err(E::from)?;
177    for row_result in T::rows(&mut stmt, []).map_err(E::from)? {
178        callback(row_result)?;
179    }
180    Ok(())
181}
182
183/// Table data that can be materialized into an in-memory map.
184pub trait Cacheable {
185    /// Key type for the cache map.
186    type K;
187    /// Value type for the cache map.
188    type V;
189    /// Build the cache from the database.
190    fn cache(db: &Connection) -> Result<HashMap<Self::K, Self::V>, TableError>;
191}
192
193// MARK: Database
194/// Open the Messages `SQLite` database read-only.
195/// # Example:
196///
197/// ```
198/// use imessage_database::{
199///     util::dirs::default_db_path,
200///     tables::table::get_connection
201/// };
202///
203/// let db_path = default_db_path();
204/// let connection = get_connection(&db_path);
205/// ```
206pub fn get_connection(path: &Path) -> Result<Connection, TableError> {
207    if path.exists() && path.is_file() {
208        return match Connection::open_with_flags(
209            path,
210            OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX,
211        ) {
212            Ok(connection) => {
213                // Read pages from the mapped region where SQLite supports it.
214                let _ = connection.pragma_update(None, "mmap_size", 8_589_934_592_i64); // up to 8 GiB
215                let _ = connection.pragma_update(None, "cache_size", -65_536_i64); // ~64 MiB
216                Ok(connection)
217            }
218            Err(why) => Err(TableError::CannotConnect(TableConnectError::Permissions(
219                why,
220            ))),
221        };
222    }
223
224    // Path does not point to a file
225    if path.exists() && !path.is_file() {
226        return Err(TableError::CannotConnect(TableConnectError::NotAFile(
227            path.to_path_buf(),
228        )));
229    }
230
231    // File is missing
232    Err(TableError::CannotConnect(TableConnectError::DoesNotExist(
233        path.to_path_buf(),
234    )))
235}
236
237/// Return the database file size on disk.
238/// # Example:
239///
240/// ```
241/// use imessage_database::{
242///     util::dirs::default_db_path,
243///     tables::table::get_db_size
244/// };
245///
246/// let db_path = default_db_path();
247/// let database_size_in_bytes = get_db_size(&db_path);
248/// ```
249pub fn get_db_size(path: &Path) -> Result<u64, TableError> {
250    Ok(metadata(path)?.len())
251}
252
253// MARK: Constants
254// Table Names
255/// Handle table name.
256pub const HANDLE: &str = "handle";
257/// Message table name.
258pub const MESSAGE: &str = "message";
259/// Chat table name.
260pub const CHAT: &str = "chat";
261/// Attachment table name.
262pub const ATTACHMENT: &str = "attachment";
263/// Chat-to-message join table name.
264pub const CHAT_MESSAGE_JOIN: &str = "chat_message_join";
265/// Message-to-attachment join table name.
266pub const MESSAGE_ATTACHMENT_JOIN: &str = "message_attachment_join";
267/// Chat-to-handle join table name.
268pub const CHAT_HANDLE_JOIN: &str = "chat_handle_join";
269/// Recently deleted messages table.
270pub const RECENTLY_DELETED: &str = "chat_recoverable_message_join";
271
272// Column names
273/// [`plist`](crate::util::plist)-encoded app-message payload column.
274pub const MESSAGE_PAYLOAD: &str = "payload_data";
275/// [`plist`](crate::util::plist)-encoded message summary column.
276pub const MESSAGE_SUMMARY_INFO: &str = "message_summary_info";
277/// [`typedstream`](crate::util::typedstream)-encoded attributed body column.
278pub const ATTRIBUTED_BODY: &str = "attributedBody";
279/// [`plist`](crate::util::plist)-encoded sticker metadata column.
280pub const STICKER_USER_INFO: &str = "sticker_user_info";
281/// [`plist`](crate::util::plist)-encoded attachment attribution column.
282pub const ATTRIBUTION_INFO: &str = "attribution_info";
283/// [`plist`](crate::util::plist)-encoded chat properties column.
284pub const PROPERTIES: &str = "properties";
285
286// Default information
287/// First-person display name for the database owner.
288pub const ME: &str = "Me";
289/// Second-person display name for the database owner.
290pub const YOU: &str = "You";
291/// Display name used when a contact or chat name is unavailable.
292pub const UNKNOWN: &str = "Unknown";
293/// Default macOS Messages database path.
294pub const DEFAULT_PATH_MACOS: &str = "Library/Messages/chat.db";
295/// Default Messages database path inside an iOS backup.
296pub const DEFAULT_PATH_IOS: &str = "3d/3d0d7e5fb2ce288813306e4d4636395e047a3d28";
297/// Chat name reserved for messages that do not belong to a chat row.
298pub const ORPHANED: &str = "orphaned";
299/// Replacement token found in Fitness.app messages.
300pub const FITNESS_RECEIVER: &str = "$(kIMTranscriptPluginBreadcrumbTextReceiverIdentifier)";
301/// Attachments directory name used in exports.
302pub const ATTACHMENTS_DIR: &str = "attachments";
303
304#[cfg(test)]
305mod tests {
306    use rusqlite::{CachedStatement, Connection, Result, Row};
307
308    use crate::error::table::TableError;
309
310    use super::Table;
311
312    struct TestRow(i64);
313
314    impl Table for TestRow {
315        fn from_row(row: &Row) -> Result<Self> {
316            Ok(Self(row.get(0)?))
317        }
318
319        fn get(db: &'_ Connection) -> Result<CachedStatement<'_>, TableError> {
320            Ok(db.prepare_cached("SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3")?)
321        }
322    }
323
324    #[derive(Debug)]
325    enum StreamError {
326        Table(TableError),
327        Stop,
328    }
329
330    impl From<TableError> for StreamError {
331        fn from(err: TableError) -> Self {
332            Self::Table(err)
333        }
334    }
335
336    #[test]
337    fn stream_propagates_callback_errors() {
338        let db = Connection::open_in_memory().unwrap();
339        let mut seen = vec![];
340
341        let result = TestRow::stream(&db, |row| {
342            let row = row.map_err(StreamError::from)?;
343            seen.push(row.0);
344            if row.0 == 2 {
345                return Err(StreamError::Stop);
346            }
347            Ok(())
348        });
349
350        assert!(matches!(result, Err(StreamError::Stop)));
351        assert_eq!(seen, vec![1, 2]);
352    }
353
354    #[test]
355    fn stream_converts_setup_errors() {
356        struct BrokenTable;
357
358        impl Table for BrokenTable {
359            fn from_row(_row: &Row) -> Result<Self> {
360                Ok(Self)
361            }
362
363            fn get(_db: &'_ Connection) -> Result<CachedStatement<'_>, TableError> {
364                Err(TableError::CannotRead(std::io::Error::other("boom")))
365            }
366        }
367
368        let db = Connection::open_in_memory().unwrap();
369        let result = BrokenTable::stream(&db, |_| Ok::<(), StreamError>(()));
370
371        assert!(matches!(
372            result,
373            Err(StreamError::Table(TableError::CannotRead(_)))
374        ));
375    }
376}