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::{CachedStatement, Connection, Error, OpenFlags, Result, Row, blob::Blob};
36
37use crate::error::table::{TableConnectError, TableError};
38
39// MARK: Traits
40/// Defines behavior for SQL Table data
41pub trait Table: Sized {
42    /// Deserialize a single row into Self, returning a [`rusqlite::Result`]
43    fn from_row(row: &Row) -> Result<Self>;
44
45    /// Prepare SELECT * statement
46    fn get(db: &'_ Connection) -> Result<CachedStatement<'_>, TableError>;
47
48    /// Map a `rusqlite::Result<Self>` into our `TableError`
49    fn extract(item: Result<Result<Self, Error>, Error>) -> Result<Self, TableError> {
50        match item {
51            Ok(Ok(row)) => Ok(row),
52            Err(why) | Ok(Err(why)) => Err(TableError::QueryError(why)),
53        }
54    }
55
56    /// Process all rows from the table using a callback.
57    /// This is the most memory-efficient approach for large tables.
58    ///
59    /// Uses the default `Table` implementation to prepare the statement and query the rows.
60    ///
61    /// To execute custom queries, see the [`message`](crate::tables::messages::message) module docs for examples.
62    ///
63    /// # Example
64    ///
65    /// ```no_run
66    /// use imessage_database::{
67    ///    error::table::TableError,
68    ///    tables::{
69    ///        table::{get_connection, Table},
70    ///        handle::Handle,
71    ///    },
72    ///    util::dirs::default_db_path
73    /// };
74    ///
75    /// // Get a connection to the database
76    /// let db_path = default_db_path();
77    /// let db = get_connection(&db_path).unwrap();
78    ///
79    /// // Stream the Handle table, processing each row with a callback
80    /// Handle::stream(&db, |handle_result| {
81    ///     match handle_result {
82    ///         Ok(handle) => println!("Handle: {}", handle.id),
83    ///         Err(e) => eprintln!("Error: {:?}", e),
84    ///     }
85    ///     Ok::<(), TableError>(())
86    /// }).unwrap();
87    /// ```
88    fn stream<F, E>(db: &Connection, callback: F) -> Result<(), E>
89    where
90        E: From<TableError>,
91        F: FnMut(Result<Self, TableError>) -> Result<(), E>,
92    {
93        stream_table_callback::<Self, F, E>(db, callback)
94    }
95
96    /// Get a BLOB from the table
97    ///
98    /// # Arguments
99    ///
100    /// * `db` - The database connection
101    /// * `table` - The name of the table
102    /// * `column` - The name of the column containing the BLOB
103    /// * `rowid` - The row ID to retrieve the BLOB from
104    fn get_blob<'a>(
105        &self,
106        db: &'a Connection,
107        table: &str,
108        column: &str,
109        rowid: i64,
110    ) -> Option<Blob<'a>> {
111        db.blob_open(rusqlite::MAIN_DB, table, column, rowid, true)
112            .ok()
113    }
114
115    /// Check if a BLOB exists in the table
116    fn has_blob(&self, db: &Connection, table: &str, column: &str, rowid: i64) -> bool {
117        let sql = std::format!(
118            "SELECT ({column} IS NOT NULL) AS not_null
119         FROM {table}
120         WHERE rowid = ?1",
121        );
122
123        // This returns 1 for true, 0 for false.
124        db.query_row(&sql, [rowid], |row| row.get(0))
125            .ok()
126            .is_some_and(|v: i32| v != 0)
127    }
128}
129
130fn stream_table_callback<T, F, E>(db: &Connection, mut callback: F) -> Result<(), E>
131where
132    T: Table + Sized,
133    E: From<TableError>,
134    F: FnMut(Result<T, TableError>) -> Result<(), E>,
135{
136    let mut stmt = T::get(db).map_err(E::from)?;
137    let rows = stmt
138        .query_map([], |row| Ok(T::from_row(row)))
139        .map_err(TableError::from)
140        .map_err(E::from)?;
141
142    for row_result in rows {
143        let item_result = T::extract(row_result);
144        callback(item_result)?;
145    }
146    Ok(())
147}
148
149/// Defines behavior for table data that can be cached in memory
150pub trait Cacheable {
151    /// The key type for the cache `HashMap`
152    type K;
153    /// The value type for the cache `HashMap`
154    type V;
155    /// Caches the table data in a `HashMap`
156    fn cache(db: &Connection) -> Result<HashMap<Self::K, Self::V>, TableError>;
157}
158
159// MARK: Database
160/// Get a connection to the iMessage `SQLite` database
161/// # Example:
162///
163/// ```
164/// use imessage_database::{
165///     util::dirs::default_db_path,
166///     tables::table::get_connection
167/// };
168///
169/// let db_path = default_db_path();
170/// let connection = get_connection(&db_path);
171/// ```
172pub fn get_connection(path: &Path) -> Result<Connection, TableError> {
173    if path.exists() && path.is_file() {
174        return match Connection::open_with_flags(
175            path,
176            OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX,
177        ) {
178            Ok(res) => Ok(res),
179            Err(why) => Err(TableError::CannotConnect(TableConnectError::Permissions(
180                why,
181            ))),
182        };
183    }
184
185    // Path does not point to a file
186    if path.exists() && !path.is_file() {
187        return Err(TableError::CannotConnect(TableConnectError::NotAFile(
188            path.to_path_buf(),
189        )));
190    }
191
192    // File is missing
193    Err(TableError::CannotConnect(TableConnectError::DoesNotExist(
194        path.to_path_buf(),
195    )))
196}
197
198/// Get the size of the database on the disk
199/// # Example:
200///
201/// ```
202/// use imessage_database::{
203///     util::dirs::default_db_path,
204///     tables::table::get_db_size
205/// };
206///
207/// let db_path = default_db_path();
208/// let database_size_in_bytes = get_db_size(&db_path);
209/// ```
210pub fn get_db_size(path: &Path) -> Result<u64, TableError> {
211    Ok(metadata(path)?.len())
212}
213
214// MARK: Constants
215// Table Names
216/// Handle table name
217pub const HANDLE: &str = "handle";
218/// Message table name
219pub const MESSAGE: &str = "message";
220/// Chat table name
221pub const CHAT: &str = "chat";
222/// Attachment table name
223pub const ATTACHMENT: &str = "attachment";
224/// Chat to message join table name
225pub const CHAT_MESSAGE_JOIN: &str = "chat_message_join";
226/// Message to attachment join table name
227pub const MESSAGE_ATTACHMENT_JOIN: &str = "message_attachment_join";
228/// Chat to handle join table name
229pub const CHAT_HANDLE_JOIN: &str = "chat_handle_join";
230/// Recently deleted messages table
231pub const RECENTLY_DELETED: &str = "chat_recoverable_message_join";
232
233// Column names
234/// The payload data column contains `plist`-encoded app message data
235pub const MESSAGE_PAYLOAD: &str = "payload_data";
236/// The message summary info column contains `plist`-encoded edited message information
237pub const MESSAGE_SUMMARY_INFO: &str = "message_summary_info";
238/// The `attributedBody` column contains [`typedstream`](crate::util::typedstream)-encoded message body text with many other attributes
239pub const ATTRIBUTED_BODY: &str = "attributedBody";
240/// The sticker user info column contains `plist`-encoded metadata for sticker attachments
241pub const STICKER_USER_INFO: &str = "sticker_user_info";
242/// The attribution info contains `plist`-encoded metadata for sticker attachments
243pub const ATTRIBUTION_INFO: &str = "attribution_info";
244/// The properties column contains `plist`-encoded metadata for a chat
245pub const PROPERTIES: &str = "properties";
246
247// Default information
248/// Name used for messages sent by the database owner in a first-person context
249pub const ME: &str = "Me";
250/// Name used for messages sent by the database owner in a second-person context
251pub const YOU: &str = "You";
252/// Name used for contacts or chats where the name cannot be discovered
253pub const UNKNOWN: &str = "Unknown";
254/// Default location for the Messages database on macOS
255pub const DEFAULT_PATH_MACOS: &str = "Library/Messages/chat.db";
256/// Default location for the Messages database in an iOS backup
257pub const DEFAULT_PATH_IOS: &str = "3d/3d0d7e5fb2ce288813306e4d4636395e047a3d28";
258/// Chat name reserved for messages that do not belong to a chat in the table
259pub const ORPHANED: &str = "orphaned";
260/// Replacement text sent in Fitness.app messages
261pub const FITNESS_RECEIVER: &str = "$(kIMTranscriptPluginBreadcrumbTextReceiverIdentifier)";
262/// Name for attachments directory in exports
263pub const ATTACHMENTS_DIR: &str = "attachments";
264
265#[cfg(test)]
266mod tests {
267    use rusqlite::{CachedStatement, Connection, Result, Row};
268
269    use crate::error::table::TableError;
270
271    use super::Table;
272
273    struct TestRow(i64);
274
275    impl Table for TestRow {
276        fn from_row(row: &Row) -> Result<Self> {
277            Ok(Self(row.get(0)?))
278        }
279
280        fn get(db: &'_ Connection) -> Result<CachedStatement<'_>, TableError> {
281            Ok(db.prepare_cached("SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3")?)
282        }
283    }
284
285    #[derive(Debug)]
286    enum StreamError {
287        Table(TableError),
288        Stop,
289    }
290
291    impl From<TableError> for StreamError {
292        fn from(err: TableError) -> Self {
293            Self::Table(err)
294        }
295    }
296
297    #[test]
298    fn stream_propagates_callback_errors() {
299        let db = Connection::open_in_memory().unwrap();
300        let mut seen = vec![];
301
302        let result = TestRow::stream(&db, |row| {
303            let row = row.map_err(StreamError::from)?;
304            seen.push(row.0);
305            if row.0 == 2 {
306                return Err(StreamError::Stop);
307            }
308            Ok(())
309        });
310
311        assert!(matches!(result, Err(StreamError::Stop)));
312        assert_eq!(seen, vec![1, 2]);
313    }
314
315    #[test]
316    fn stream_converts_setup_errors() {
317        struct BrokenTable;
318
319        impl Table for BrokenTable {
320            fn from_row(_row: &Row) -> Result<Self> {
321                Ok(Self)
322            }
323
324            fn get(_db: &'_ Connection) -> Result<CachedStatement<'_>, TableError> {
325                Err(TableError::CannotRead(std::io::Error::other("boom")))
326            }
327        }
328
329        let db = Connection::open_in_memory().unwrap();
330        let result = BrokenTable::stream(&db, |_| Ok::<(), StreamError>(()));
331
332        assert!(matches!(
333            result,
334            Err(StreamError::Table(TableError::CannotRead(_)))
335        ));
336    }
337}
imessage_database/tables/table.rs

imessage_database/tables/
table.rs
