use crate::model::{
Cohort, DeleteOutcome, Error, FeatureRow, IdentifyOutcome, ObservationKey, TrainOutcome,
};
use crate::scoring::{self, LabeledRow};
use rusqlite::{Connection, OptionalExtension, Transaction, TransactionBehavior, params};
use std::collections::BTreeMap;
use std::path::Path;
use std::sync::{Mutex, MutexGuard};
use std::time::Duration;
const DATABASE_SCHEMA_VERSION: i64 = 1;
const BUSY_TIMEOUT: Duration = Duration::from_secs(5);
pub(crate) struct Store {
connection: Mutex<Connection>,
}
#[derive(Clone, Debug, PartialEq)]
struct StoredObservation {
cohort: Cohort,
row: FeatureRow,
speaker_id: String,
}
impl Store {
pub(crate) fn open(path: impl AsRef<Path>) -> Result<Self, Error> {
let mut connection = Connection::open(path)?;
connection.busy_timeout(BUSY_TIMEOUT)?;
let journal_mode: String =
connection.query_row("PRAGMA journal_mode = WAL", [], |row| row.get(0))?;
if !journal_mode.eq_ignore_ascii_case("wal") {
return Err(Error::Storage(format!(
"SQLite refused WAL mode and selected {journal_mode}"
)));
}
connection.pragma_update(None, "synchronous", "FULL")?;
connection.pragma_update(None, "foreign_keys", "ON")?;
initialize_or_check_schema(&mut connection)?;
rebuild_current(&mut connection)?;
Ok(Self {
connection: Mutex::new(connection),
})
}
pub(crate) fn identify(
&self,
key: &ObservationKey,
cohort: &Cohort,
row: &FeatureRow,
threshold: f64,
) -> Result<IdentifyOutcome, Error> {
let mut connection = self.lock()?;
let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?;
let existing = find_current(&transaction, key)?;
if existing
.as_ref()
.is_some_and(|stored| stored.cohort != *cohort || stored.row != *row)
{
return Err(Error::conflict(key));
}
let observations = load_cohort(&transaction, cohort, existing.as_ref().map(|_| key))?;
let evidence = scoring::score(&observations, row)?;
let speaker_id = evidence.as_ref().and_then(|value| {
(value.confidence_score >= threshold).then(|| value.best.speaker_id.clone())
});
if let Some(stored) = existing {
if speaker_id.as_deref() != Some(stored.speaker_id.as_str()) {
return Err(Error::conflict(key));
}
} else if let Some(speaker_id) = &speaker_id {
append_add(&transaction, key, cohort, row, speaker_id)?;
}
transaction.commit()?;
Ok(IdentifyOutcome {
speaker_id,
evidence,
})
}
pub(crate) fn train(
&self,
key: &ObservationKey,
cohort: &Cohort,
row: &FeatureRow,
speaker_id: &str,
) -> Result<TrainOutcome, Error> {
let mut connection = self.lock()?;
let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?;
let existing = find_current(&transaction, key)?;
let outcome = match existing {
None => {
append_add(&transaction, key, cohort, row, speaker_id)?;
TrainOutcome::Added
}
Some(stored)
if stored.cohort == *cohort
&& stored.row == *row
&& stored.speaker_id == speaker_id =>
{
TrainOutcome::Unchanged
}
Some(_) => {
append_delete(&transaction, key)?;
append_add(&transaction, key, cohort, row, speaker_id)?;
TrainOutcome::Corrected
}
};
transaction.commit()?;
Ok(outcome)
}
pub(crate) fn delete(&self, key: &ObservationKey) -> Result<DeleteOutcome, Error> {
let mut connection = self.lock()?;
let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?;
let outcome = if find_current(&transaction, key)?.is_some() {
append_delete(&transaction, key)?;
DeleteOutcome::Deleted
} else {
DeleteOutcome::NotFound
};
transaction.commit()?;
Ok(outcome)
}
fn lock(&self) -> Result<MutexGuard<'_, Connection>, Error> {
self.connection
.lock()
.map_err(|_| Error::Storage("database mutex was poisoned".to_owned()))
}
}
fn initialize_or_check_schema(connection: &mut Connection) -> Result<(), Error> {
let version: i64 = connection.pragma_query_value(None, "user_version", |row| row.get(0))?;
match version {
DATABASE_SCHEMA_VERSION => Ok(()),
0 => {
let has_user_tables: bool = connection.query_row(
"SELECT EXISTS(
SELECT 1 FROM sqlite_master
WHERE type = 'table' AND name NOT LIKE 'sqlite_%'
)",
[],
|row| row.get(0),
)?;
if has_user_tables {
return Err(Error::corrupt(
"unversioned database already contains application tables",
));
}
let transaction = connection.transaction()?;
transaction.execute_batch(
"CREATE TABLE events (
sequence INTEGER PRIMARY KEY AUTOINCREMENT,
kind TEXT NOT NULL CHECK (kind IN ('ADD', 'DELETE')),
object_id TEXT NOT NULL,
piece_index INTEGER NOT NULL,
provider TEXT,
model TEXT,
prompt_version TEXT,
feature_schema_version TEXT,
primary_language TEXT,
row_json TEXT,
speaker_id TEXT,
CHECK (
(kind = 'ADD'
AND provider IS NOT NULL
AND model IS NOT NULL
AND prompt_version IS NOT NULL
AND feature_schema_version IS NOT NULL
AND primary_language IS NOT NULL
AND row_json IS NOT NULL
AND speaker_id IS NOT NULL)
OR
(kind = 'DELETE'
AND provider IS NULL
AND model IS NULL
AND prompt_version IS NULL
AND feature_schema_version IS NULL
AND primary_language IS NULL
AND row_json IS NULL
AND speaker_id IS NULL)
)
);
CREATE TRIGGER events_no_update
BEFORE UPDATE ON events
BEGIN
SELECT RAISE(ABORT, 'events is append-only');
END;
CREATE TRIGGER events_no_delete
BEFORE DELETE ON events
BEGIN
SELECT RAISE(ABORT, 'events is append-only');
END;
CREATE TABLE current_observations (
object_id TEXT NOT NULL,
piece_index INTEGER NOT NULL,
provider TEXT NOT NULL,
model TEXT NOT NULL,
prompt_version TEXT NOT NULL,
feature_schema_version TEXT NOT NULL,
primary_language TEXT NOT NULL,
row_json TEXT NOT NULL,
speaker_id TEXT NOT NULL,
PRIMARY KEY (object_id, piece_index)
);
CREATE INDEX current_observations_cohort
ON current_observations (
provider,
model,
prompt_version,
feature_schema_version,
primary_language
);",
)?;
transaction.pragma_update(None, "user_version", DATABASE_SCHEMA_VERSION)?;
transaction.commit()?;
Ok(())
}
found => {
let found = u32::try_from(found).unwrap_or(u32::MAX);
Err(Error::UnsupportedSchema { found })
}
}
}
fn rebuild_current(connection: &mut Connection) -> Result<(), Error> {
let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?;
let state = {
let mut statement = transaction.prepare(
"SELECT
sequence,
kind,
object_id,
piece_index,
provider,
model,
prompt_version,
feature_schema_version,
primary_language,
row_json,
speaker_id
FROM events
ORDER BY sequence",
)?;
let mut rows = statement.query([])?;
let mut state = BTreeMap::<(String, u32), StoredObservation>::new();
let mut previous_sequence = None;
while let Some(record) = rows.next()? {
let sequence: i64 = record.get(0)?;
if sequence <= 0 || previous_sequence.is_some_and(|previous| sequence <= previous) {
return Err(Error::corrupt("event sequence is not strictly ordered"));
}
previous_sequence = Some(sequence);
let kind: String = record.get(1)?;
let object_id: String = record.get(2)?;
let piece_index = read_piece_index(record.get::<_, i64>(3)?)?;
let key = ObservationKey {
object_id,
piece_index,
};
key.validate()
.map_err(|error| Error::corrupt(format!("invalid event key: {error}")))?;
let provider: Option<String> = record.get(4)?;
let model: Option<String> = record.get(5)?;
let prompt_version: Option<String> = record.get(6)?;
let schema_version: Option<String> = record.get(7)?;
let primary_language: Option<String> = record.get(8)?;
let row_json: Option<String> = record.get(9)?;
let speaker_id: Option<String> = record.get(10)?;
let map_key = (key.object_id.clone(), key.piece_index);
match (
kind.as_str(),
provider,
model,
prompt_version,
schema_version,
primary_language,
row_json,
speaker_id,
) {
(
"ADD",
Some(provider),
Some(model),
Some(prompt_version),
Some(schema_version),
Some(primary_language),
Some(row_json),
Some(speaker_id),
) => {
let cohort = Cohort {
provider,
model,
prompt_version,
schema_version,
primary_language,
};
let row = parse_row(&row_json)?;
validate_stored(&cohort, &row, &speaker_id)?;
if state
.insert(
map_key,
StoredObservation {
cohort,
row,
speaker_id,
},
)
.is_some()
{
return Err(Error::corrupt(format!(
"ADD event {sequence} overwrites a current observation"
)));
}
}
("DELETE", None, None, None, None, None, None, None) => {
if state.remove(&map_key).is_none() {
return Err(Error::corrupt(format!(
"DELETE event {sequence} has no current observation"
)));
}
}
_ => {
return Err(Error::corrupt(format!(
"event {sequence} has an invalid kind or payload"
)));
}
}
}
state
};
transaction.execute("DELETE FROM current_observations", [])?;
for ((object_id, piece_index), stored) in state {
let row_json = serialize_row(&stored.row)?;
transaction.execute(
"INSERT INTO current_observations (
object_id,
piece_index,
provider,
model,
prompt_version,
feature_schema_version,
primary_language,
row_json,
speaker_id
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
params![
object_id,
i64::from(piece_index),
stored.cohort.provider,
stored.cohort.model,
stored.cohort.prompt_version,
stored.cohort.schema_version,
stored.cohort.primary_language,
row_json,
stored.speaker_id,
],
)?;
}
transaction.commit()?;
Ok(())
}
fn find_current(
connection: &Connection,
key: &ObservationKey,
) -> Result<Option<StoredObservation>, Error> {
let stored = connection
.query_row(
"SELECT
provider,
model,
prompt_version,
feature_schema_version,
primary_language,
row_json,
speaker_id
FROM current_observations
WHERE object_id = ? AND piece_index = ?",
params![key.object_id, i64::from(key.piece_index)],
|record| {
Ok((
Cohort {
provider: record.get(0)?,
model: record.get(1)?,
prompt_version: record.get(2)?,
schema_version: record.get(3)?,
primary_language: record.get(4)?,
},
record.get::<_, String>(5)?,
record.get::<_, String>(6)?,
))
},
)
.optional()?;
stored
.map(|(cohort, row_json, speaker_id)| {
let row = parse_row(&row_json)?;
validate_stored(&cohort, &row, &speaker_id)?;
Ok(StoredObservation {
cohort,
row,
speaker_id,
})
})
.transpose()
}
fn load_cohort(
connection: &Connection,
cohort: &Cohort,
exclude: Option<&ObservationKey>,
) -> Result<Vec<LabeledRow>, Error> {
let mut statement = connection.prepare(
"SELECT object_id, piece_index, row_json, speaker_id
FROM current_observations
WHERE provider = ?
AND model = ?
AND prompt_version = ?
AND feature_schema_version = ?
AND primary_language = ?
ORDER BY object_id, piece_index",
)?;
let mut rows = statement.query(params![
cohort.provider,
cohort.model,
cohort.prompt_version,
cohort.schema_version,
cohort.primary_language,
])?;
let mut observations = Vec::new();
while let Some(record) = rows.next()? {
let object_id: String = record.get(0)?;
let piece_index = read_piece_index(record.get::<_, i64>(1)?)?;
if exclude.is_some_and(|key| key.object_id == object_id && key.piece_index == piece_index) {
continue;
}
let row_json: String = record.get(2)?;
let speaker_id: String = record.get(3)?;
let row = parse_row(&row_json)?;
validate_stored(cohort, &row, &speaker_id)?;
observations.push(LabeledRow { speaker_id, row });
}
Ok(observations)
}
fn append_add(
transaction: &Transaction<'_>,
key: &ObservationKey,
cohort: &Cohort,
row: &FeatureRow,
speaker_id: &str,
) -> Result<(), Error> {
let row_json = serialize_row(row)?;
transaction.execute(
"INSERT INTO events (
kind,
object_id,
piece_index,
provider,
model,
prompt_version,
feature_schema_version,
primary_language,
row_json,
speaker_id
) VALUES ('ADD', ?, ?, ?, ?, ?, ?, ?, ?, ?)",
params![
key.object_id,
i64::from(key.piece_index),
cohort.provider,
cohort.model,
cohort.prompt_version,
cohort.schema_version,
cohort.primary_language,
row_json,
speaker_id,
],
)?;
transaction.execute(
"INSERT INTO current_observations (
object_id,
piece_index,
provider,
model,
prompt_version,
feature_schema_version,
primary_language,
row_json,
speaker_id
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
params![
key.object_id,
i64::from(key.piece_index),
cohort.provider,
cohort.model,
cohort.prompt_version,
cohort.schema_version,
cohort.primary_language,
row_json,
speaker_id,
],
)?;
Ok(())
}
fn append_delete(transaction: &Transaction<'_>, key: &ObservationKey) -> Result<(), Error> {
transaction.execute(
"INSERT INTO events (kind, object_id, piece_index)
VALUES ('DELETE', ?, ?)",
params![key.object_id, i64::from(key.piece_index)],
)?;
let changed = transaction.execute(
"DELETE FROM current_observations
WHERE object_id = ? AND piece_index = ?",
params![key.object_id, i64::from(key.piece_index)],
)?;
if changed != 1 {
return Err(Error::corrupt(
"current observation disappeared during an atomic mutation",
));
}
Ok(())
}
fn serialize_row(row: &FeatureRow) -> Result<String, Error> {
serde_json::to_string(row)
.map_err(|error| Error::Storage(format!("could not serialize feature row: {error}")))
}
fn parse_row(json: &str) -> Result<FeatureRow, Error> {
serde_json::from_str(json)
.map_err(|error| Error::corrupt(format!("invalid stored feature row: {error}")))
}
fn validate_stored(cohort: &Cohort, row: &FeatureRow, speaker_id: &str) -> Result<(), Error> {
cohort
.validate()
.map_err(|error| Error::corrupt(format!("invalid stored cohort: {error}")))?;
row.validate()
.map_err(|error| Error::corrupt(format!("invalid stored feature row: {error}")))?;
if speaker_id.trim().is_empty() {
return Err(Error::corrupt("stored speaker_id is empty"));
}
Ok(())
}
fn read_piece_index(value: i64) -> Result<u32, Error> {
u32::try_from(value).map_err(|_| Error::corrupt("stored piece_index is outside u32 range"))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::model::tests::{sample_cohort, sample_row};
use std::fs;
use std::sync::atomic::{AtomicU64, Ordering};
static NEXT_PATH: AtomicU64 = AtomicU64::new(0);
fn database_path(label: &str) -> std::path::PathBuf {
std::env::temp_dir().join(format!(
"kcode-speech-classification-store-{}-{label}-{}.sqlite3",
std::process::id(),
NEXT_PATH.fetch_add(1, Ordering::Relaxed)
))
}
fn key() -> ObservationKey {
ObservationKey {
object_id: "object".to_owned(),
piece_index: 7,
}
}
fn event_count(path: &Path) -> i64 {
let connection = Connection::open(path).unwrap();
connection
.query_row("SELECT COUNT(*) FROM events", [], |row| row.get(0))
.unwrap()
}
fn remove_database(path: &Path) {
for candidate in [
path.to_path_buf(),
path.with_extension("sqlite3-wal"),
path.with_extension("sqlite3-shm"),
] {
let _ = fs::remove_file(candidate);
}
}
#[test]
fn train_correction_and_delete_have_exact_event_semantics() {
let path = database_path("events");
let store = Store::open(&path).unwrap();
let cohort = sample_cohort();
let row = sample_row();
assert_eq!(
store.train(&key(), &cohort, &row, "speaker-a").unwrap(),
TrainOutcome::Added
);
assert_eq!(event_count(&path), 1);
assert_eq!(
store.train(&key(), &cohort, &row, "speaker-a").unwrap(),
TrainOutcome::Unchanged
);
assert_eq!(event_count(&path), 1);
assert_eq!(
store.train(&key(), &cohort, &row, "speaker-b").unwrap(),
TrainOutcome::Corrected
);
assert_eq!(event_count(&path), 3);
assert_eq!(store.delete(&key()).unwrap(), DeleteOutcome::Deleted);
assert_eq!(event_count(&path), 4);
assert_eq!(store.delete(&key()).unwrap(), DeleteOutcome::NotFound);
assert_eq!(event_count(&path), 4);
drop(store);
remove_database(&path);
}
#[test]
fn open_rebuilds_current_state_from_authoritative_events() {
let path = database_path("replay");
let cohort = sample_cohort();
let row = sample_row();
{
let store = Store::open(&path).unwrap();
store.train(&key(), &cohort, &row, "speaker-a").unwrap();
}
{
let connection = Connection::open(&path).unwrap();
connection
.execute(
"UPDATE current_observations
SET speaker_id = 'wrong'
WHERE object_id = 'object' AND piece_index = 7",
[],
)
.unwrap();
}
let store = Store::open(&path).unwrap();
assert_eq!(
store.train(&key(), &cohort, &row, "speaker-a").unwrap(),
TrainOutcome::Unchanged
);
assert_eq!(event_count(&path), 1);
drop(store);
remove_database(&path);
}
#[test]
fn rejects_an_unknown_schema_version() {
let path = database_path("schema");
{
let connection = Connection::open(&path).unwrap();
connection.pragma_update(None, "user_version", 99).unwrap();
}
assert!(matches!(
Store::open(&path),
Err(Error::UnsupportedSchema { found: 99 })
));
remove_database(&path);
}
}