kcode_k1_audio_classification_projection/
lib.rs

1use kcode_k1_audio_classification_format::decode_event;
2pub use kcode_k1_audio_classification_projection_reducer::ProjectionEffect;
3use kcode_k1_audio_classification_projection_reducer::{event_fragment_id, reduce};
4#[cfg(feature = "testkit")]
5use kcode_k1_audio_classification_projection_state::append_error;
6pub use kcode_k1_audio_classification_projection_state::{
7    ExecutedAnalysis, FragmentId, FragmentStageV1, FragmentStatus, LlmJobState, LlmJobStatus,
8    OverallState, SpeakerLabelV1, StageState, StageStatus,
9};
10use kcode_k1_audio_classification_projection_state::{
11    actionable_state, interrupted_stage, validate_status,
12};
13use kcode_k1_transaction::Transaction;
14use kcode_k1_txn_ordering::{K1TxnOrdering, TxId};
15use rusqlite::{Connection, Error as SqlError, ErrorCode, OptionalExtension, params};
16use std::ffi::OsString;
17use std::fmt::Display;
18use std::fs;
19use std::ops::Deref;
20use std::path::{Path, PathBuf};
21use std::sync::{Mutex, MutexGuard};
22
23const DATABASE_NAME: &str = "audio-classification.sqlite3";
24const SUBSYSTEM: &str = "audio-classification";
25const SCHEMA_VERSION: i64 = 2;
26const METADATA_SQL: &str = "CREATE TABLE metadata(singleton INTEGER PRIMARY KEY CHECK(singleton = 1), schema_version INTEGER NOT NULL, last_applied_txid BLOB)";
27const FRAGMENTS_SQL: &str = "CREATE TABLE fragments(fragment_id BLOB PRIMARY KEY, actionable_state INTEGER NOT NULL, encoded_status BLOB NOT NULL)";
28const INDEX_SQL: &str = "CREATE INDEX fragments_actionable_state ON fragments(actionable_state)";
29
30#[derive(Clone, Debug, PartialEq, Eq)]
31pub struct AppliedEvent {
32    pub fragment_id: FragmentId,
33    pub effect: ProjectionEffect,
34}
35
36#[derive(Clone, Debug, PartialEq, Eq)]
37pub struct InterruptedFragment {
38    pub fragment_id: FragmentId,
39    pub stage: FragmentStageV1,
40}
41
42pub struct Projection {
43    connection: Mutex<Connection>,
44}
45
46impl Projection {
47    pub fn open(root: &Path, ordering: &K1TxnOrdering) -> Result<(Self, Option<TxId>), String> {
48        fs::create_dir_all(root).map_err(text)?;
49        let path = root.join(DATABASE_NAME);
50        let exists = path.try_exists().map_err(text)?;
51        let (connection, cursor) = if exists {
52            match open_existing(&path) {
53                Ok(value) => value,
54                Err(OpenIssue::Recoverable) => recreate(&path)?,
55                Err(OpenIssue::Fatal(error)) => return Err(error),
56            }
57        } else {
58            remove_sidecars(&path)?;
59            create_database(&path)?
60        };
61        if let Some(cursor) = cursor {
62            let bytes = ordering.get_txn(cursor).map_err(text)?;
63            let valid = bytes.as_deref().is_some_and(|bytes| {
64                Transaction::parse(bytes).is_ok_and(|value| value.subsystem().as_str() == SUBSYSTEM)
65            });
66            if !valid {
67                drop(connection);
68                let (connection, cursor) = recreate(&path)?;
69                return Ok((
70                    Self {
71                        connection: Mutex::new(connection),
72                    },
73                    cursor,
74                ));
75            }
76        }
77        Ok((
78            Self {
79                connection: Mutex::new(connection),
80            },
81            cursor,
82        ))
83    }
84
85    pub fn apply(&self, callback_txid: TxId, payload: &[u8]) -> Result<AppliedEvent, String> {
86        let event = decode_event(payload).map_err(text)?;
87        let fragment_id = event_fragment_id(&event);
88        let mut connection = self.lock()?;
89        let transaction = connection.transaction().map_err(text)?;
90        let reduction = reduce(
91            load_status(&transaction, fragment_id)?,
92            &event,
93            callback_txid,
94        )?;
95        write_status(&transaction, fragment_id, &reduction.status)?;
96        let updated = transaction
97            .execute(
98                "UPDATE metadata SET last_applied_txid = ?1 WHERE singleton = 1",
99                params![&callback_txid.as_bytes()[..]],
100            )
101            .map_err(text)?;
102        if updated != 1 {
103            return Err("projection metadata row is missing".to_string());
104        }
105        transaction.commit().map_err(text)?;
106        Ok(AppliedEvent {
107            fragment_id,
108            effect: reduction.effect,
109        })
110    }
111
112    pub fn status(&self, fragment_id: FragmentId) -> Result<Option<FragmentStatus>, String> {
113        load_status(&self.lock()?, fragment_id)
114    }
115
116    pub fn queued(&self) -> Result<Vec<FragmentId>, String> {
117        Ok(actionable_statuses(&self.lock()?, 1)?
118            .into_iter()
119            .map(|value| value.0)
120            .collect())
121    }
122
123    pub fn running(&self) -> Result<Vec<InterruptedFragment>, String> {
124        Ok(actionable_statuses(&self.lock()?, 2)?
125            .into_iter()
126            .map(|(fragment_id, status)| InterruptedFragment {
127                fragment_id,
128                stage: interrupted_stage(&status),
129            })
130            .collect())
131    }
132
133    pub fn validate_labels(
134        &self,
135        fragment_id: FragmentId,
136        labels: &[SpeakerLabelV1],
137    ) -> Result<TxId, String> {
138        let status = self.status(fragment_id)?.ok_or("unknown fragment")?;
139        kcode_k1_audio_classification_projection_state::validate_labels(&status, labels)
140    }
141
142    pub fn clear(&self) -> Result<(), String> {
143        let mut connection = self.lock()?;
144        let transaction = connection.transaction().map_err(text)?;
145        transaction
146            .execute("DELETE FROM fragments", [])
147            .map_err(text)?;
148        transaction
149            .execute(
150                "UPDATE metadata SET last_applied_txid = NULL WHERE singleton = 1",
151                [],
152            )
153            .map_err(text)?;
154        transaction.commit().map_err(text)
155    }
156
157    #[cfg(feature = "testkit")]
158    pub fn inject_errors(
159        &self,
160        fragment_id: FragmentId,
161        errors: Vec<String>,
162    ) -> Result<(), String> {
163        let mut connection = self.lock()?;
164        let transaction = connection.transaction().map_err(text)?;
165        let mut status = load_status(&transaction, fragment_id)?.ok_or("unknown fragment")?;
166        for error in errors {
167            append_error(&mut status, error);
168        }
169        write_status(&transaction, fragment_id, &status)?;
170        transaction.commit().map_err(text)
171    }
172
173    fn lock(&self) -> Result<MutexGuard<'_, Connection>, String> {
174        self.connection.lock().map_err(text)
175    }
176}
177
178fn load_status<C: Deref<Target = Connection>>(
179    connection: &C,
180    fragment_id: FragmentId,
181) -> Result<Option<FragmentStatus>, String> {
182    let stored = connection
183        .query_row(
184            "SELECT actionable_state, encoded_status FROM fragments WHERE fragment_id = ?1",
185            params![&fragment_id.as_bytes()[..]],
186            |row| Ok((row.get::<_, i64>(0)?, row.get::<_, Vec<u8>>(1)?)),
187        )
188        .optional()
189        .map_err(text)?;
190    stored
191        .map(|(actionable, bytes)| decode_status(&bytes, actionable))
192        .transpose()
193}
194
195fn write_status(
196    connection: &Connection,
197    fragment_id: FragmentId,
198    status: &FragmentStatus,
199) -> Result<(), String> {
200    let encoded = postcard::to_allocvec(status).map_err(text)?;
201    connection.execute(
202        "INSERT INTO fragments(fragment_id, actionable_state, encoded_status) VALUES(?1, ?2, ?3) ON CONFLICT(fragment_id) DO UPDATE SET actionable_state = excluded.actionable_state, encoded_status = excluded.encoded_status",
203        params![&fragment_id.as_bytes()[..], actionable_state(status.state), encoded],
204    ).map_err(text)?;
205    Ok(())
206}
207
208fn actionable_statuses<C: Deref<Target = Connection>>(
209    connection: &C,
210    actionable: i64,
211) -> Result<Vec<(FragmentId, FragmentStatus)>, String> {
212    let mut statement = connection
213        .prepare("SELECT fragment_id, encoded_status FROM fragments WHERE actionable_state = ?1")
214        .map_err(text)?;
215    let rows = statement
216        .query_map(params![actionable], |row| {
217            Ok((row.get::<_, Vec<u8>>(0)?, row.get::<_, Vec<u8>>(1)?))
218        })
219        .map_err(text)?;
220    rows.map(|row| {
221        let (id, encoded) = row.map_err(text)?;
222        Ok((fragment_id(&id)?, decode_status(&encoded, actionable)?))
223    })
224    .collect()
225}
226
227fn fragment_id(bytes: &[u8]) -> Result<FragmentId, String> {
228    let bytes: [u8; 12] = bytes.try_into().map_err(text)?;
229    Ok(FragmentId::from_bytes(bytes))
230}
231
232fn decode_status(bytes: &[u8], actionable: i64) -> Result<FragmentStatus, String> {
233    let status: FragmentStatus = postcard::from_bytes(bytes).map_err(text)?;
234    if postcard::to_allocvec(&status).map_err(text)? != bytes {
235        return Err("stored fragment status is noncanonical".to_string());
236    }
237    validate_status(&status, actionable)?;
238    Ok(status)
239}
240
241fn create_database(path: &Path) -> Result<(Connection, Option<TxId>), String> {
242    let connection = Connection::open(path).map_err(text)?;
243    configure(&connection).map_err(text)?;
244    connection.execute_batch(&format!(
245        "{METADATA_SQL};{FRAGMENTS_SQL};{INDEX_SQL};INSERT INTO metadata(singleton, schema_version, last_applied_txid) VALUES(1, {SCHEMA_VERSION}, NULL);"
246    )).map_err(text)?;
247    Ok((connection, None))
248}
249
250fn open_existing(path: &Path) -> Result<(Connection, Option<TxId>), OpenIssue> {
251    let connection = Connection::open(path).map_err(classify)?;
252    configure(&connection).map_err(classify)?;
253    let check: String = connection
254        .query_row("PRAGMA quick_check", [], |row| row.get(0))
255        .map_err(classify)?;
256    if check != "ok" {
257        return Err(OpenIssue::Recoverable);
258    }
259    validate_schema(&connection)?;
260    let cursor = validate_rows(&connection)?;
261    Ok((connection, cursor))
262}
263
264fn configure(connection: &Connection) -> Result<(), SqlError> {
265    connection.pragma_update(None, "journal_mode", "WAL")?;
266    connection.pragma_update(None, "synchronous", "FULL")
267}
268
269fn validate_schema(connection: &Connection) -> Result<(), OpenIssue> {
270    let mut statement = connection.prepare(
271        "SELECT type, name, sql FROM sqlite_schema WHERE name NOT LIKE 'sqlite_%' ORDER BY type, name",
272    ).map_err(classify)?;
273    let rows = statement
274        .query_map([], |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)))
275        .map_err(classify)?;
276    let schema = rows
277        .collect::<Result<Vec<(String, String, String)>, _>>()
278        .map_err(classify)?;
279    let expected = vec![
280        (
281            "index".into(),
282            "fragments_actionable_state".into(),
283            INDEX_SQL.into(),
284        ),
285        ("table".into(), "fragments".into(), FRAGMENTS_SQL.into()),
286        ("table".into(), "metadata".into(), METADATA_SQL.into()),
287    ];
288    if schema != expected {
289        return Err(OpenIssue::Recoverable);
290    }
291    Ok(())
292}
293
294fn validate_rows(connection: &Connection) -> Result<Option<TxId>, OpenIssue> {
295    let mut metadata = connection
296        .prepare("SELECT singleton, schema_version, last_applied_txid FROM metadata")
297        .map_err(classify)?;
298    let rows = metadata
299        .query_map([], |row| {
300            Ok((
301                row.get::<_, i64>(0)?,
302                row.get::<_, i64>(1)?,
303                row.get::<_, Option<Vec<u8>>>(2)?,
304            ))
305        })
306        .map_err(classify)?;
307    let rows = rows.collect::<Result<Vec<_>, _>>().map_err(classify)?;
308    if rows.len() != 1 || rows[0].0 != 1 || rows[0].1 != SCHEMA_VERSION {
309        return Err(OpenIssue::Recoverable);
310    }
311    let mut fragments = connection
312        .prepare("SELECT fragment_id, actionable_state, encoded_status FROM fragments")
313        .map_err(classify)?;
314    let values = fragments
315        .query_map([], |row| {
316            Ok((
317                row.get::<_, Vec<u8>>(0)?,
318                row.get::<_, i64>(1)?,
319                row.get::<_, Vec<u8>>(2)?,
320            ))
321        })
322        .map_err(classify)?;
323    for value in values {
324        let (id, actionable, encoded) = value.map_err(classify)?;
325        fragment_id(&id).map_err(|_| OpenIssue::Recoverable)?;
326        decode_status(&encoded, actionable).map_err(|_| OpenIssue::Recoverable)?;
327    }
328    rows[0]
329        .2
330        .as_deref()
331        .map(|bytes| {
332            let bytes: [u8; 12] = bytes.try_into().map_err(|_| OpenIssue::Recoverable)?;
333            Ok(TxId::from_bytes(bytes))
334        })
335        .transpose()
336}
337
338fn classify(error: SqlError) -> OpenIssue {
339    if matches!(
340        error.sqlite_error_code(),
341        Some(ErrorCode::DatabaseCorrupt | ErrorCode::NotADatabase)
342    ) {
343        OpenIssue::Recoverable
344    } else {
345        OpenIssue::Fatal(error.to_string())
346    }
347}
348
349fn recreate(path: &Path) -> Result<(Connection, Option<TxId>), String> {
350    remove_sidecars(path)?;
351    remove_if_present(path)?;
352    create_database(path)
353}
354
355fn remove_sidecars(path: &Path) -> Result<(), String> {
356    for suffix in ["-wal", "-shm"] {
357        remove_if_present(&path_with_suffix(path, suffix))?;
358    }
359    Ok(())
360}
361
362fn remove_if_present(path: &Path) -> Result<(), String> {
363    match fs::remove_file(path) {
364        Ok(()) => Ok(()),
365        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
366        Err(error) => Err(error.to_string()),
367    }
368}
369
370fn path_with_suffix(path: &Path, suffix: &str) -> PathBuf {
371    let mut value = OsString::from(path.as_os_str());
372    value.push(suffix);
373    PathBuf::from(value)
374}
375
376fn text(error: impl Display) -> String {
377    error.to_string()
378}
379
380enum OpenIssue {
381    Recoverable,
382    Fatal(String),
383}
384
385#[cfg(test)]
386mod tests;
kcode_k1_audio_classification_projection/lib.rs

kcode_k1_audio_classification_projection/
lib.rs
