imessage_database/tables/
table.rs

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

imessage_database/tables/
table.rs
