use anyhow::Result;
use rusqlite::Connection;
use crate::state::session::{self, SessionMode, SessionOwnerKind, SessionStateFile};
pub(crate) fn read(conn: &Connection) -> Result<Option<SessionStateFile>> {
let mut stmt = conn.prepare(
"SELECT schema_version, started_at_epoch_s, last_started_at_epoch_s,
start_count, session_id, mode, owner_kind, owner_id,
supervisor_id, lease_ttl_secs, last_heartbeat_at_epoch_s, revision
FROM session WHERE id = 1",
)?;
let result = stmt.query_row([], |row| {
let mode: String = row.get(5)?;
let mode = SessionMode::from_str(&mode).ok_or_else(|| {
rusqlite::Error::FromSqlConversionFailure(
5,
rusqlite::types::Type::Text,
Box::new(std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("invalid session mode `{mode}`"),
)),
)
})?;
let owner_kind = match row.get::<_, Option<String>>(6)? {
Some(owner_kind) => SessionOwnerKind::from_str(&owner_kind).ok_or_else(|| {
rusqlite::Error::FromSqlConversionFailure(
6,
rusqlite::types::Type::Text,
Box::new(std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("invalid session owner_kind `{owner_kind}`"),
)),
)
})?,
None => SessionOwnerKind::Interactive,
};
let session_id: Option<String> = row.get(4)?;
let owner_id = match row.get::<_, Option<String>>(7)? {
Some(owner_id) => Some(owner_id),
None if session_id.is_some() && owner_kind == SessionOwnerKind::Interactive => {
Some("interactive".to_owned())
}
None => None,
};
Ok(SessionStateFile {
schema_version: row.get(0)?,
started_at_epoch_s: row.get(1)?,
last_started_at_epoch_s: row.get(2)?,
start_count: row.get(3)?,
session_id,
mode,
owner_kind,
owner_id,
supervisor_id: row.get(8)?,
lease_ttl_secs: row.get(9)?,
last_heartbeat_at_epoch_s: row.get(10)?,
revision: row.get(11)?,
})
});
match result {
Ok(state) => Ok(Some(state)),
Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
Err(e) => Err(e.into()),
}
}
pub(crate) fn write(conn: &Connection, state: &SessionStateFile) -> Result<()> {
let owner_id = match state.owner_kind {
SessionOwnerKind::Interactive => state
.owner_id
.clone()
.or_else(|| state.session_id.as_ref().map(|_| "interactive".to_owned())),
SessionOwnerKind::RuntimeSupervisor | SessionOwnerKind::RuntimeWorker => Some(
state
.owner_id
.clone()
.ok_or_else(|| anyhow::anyhow!("autonomous session rows require owner_id"))?,
),
};
if state.owner_kind.lifecycle() == session::SessionLifecycle::Autonomous
&& (state.lease_ttl_secs.is_none() || state.last_heartbeat_at_epoch_s.is_none())
{
return Err(anyhow::anyhow!(
"autonomous session rows require lease_ttl_secs and last_heartbeat_at_epoch_s"
));
}
conn.execute(
"INSERT OR REPLACE INTO session
(id, schema_version, started_at_epoch_s, last_started_at_epoch_s,
start_count, session_id, mode, owner_kind, owner_id, supervisor_id,
lease_ttl_secs, last_heartbeat_at_epoch_s, revision)
VALUES (1, ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12)",
rusqlite::params![
state.schema_version,
state.started_at_epoch_s,
state.last_started_at_epoch_s,
state.start_count,
state.session_id,
state.mode.as_str(),
state.owner_kind.as_str(),
owner_id,
state.supervisor_id,
state.lease_ttl_secs,
state.last_heartbeat_at_epoch_s,
state.revision,
],
)?;
Ok(())
}
pub(crate) fn delete(conn: &Connection) -> Result<()> {
conn.execute("DELETE FROM session WHERE id = 1", [])?;
Ok(())
}
pub(crate) fn load_session_id(conn: &Connection, now_epoch_s: u64) -> Result<Option<String>> {
let Some(state) = read(conn)? else {
return Ok(None);
};
if session::is_stale(&state, now_epoch_s) {
return Ok(None);
}
Ok(state.session_id)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::db::schema;
use crate::state::session::{SessionMode, SessionOwnerKind};
fn test_conn() -> Connection {
let conn = Connection::open_in_memory().unwrap();
schema::initialize(&conn).unwrap();
conn
}
fn interactive_state(
schema_version: u32,
started_at_epoch_s: u64,
last_started_at_epoch_s: u64,
start_count: u32,
session_id: Option<&str>,
mode: SessionMode,
) -> SessionStateFile {
SessionStateFile {
schema_version,
started_at_epoch_s,
last_started_at_epoch_s,
start_count,
session_id: session_id.map(str::to_owned),
mode,
owner_kind: SessionOwnerKind::Interactive,
owner_id: session_id.map(|_| "interactive".to_owned()),
supervisor_id: None,
lease_ttl_secs: None,
last_heartbeat_at_epoch_s: None,
revision: u64::from(session_id.is_some()),
}
}
#[test]
fn read_returns_none_when_empty() {
let conn = test_conn();
assert!(read(&conn).unwrap().is_none());
}
#[test]
fn write_and_read_roundtrip() {
let conn = test_conn();
let state = interactive_state(
3,
1_000_000,
1_000_050,
2,
Some("ses_01ABC"),
SessionMode::Research,
);
write(&conn, &state).unwrap();
let loaded = read(&conn).unwrap().expect("should exist");
assert_eq!(loaded.schema_version, 3);
assert_eq!(loaded.started_at_epoch_s, 1_000_000);
assert_eq!(loaded.last_started_at_epoch_s, 1_000_050);
assert_eq!(loaded.start_count, 2);
assert_eq!(loaded.session_id.as_deref(), Some("ses_01ABC"));
assert_eq!(loaded.mode, crate::state::session::SessionMode::Research);
}
#[test]
fn write_without_session_id() {
let conn = test_conn();
let state = interactive_state(3, 1_000_000, 1_000_000, 1, None, SessionMode::General);
write(&conn, &state).unwrap();
let loaded = read(&conn).unwrap().expect("should exist");
assert!(loaded.session_id.is_none());
}
#[test]
fn delete_removes_session() {
let conn = test_conn();
let state = interactive_state(
3,
1_000_000,
1_000_000,
1,
Some("ses_DEL"),
SessionMode::General,
);
write(&conn, &state).unwrap();
delete(&conn).unwrap();
assert!(read(&conn).unwrap().is_none());
}
#[test]
fn load_session_id_returns_none_for_stale() {
let conn = test_conn();
let now = 1_000_000u64;
let stale_after_secs = 8 * 60 * 60;
let state = interactive_state(
3,
now - stale_after_secs - 200,
now - stale_after_secs - 100,
1,
Some("ses_STALE"),
SessionMode::General,
);
write(&conn, &state).unwrap();
assert!(load_session_id(&conn, now).unwrap().is_none());
}
#[test]
fn load_session_id_returns_id_for_active() {
let conn = test_conn();
let now = 1_000_000u64;
let state = interactive_state(
3,
now - 100,
now - 50,
1,
Some("ses_ACTIVE"),
SessionMode::General,
);
write(&conn, &state).unwrap();
assert_eq!(
load_session_id(&conn, now).unwrap().as_deref(),
Some("ses_ACTIVE")
);
}
}