mod maintenance; mod read; mod write; pub use maintenance::{FeedHealth, StorageHealth, StorageReadFailure}; use crate::provider::RawFeed; use anyhow::{ensure, Context}; use deadpool_sqlite::{Hook, HookError, Manager, Pool, Runtime}; use deadpool_sqlite::rusqlite::{functions::FunctionFlags, Error as SqliteError}; use std::{ collections::{HashMap, HashSet}, path::{Path, PathBuf}, sync::{Arc, Mutex}, time::Duration, }; const MIGRATIONS: [&str; 5] = [ include_str!("../migrations/0001_spatial_storage.sql"), include_str!("../migrations/0002_maintenance_health.sql"), include_str!("../migrations/0003_shared_place_alias.sql"), include_str!("../migrations/0004_sequence_only_stops.sql"), include_str!("../migrations/0005_realtime_external_ids.sql"), ]; const CURRENT_SCHEMA_VERSION: i64 = MIGRATIONS.len() as i64; #[derive(Clone)] pub struct Store { pool: Pool, path: Arc, last_read_failure: Arc>>, } pub struct StaticBatch { pub raw: Vec, pub stops: Vec, pub routes: HashMap, crate::model::RouteInfo>, pub sequences: HashMap, Vec>, pub shapes: crate::shapes::ShapeIndex, } pub struct WaitingBatch { pub raw: RawFeed, pub passages: HashMap, Vec>, pub end_notices: HashMap, Vec>, } pub struct PositionsBatch { pub raw: RawFeed, pub positions: HashMap, Vec>, } pub struct DisruptionsBatch { pub raw: RawFeed, pub disruptions: Vec, } pub struct VilloDiscoveryBatch { pub raw: Vec, pub status_url: String, pub stations: crate::villo::VilloInfo, } pub struct VilloStatusBatch { pub raw: RawFeed, pub status: crate::villo::VilloStatus, } pub struct TrainStationsBatch { pub raw: RawFeed, pub stations: Vec, } pub struct TrainDeparturesBatch { pub raw: RawFeed, pub station_id: String, pub departures: Vec, } #[derive(Clone, Debug, PartialEq)] pub struct TrainBoardData { pub stops: Vec, pub refresh_station_ids: Vec, } #[derive(Clone, Debug, PartialEq)] pub struct TrainStopData { pub id: String, pub name: String, pub distance_m: f64, pub departures: Vec, } pub struct TransitBoardData { pub data_age_seconds: Option, pub stale: bool, pub stops: Vec, } pub struct TransitStopData { pub id: String, pub name: crate::model::Bilingual, pub distance_m: u32, pub rank: u8, pub departures: Vec, } pub struct SearchHitData { pub id: String, pub name: crate::model::Bilingual, pub line_count: usize, pub direction: Option, } #[derive(Clone, Debug, PartialEq)] pub struct LineStripStopData { pub id: String, pub name: crate::model::Bilingual, pub you: bool, pub vehicle_after_m: Option, } #[derive(Clone, Debug, PartialEq)] pub struct LineDisruptionData { pub en: Option, pub fr: Option, pub nl: Option, } #[derive(Clone, Debug, PartialEq)] pub struct LineDetailData { pub line: String, pub kind: crate::model::VehicleKind, pub color: String, pub text_color: String, pub destination: crate::model::Bilingual, pub strip: Vec, pub in_seconds: Vec, pub disruption: Option, pub shape: Option>, } #[derive(Debug)] pub(crate) struct LineDetailUnavailable; impl std::fmt::Display for LineDetailUnavailable { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { formatter.write_str("static transit data unavailable") } } impl std::error::Error for LineDetailUnavailable {} #[derive(Clone, Debug, PartialEq)] pub struct VilloStopData { pub id: String, pub name: crate::model::Bilingual, pub distance_m: f64, pub bikes: u16, pub docks: u16, pub mechanical: Option, pub electrical: Option, } pub struct VilloBoardData { pub discovery_ready: bool, pub data_age_seconds: Option, pub stops: Vec, } impl StaticBatch { pub fn validate(&self) -> anyhow::Result<()> { let scopes: HashSet<&str> = self.raw.iter().map(|raw| raw.scope.as_ref()).collect(); ensure!( self.raw.len() == 3 && scopes == HashSet::from(["stop-details", "gtfs", "stops-by-line"]), "static batch must contain exactly stop-details, gtfs, and stops-by-line" ); ensure!( self.raw .iter() .all(|raw| raw.provider == "bmc" && raw.feed == "static"), "static raw feeds must belong to bmc/static" ); ensure!(!self.stops.is_empty(), "static stops are empty"); ensure!(!self.routes.is_empty(), "static routes are empty"); ensure!( self.sequences.values().any(|sequences| !sequences.is_empty()), "static sequences are empty" ); self.alias_targets()?; for (line, route) in &self.routes { ensure!( line.as_ref() == route.line, "route key {line} does not match line {}", route.line ); } for (line, sequences) in &self.sequences { ensure!( self.routes.contains_key(line.as_ref()), "sequence line {line} is unknown" ); for sequence in sequences { ensure!( sequence.line == line.as_ref(), "sequence key {line} does not match line {}", sequence.line ); ensure!(!sequence.stops.is_empty(), "sequence for line {line} is empty"); } } for ((line, destination), points) in &self.shapes { ensure!(self.routes.contains_key(line.as_str()), "shape line {line} is unknown"); ensure!( points.len() >= 2, "shape {line}/{destination} has fewer than two points" ); } Ok(()) } pub(super) fn alias_targets(&self) -> anyhow::Result> { let mut aliases = HashMap::with_capacity(self.stops.len() * 2); for (index, stop) in self.stops.iter().enumerate() { ensure!( aliases.insert(stop.id.as_str(), index).is_none(), "ambiguous exact stop alias {}", stop.id ); } for (index, stop) in self.stops.iter().enumerate() { let Some(alias) = crate::bmc::strip_platform_suffix(&stop.id) else { continue; }; aliases.entry(alias).or_insert(index); } Ok(aliases) } } impl Store { pub async fn open( path: impl AsRef, spatialite_extension: impl AsRef, ) -> anyhow::Result { let path = Arc::new(path.as_ref().to_path_buf()); let extension_path = Arc::new(spatialite_extension.as_ref().to_path_buf()); let config = deadpool_sqlite::Config::new(path.as_ref().clone()); let manager = Manager::from_config(&config, Runtime::Tokio1); let pool = Pool::builder(manager) .runtime(Runtime::Tokio1) .max_size(8) .wait_timeout(Some(Duration::from_secs(5))) .post_create(Hook::async_fn(move |connection, _| { let extension_path = Arc::clone(&extension_path); Box::pin(async move { connection .interact(move |connection| { connection.busy_timeout(Duration::from_secs(5))?; connection.execute_batch( "PRAGMA foreign_keys=ON; PRAGMA journal_mode=WAL; PRAGMA synchronous=FULL;", )?; unsafe { connection.load_extension_enable()?; connection.load_extension(extension_path.as_path(), None::<&str>)?; connection.load_extension_disable()?; } connection.create_scalar_function( "bmc_ids_match", 2, FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC, |context| { let static_id = context .get_raw(0) .as_str() .map_err(|error| SqliteError::UserFunctionError(error.into()))?; let realtime_id = context .get_raw(1) .as_str() .map_err(|error| SqliteError::UserFunctionError(error.into()))?; Ok(crate::bmc::ids_match(static_id, realtime_id)) }, )?; Ok(()) }) .await .map_err(|error| HookError::message(error.to_string()))? .map_err(HookError::Backend) }) })) .build() .context("failed to build SQLite connection pool")?; let store = Self { pool, path, last_read_failure: Arc::new(Mutex::new(None)), }; store.migrate().await?; Ok(store) } pub async fn schema_version(&self) -> anyhow::Result { self.read("schema_version", |connection| { Ok(connection.query_row("SELECT version FROM schema_meta", [], |row| row.get(0))?) }) .await } fn record_read_failure(&self, stage: &str, operation: &str, error: &anyhow::Error) { let mut failure = self .last_read_failure .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); *failure = Some(StorageReadFailure { stage: stage.into(), operation: operation.into(), message: format!("{error:#}"), failed_at: chrono::Utc::now().timestamp_millis(), }); } pub(super) async fn read(&self, operation: &'static str, read: F) -> anyhow::Result where T: Send + 'static, F: FnOnce(&mut deadpool_sqlite::rusqlite::Connection) -> anyhow::Result + Send + 'static, { let connection = match self.pool.get().await { Ok(connection) => connection, Err(error) => { let error = anyhow::Error::new(error).context(format!( "failed to acquire connection for {operation} from {}", self.path.display() )); self.record_read_failure("acquire", operation, &error); return Err(error); } }; let result = match connection.interact(read).await { Ok(result) => result, Err(error) => Err(anyhow::anyhow!("SQLite interaction failed: {error}")), } .with_context(|| { format!( "failed to run storage read {operation} from {}", self.path.display() ) }); if let Err(error) = &result { self.record_read_failure("query", operation, error); } result } #[cfg(test)] pub(crate) async fn hold_connection_for_test( &self, ready: tokio::sync::mpsc::Sender<()>, release: Arc, ) -> anyhow::Result<()> { let connection = self.pool.get().await?; connection .interact(move |_| { ready.blocking_send(()).unwrap(); release.wait(); }) .await .map_err(|error| anyhow::anyhow!("SQLite interaction failed: {error}")) } async fn migrate(&self) -> anyhow::Result<()> { let connection = self.pool.get().await.with_context(|| { format!("failed to get connection for {}", self.path.display()) })?; connection .interact(|connection| migrate_connection(connection, || Ok(()))) .await .map_err(|error| anyhow::anyhow!("SQLite interaction failed: {error}"))? .with_context(|| format!("failed to migrate {}", self.path.display())) } } const V1_VERSION_SQL: &str = "INSERT INTO schema_meta(version, applied_at)\nVALUES (1, unixepoch('subsec') * 1000);\n"; const V1_APPLICATION_START: &str = "DROP TABLE IF EXISTS knn2;"; fn migrate_connection( connection: &mut deadpool_sqlite::rusqlite::Connection, mut before_v1_version: impl FnMut() -> anyhow::Result<()>, ) -> anyhow::Result<()> { let schema_exists = connection.query_row( "SELECT EXISTS( SELECT 1 FROM sqlite_schema WHERE type = 'table' AND name = 'schema_meta' )", [], |row| row.get::<_, bool>(0), )?; let mut version = if schema_exists { let version = connection.query_row( "SELECT CASE WHEN COUNT(*) = 1 THEN MIN(version) END FROM schema_meta", [], |row| row.get::<_, Option>(0), )?; ensure!(version.is_some(), "unsupported schema version: {version:?}"); version.unwrap() } else { 0 }; ensure!( (0..=CURRENT_SCHEMA_VERSION).contains(&version), "unsupported schema version: {version}" ); while version < CURRENT_SCHEMA_VERSION { let next = version + 1; if next == 1 { let bootstrap = MIGRATIONS[0] .strip_suffix(V1_VERSION_SQL) .context("v1 migration version statement changed")?; let application_start = bootstrap .find(V1_APPLICATION_START) .context("v1 migration application start changed")?; let (spatial_init, application) = bootstrap.split_at(application_start); let spatial_initialized = connection.query_row( "SELECT count(*) = 2 FROM sqlite_schema WHERE type = 'table' AND name IN ('geometry_columns', 'spatial_ref_sys')", [], |row| row.get::<_, bool>(0), )?; if spatial_initialized { connection.execute_batch("PRAGMA auto_vacuum = INCREMENTAL;")?; } else { connection.execute_batch(spatial_init)?; } let transaction = connection.transaction()?; transaction.execute_batch(application)?; before_v1_version()?; transaction.execute_batch(V1_VERSION_SQL)?; transaction.commit()?; } else { if version == 1 { connection.execute_batch( "CREATE INDEX IF NOT EXISTS place_normalized_lookup ON place(provider, kind, normalized_id, active, source_order);", )?; } let transaction = connection.transaction()?; transaction.execute_batch(MIGRATIONS[(next - 1) as usize])?; transaction.commit()?; } version = connection .query_row( "SELECT CASE WHEN COUNT(*) = 1 THEN MIN(version) END FROM schema_meta", [], |row| row.get::<_, Option>(0), )? .with_context(|| format!("migration {next} did not record its version"))?; ensure!(version == next, "migration {next} recorded version {version}"); } connection.execute_batch( "CREATE INDEX IF NOT EXISTS place_normalized_lookup ON place(provider, kind, normalized_id, active, source_order); CREATE INDEX IF NOT EXISTS provider_snapshot_retention ON provider_snapshot(captured_at); CREATE INDEX IF NOT EXISTS arrival_fact_retention ON arrival_fact(last_seen_at); CREATE INDEX IF NOT EXISTS train_departure_fact_retention ON train_departure_fact(last_seen_at); CREATE INDEX IF NOT EXISTS disruption_fact_retention ON disruption_fact(last_seen_at); CREATE INDEX IF NOT EXISTS vehicle_position_fact_retention ON vehicle_position_fact(observed_at); CREATE INDEX IF NOT EXISTS villo_status_fact_retention ON villo_status_fact(observed_at);", )?; Ok(()) } #[cfg(test)] mod tests { use super::*; #[tokio::test] async fn v1_bootstrap_failure_rolls_back_application_schema_and_reopens() { let directory = tempfile::tempdir().unwrap(); let path = directory.path().join("stibble.sqlite"); let extension = PathBuf::from(std::env::var("SPATIALITE_EXTENSION").unwrap()); let mut connection = deadpool_sqlite::rusqlite::Connection::open(&path).unwrap(); unsafe { connection.load_extension_enable().unwrap(); connection .load_extension(&extension, None::<&str>) .unwrap(); connection.load_extension_disable().unwrap(); } let failure = migrate_connection(&mut connection, || { anyhow::bail!("injected bootstrap failure") }); assert!(failure.is_err()); assert_eq!( connection .query_row( "SELECT count(*) FROM sqlite_schema WHERE name = 'schema_meta'", [], |row| row.get::<_, i64>(0), ) .unwrap(), 0, ); assert!( connection .query_row( "SELECT count(*) FROM spatial_ref_sys", [], |row| row.get::<_, i64>(0), ) .unwrap() > 0 ); drop(connection); let reopened = Store::open(&path, &extension).await.unwrap(); assert_eq!(reopened.schema_version().await.unwrap(), CURRENT_SCHEMA_VERSION); assert!(path.exists()); } }