Skip to main content

rivet/state/
file_log.rs

1use crate::error::Result;
2
3use super::{StateConn, StateStore, pg_sql};
4
5/// One row from `file_log` (formerly `file_manifest`; renamed in schema v8).
6#[derive(Debug, serde::Serialize)]
7#[allow(dead_code)]
8pub struct FileRecord {
9    pub run_id: String,
10    pub export_name: String,
11    pub file_name: String,
12    pub row_count: i64,
13    pub bytes: i64,
14    pub format: String,
15    pub compression: Option<String>,
16    pub created_at: String,
17}
18
19/// File log store — reads and writes `file_log`.
20///
21/// Historical note: this table was named `file_manifest` prior to schema v8.
22/// The name was reclaimed for the 0.7.0 cloud-output JSON manifest contract;
23/// the internal SQLite log was renamed to `file_log` to remove the overload.
24///
25/// Invariant I2 (Write Before Log) governs when `record_file` is called:
26/// only after a destination write succeeds.  Failed writes produce no log entry.
27/// Invariant I7 (File-Log Failure Is Non-Fatal) means callers use `let _ = record_file(...)`.
28impl StateStore {
29    #[allow(clippy::too_many_arguments)]
30    pub fn record_file(
31        &self,
32        run_id: &str,
33        export_name: &str,
34        file_name: &str,
35        row_count: i64,
36        bytes: i64,
37        format: &str,
38        compression: Option<&str>,
39    ) -> Result<()> {
40        let now = chrono::Utc::now().to_rfc3339();
41        let sql = "INSERT INTO file_log (run_id, export_name, file_name, row_count, bytes, format, compression, created_at)
42             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)";
43        match &self.conn {
44            StateConn::Sqlite(c) => {
45                c.execute(
46                    sql,
47                    rusqlite::params![
48                        run_id,
49                        export_name,
50                        file_name,
51                        row_count,
52                        bytes,
53                        format,
54                        compression,
55                        now
56                    ],
57                )?;
58            }
59            StateConn::Postgres(client) => {
60                let mut c = client.borrow_mut();
61                c.execute(
62                    &pg_sql(sql),
63                    &[
64                        &run_id,
65                        &export_name,
66                        &file_name,
67                        &row_count,
68                        &bytes,
69                        &format,
70                        &compression,
71                        &now,
72                    ],
73                )?;
74            }
75        }
76        Ok(())
77    }
78
79    pub fn get_files(&self, export_name: Option<&str>, limit: usize) -> Result<Vec<FileRecord>> {
80        let cols =
81            "run_id, export_name, file_name, row_count, bytes, format, compression, created_at";
82        let limit_i64 = limit as i64;
83        match &self.conn {
84            StateConn::Sqlite(c) => {
85                let (sql, params): (&str, Vec<Box<dyn rusqlite::types::ToSql>>) = if let Some(
86                    name,
87                ) = export_name
88                {
89                    (
90                        "SELECT run_id, export_name, file_name, row_count, bytes, format, compression, created_at \
91                             FROM file_log WHERE export_name = ?1 ORDER BY id DESC LIMIT ?2",
92                        vec![Box::new(name.to_string()), Box::new(limit_i64)],
93                    )
94                } else {
95                    (
96                        "SELECT run_id, export_name, file_name, row_count, bytes, format, compression, created_at \
97                             FROM file_log ORDER BY id DESC LIMIT ?1",
98                        vec![Box::new(limit_i64)],
99                    )
100                };
101                let mut stmt = c.prepare(sql)?;
102                let params_refs: Vec<&dyn rusqlite::types::ToSql> =
103                    params.iter().map(|p| p.as_ref()).collect();
104                let rows = stmt.query_map(params_refs.as_slice(), |row| {
105                    Ok(FileRecord {
106                        run_id: row.get(0)?,
107                        export_name: row.get(1)?,
108                        file_name: row.get(2)?,
109                        row_count: row.get(3)?,
110                        bytes: row.get(4)?,
111                        format: row.get(5)?,
112                        compression: row.get(6)?,
113                        created_at: row.get(7)?,
114                    })
115                })?;
116                rows.collect::<std::result::Result<Vec<_>, _>>()
117                    .map_err(Into::into)
118            }
119            StateConn::Postgres(client) => {
120                // Single borrow for the duration of this call; safe because all Postgres
121                // operations in StateStore are sequential (no re-entrant borrows).
122                let mut c = client.borrow_mut();
123                let rows = if let Some(name) = export_name {
124                    c.query(
125                        &format!("SELECT {} FROM file_log WHERE export_name = $1 ORDER BY id DESC LIMIT $2", cols),
126                        &[&name, &limit_i64],
127                    )?
128                } else {
129                    c.query(
130                        &format!("SELECT {} FROM file_log ORDER BY id DESC LIMIT $1", cols),
131                        &[&limit_i64],
132                    )?
133                };
134                Ok(rows
135                    .iter()
136                    .map(|row| FileRecord {
137                        run_id: row.get(0),
138                        export_name: row.get(1),
139                        file_name: row.get(2),
140                        row_count: row.get(3),
141                        bytes: row.get(4),
142                        format: row.get(5),
143                        compression: row.get(6),
144                        created_at: row.get(7),
145                    })
146                    .collect())
147            }
148        }
149    }
150
151    /// Every `file_log` row for one `run_id`, in write order. Unlike `chunk_task`
152    /// (one `file_name` per chunk, the FIRST rotation sibling only), file_log
153    /// records EVERY committed part — including all `max_file_size` rotation
154    /// siblings — with its real `bytes`, written per-part before `complete_chunk_
155    /// task`. This is the crash-consistent source the chunked resume reconstructs a
156    /// COMPLETE destination manifest from when none exists (round-5 fix).
157    pub fn list_files_for_run(&self, run_id: &str) -> Result<Vec<FileRecord>> {
158        let cols =
159            "run_id, export_name, file_name, row_count, bytes, format, compression, created_at";
160        match &self.conn {
161            StateConn::Sqlite(c) => {
162                let mut stmt = c.prepare(&format!(
163                    "SELECT {cols} FROM file_log WHERE run_id = ?1 ORDER BY id ASC"
164                ))?;
165                let rows = stmt.query_map([run_id], |row| {
166                    Ok(FileRecord {
167                        run_id: row.get(0)?,
168                        export_name: row.get(1)?,
169                        file_name: row.get(2)?,
170                        row_count: row.get(3)?,
171                        bytes: row.get(4)?,
172                        format: row.get(5)?,
173                        compression: row.get(6)?,
174                        created_at: row.get(7)?,
175                    })
176                })?;
177                rows.collect::<std::result::Result<Vec<_>, _>>()
178                    .map_err(Into::into)
179            }
180            StateConn::Postgres(client) => {
181                let mut c = client.borrow_mut();
182                let rows = c.query(
183                    &format!("SELECT {cols} FROM file_log WHERE run_id = $1 ORDER BY id ASC"),
184                    &[&run_id],
185                )?;
186                Ok(rows
187                    .iter()
188                    .map(|row| FileRecord {
189                        run_id: row.get(0),
190                        export_name: row.get(1),
191                        file_name: row.get(2),
192                        row_count: row.get(3),
193                        bytes: row.get(4),
194                        format: row.get(5),
195                        compression: row.get(6),
196                        created_at: row.get(7),
197                    })
198                    .collect())
199            }
200        }
201    }
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207
208    fn store() -> StateStore {
209        StateStore::open_in_memory().expect("in-memory store")
210    }
211
212    #[test]
213    fn record_and_query_files() {
214        let s = store();
215        s.record_file(
216            "run_001",
217            "orders",
218            "orders_20260329.parquet",
219            50000,
220            4096,
221            "parquet",
222            Some("zstd"),
223        )
224        .unwrap();
225        s.record_file(
226            "run_001",
227            "orders",
228            "orders_20260329_chunk1.parquet",
229            25000,
230            2048,
231            "parquet",
232            Some("zstd"),
233        )
234        .unwrap();
235        s.record_file(
236            "run_002",
237            "users",
238            "users_20260329.csv",
239            1000,
240            500,
241            "csv",
242            None,
243        )
244        .unwrap();
245
246        let files = s.get_files(Some("orders"), 10).unwrap();
247        assert_eq!(files.len(), 2);
248        assert_eq!(files[0].run_id, "run_001");
249        assert_eq!(files[0].row_count, 25000);
250
251        let all = s.get_files(None, 10).unwrap();
252        assert_eq!(all.len(), 3);
253    }
254
255    #[test]
256    fn files_limit_works() {
257        let s = store();
258        for i in 0..10 {
259            s.record_file(
260                &format!("r{}", i),
261                "t",
262                &format!("f{}.parquet", i),
263                i,
264                i * 100,
265                "parquet",
266                None,
267            )
268            .unwrap();
269        }
270        let files = s.get_files(Some("t"), 3).unwrap();
271        assert_eq!(files.len(), 3);
272    }
273}