use kcode_k1_audio_classification_format::decode_event;
pub use kcode_k1_audio_classification_projection_reducer::ProjectionEffect;
use kcode_k1_audio_classification_projection_reducer::{event_fragment_id, reduce};
#[cfg(feature = "testkit")]
use kcode_k1_audio_classification_projection_state::append_error;
pub use kcode_k1_audio_classification_projection_state::{
ExecutedAnalysis, FragmentId, FragmentStageV1, FragmentStatus, LlmJobState, LlmJobStatus,
OverallState, SpeakerLabelV1, StageState, StageStatus,
};
use kcode_k1_audio_classification_projection_state::{
actionable_state, interrupted_stage, validate_status,
};
use kcode_k1_transaction::Transaction;
use kcode_k1_txn_ordering::{K1TxnOrdering, TxId};
use rusqlite::{Connection, Error as SqlError, ErrorCode, OptionalExtension, params};
use std::ffi::OsString;
use std::fmt::Display;
use std::fs;
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::sync::{Mutex, MutexGuard};
const DATABASE_NAME: &str = "audio-classification.sqlite3";
const SUBSYSTEM: &str = "audio-classification";
const SCHEMA_VERSION: i64 = 2;
const METADATA_SQL: &str = "CREATE TABLE metadata(singleton INTEGER PRIMARY KEY CHECK(singleton = 1), schema_version INTEGER NOT NULL, last_applied_txid BLOB)";
const FRAGMENTS_SQL: &str = "CREATE TABLE fragments(fragment_id BLOB PRIMARY KEY, actionable_state INTEGER NOT NULL, encoded_status BLOB NOT NULL)";
const INDEX_SQL: &str = "CREATE INDEX fragments_actionable_state ON fragments(actionable_state)";
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AppliedEvent {
pub fragment_id: FragmentId,
pub effect: ProjectionEffect,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct InterruptedFragment {
pub fragment_id: FragmentId,
pub stage: FragmentStageV1,
}
pub struct Projection {
connection: Mutex<Connection>,
}
impl Projection {
pub fn open(root: &Path, ordering: &K1TxnOrdering) -> Result<(Self, Option<TxId>), String> {
fs::create_dir_all(root).map_err(text)?;
let path = root.join(DATABASE_NAME);
let exists = path.try_exists().map_err(text)?;
let (connection, cursor) = if exists {
match open_existing(&path) {
Ok(value) => value,
Err(OpenIssue::Recoverable) => recreate(&path)?,
Err(OpenIssue::Fatal(error)) => return Err(error),
}
} else {
remove_sidecars(&path)?;
create_database(&path)?
};
if let Some(cursor) = cursor {
let bytes = ordering.get_txn(cursor).map_err(text)?;
let valid = bytes.as_deref().is_some_and(|bytes| {
Transaction::parse(bytes).is_ok_and(|value| value.subsystem().as_str() == SUBSYSTEM)
});
if !valid {
drop(connection);
let (connection, cursor) = recreate(&path)?;
return Ok((
Self {
connection: Mutex::new(connection),
},
cursor,
));
}
}
Ok((
Self {
connection: Mutex::new(connection),
},
cursor,
))
}
pub fn apply(&self, callback_txid: TxId, payload: &[u8]) -> Result<AppliedEvent, String> {
let event = decode_event(payload).map_err(text)?;
let fragment_id = event_fragment_id(&event);
let mut connection = self.lock()?;
let transaction = connection.transaction().map_err(text)?;
let reduction = reduce(
load_status(&transaction, fragment_id)?,
&event,
callback_txid,
)?;
write_status(&transaction, fragment_id, &reduction.status)?;
let updated = transaction
.execute(
"UPDATE metadata SET last_applied_txid = ?1 WHERE singleton = 1",
params![&callback_txid.as_bytes()[..]],
)
.map_err(text)?;
if updated != 1 {
return Err("projection metadata row is missing".to_string());
}
transaction.commit().map_err(text)?;
Ok(AppliedEvent {
fragment_id,
effect: reduction.effect,
})
}
pub fn status(&self, fragment_id: FragmentId) -> Result<Option<FragmentStatus>, String> {
load_status(&self.lock()?, fragment_id)
}
pub fn queued(&self) -> Result<Vec<FragmentId>, String> {
Ok(actionable_statuses(&self.lock()?, 1)?
.into_iter()
.map(|value| value.0)
.collect())
}
pub fn running(&self) -> Result<Vec<InterruptedFragment>, String> {
Ok(actionable_statuses(&self.lock()?, 2)?
.into_iter()
.map(|(fragment_id, status)| InterruptedFragment {
fragment_id,
stage: interrupted_stage(&status),
})
.collect())
}
pub fn validate_labels(
&self,
fragment_id: FragmentId,
labels: &[SpeakerLabelV1],
) -> Result<TxId, String> {
let status = self.status(fragment_id)?.ok_or("unknown fragment")?;
kcode_k1_audio_classification_projection_state::validate_labels(&status, labels)
}
pub fn clear(&self) -> Result<(), String> {
let mut connection = self.lock()?;
let transaction = connection.transaction().map_err(text)?;
transaction
.execute("DELETE FROM fragments", [])
.map_err(text)?;
transaction
.execute(
"UPDATE metadata SET last_applied_txid = NULL WHERE singleton = 1",
[],
)
.map_err(text)?;
transaction.commit().map_err(text)
}
#[cfg(feature = "testkit")]
pub fn inject_errors(
&self,
fragment_id: FragmentId,
errors: Vec<String>,
) -> Result<(), String> {
let mut connection = self.lock()?;
let transaction = connection.transaction().map_err(text)?;
let mut status = load_status(&transaction, fragment_id)?.ok_or("unknown fragment")?;
for error in errors {
append_error(&mut status, error);
}
write_status(&transaction, fragment_id, &status)?;
transaction.commit().map_err(text)
}
fn lock(&self) -> Result<MutexGuard<'_, Connection>, String> {
self.connection.lock().map_err(text)
}
}
fn load_status<C: Deref<Target = Connection>>(
connection: &C,
fragment_id: FragmentId,
) -> Result<Option<FragmentStatus>, String> {
let stored = connection
.query_row(
"SELECT actionable_state, encoded_status FROM fragments WHERE fragment_id = ?1",
params![&fragment_id.as_bytes()[..]],
|row| Ok((row.get::<_, i64>(0)?, row.get::<_, Vec<u8>>(1)?)),
)
.optional()
.map_err(text)?;
stored
.map(|(actionable, bytes)| decode_status(&bytes, actionable))
.transpose()
}
fn write_status(
connection: &Connection,
fragment_id: FragmentId,
status: &FragmentStatus,
) -> Result<(), String> {
let encoded = postcard::to_allocvec(status).map_err(text)?;
connection.execute(
"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",
params![&fragment_id.as_bytes()[..], actionable_state(status.state), encoded],
).map_err(text)?;
Ok(())
}
fn actionable_statuses<C: Deref<Target = Connection>>(
connection: &C,
actionable: i64,
) -> Result<Vec<(FragmentId, FragmentStatus)>, String> {
let mut statement = connection
.prepare("SELECT fragment_id, encoded_status FROM fragments WHERE actionable_state = ?1")
.map_err(text)?;
let rows = statement
.query_map(params![actionable], |row| {
Ok((row.get::<_, Vec<u8>>(0)?, row.get::<_, Vec<u8>>(1)?))
})
.map_err(text)?;
rows.map(|row| {
let (id, encoded) = row.map_err(text)?;
Ok((fragment_id(&id)?, decode_status(&encoded, actionable)?))
})
.collect()
}
fn fragment_id(bytes: &[u8]) -> Result<FragmentId, String> {
let bytes: [u8; 12] = bytes.try_into().map_err(text)?;
Ok(FragmentId::from_bytes(bytes))
}
fn decode_status(bytes: &[u8], actionable: i64) -> Result<FragmentStatus, String> {
let status: FragmentStatus = postcard::from_bytes(bytes).map_err(text)?;
if postcard::to_allocvec(&status).map_err(text)? != bytes {
return Err("stored fragment status is noncanonical".to_string());
}
validate_status(&status, actionable)?;
Ok(status)
}
fn create_database(path: &Path) -> Result<(Connection, Option<TxId>), String> {
let connection = Connection::open(path).map_err(text)?;
configure(&connection).map_err(text)?;
connection.execute_batch(&format!(
"{METADATA_SQL};{FRAGMENTS_SQL};{INDEX_SQL};INSERT INTO metadata(singleton, schema_version, last_applied_txid) VALUES(1, {SCHEMA_VERSION}, NULL);"
)).map_err(text)?;
Ok((connection, None))
}
fn open_existing(path: &Path) -> Result<(Connection, Option<TxId>), OpenIssue> {
let connection = Connection::open(path).map_err(classify)?;
configure(&connection).map_err(classify)?;
let check: String = connection
.query_row("PRAGMA quick_check", [], |row| row.get(0))
.map_err(classify)?;
if check != "ok" {
return Err(OpenIssue::Recoverable);
}
validate_schema(&connection)?;
let cursor = validate_rows(&connection)?;
Ok((connection, cursor))
}
fn configure(connection: &Connection) -> Result<(), SqlError> {
connection.pragma_update(None, "journal_mode", "WAL")?;
connection.pragma_update(None, "synchronous", "FULL")
}
fn validate_schema(connection: &Connection) -> Result<(), OpenIssue> {
let mut statement = connection.prepare(
"SELECT type, name, sql FROM sqlite_schema WHERE name NOT LIKE 'sqlite_%' ORDER BY type, name",
).map_err(classify)?;
let rows = statement
.query_map([], |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)))
.map_err(classify)?;
let schema = rows
.collect::<Result<Vec<(String, String, String)>, _>>()
.map_err(classify)?;
let expected = vec![
(
"index".into(),
"fragments_actionable_state".into(),
INDEX_SQL.into(),
),
("table".into(), "fragments".into(), FRAGMENTS_SQL.into()),
("table".into(), "metadata".into(), METADATA_SQL.into()),
];
if schema != expected {
return Err(OpenIssue::Recoverable);
}
Ok(())
}
fn validate_rows(connection: &Connection) -> Result<Option<TxId>, OpenIssue> {
let mut metadata = connection
.prepare("SELECT singleton, schema_version, last_applied_txid FROM metadata")
.map_err(classify)?;
let rows = metadata
.query_map([], |row| {
Ok((
row.get::<_, i64>(0)?,
row.get::<_, i64>(1)?,
row.get::<_, Option<Vec<u8>>>(2)?,
))
})
.map_err(classify)?;
let rows = rows.collect::<Result<Vec<_>, _>>().map_err(classify)?;
if rows.len() != 1 || rows[0].0 != 1 || rows[0].1 != SCHEMA_VERSION {
return Err(OpenIssue::Recoverable);
}
let mut fragments = connection
.prepare("SELECT fragment_id, actionable_state, encoded_status FROM fragments")
.map_err(classify)?;
let values = fragments
.query_map([], |row| {
Ok((
row.get::<_, Vec<u8>>(0)?,
row.get::<_, i64>(1)?,
row.get::<_, Vec<u8>>(2)?,
))
})
.map_err(classify)?;
for value in values {
let (id, actionable, encoded) = value.map_err(classify)?;
fragment_id(&id).map_err(|_| OpenIssue::Recoverable)?;
decode_status(&encoded, actionable).map_err(|_| OpenIssue::Recoverable)?;
}
rows[0]
.2
.as_deref()
.map(|bytes| {
let bytes: [u8; 12] = bytes.try_into().map_err(|_| OpenIssue::Recoverable)?;
Ok(TxId::from_bytes(bytes))
})
.transpose()
}
fn classify(error: SqlError) -> OpenIssue {
if matches!(
error.sqlite_error_code(),
Some(ErrorCode::DatabaseCorrupt | ErrorCode::NotADatabase)
) {
OpenIssue::Recoverable
} else {
OpenIssue::Fatal(error.to_string())
}
}
fn recreate(path: &Path) -> Result<(Connection, Option<TxId>), String> {
remove_sidecars(path)?;
remove_if_present(path)?;
create_database(path)
}
fn remove_sidecars(path: &Path) -> Result<(), String> {
for suffix in ["-wal", "-shm"] {
remove_if_present(&path_with_suffix(path, suffix))?;
}
Ok(())
}
fn remove_if_present(path: &Path) -> Result<(), String> {
match fs::remove_file(path) {
Ok(()) => Ok(()),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(error) => Err(error.to_string()),
}
}
fn path_with_suffix(path: &Path, suffix: &str) -> PathBuf {
let mut value = OsString::from(path.as_os_str());
value.push(suffix);
PathBuf::from(value)
}
fn text(error: impl Display) -> String {
error.to_string()
}
enum OpenIssue {
Recoverable,
Fatal(String),
}
#[cfg(test)]
mod tests;