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}