1use crate::error::Result;
2
3use super::{StateConn, StateStore, pg_sql};
4
5#[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
19impl 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 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 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}