70 lines
2.1 KiB
Rust
70 lines
2.1 KiB
Rust
|
use sqlx::any::AnyPoolOptions;
|
||
|
|
use sqlx::{AnyPool, Row};
|
||
|
|
|
||
|
|
#[derive(Clone)]
|
||
|
|
pub struct Storage {
|
||
|
|
pool: AnyPool,
|
||
|
|
postgres: bool,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl Storage {
|
||
|
|
pub async fn connect(database_url: &str) -> Result<Self, sqlx::Error> {
|
||
|
|
sqlx::any::install_default_drivers();
|
||
|
|
let postgres =
|
||
|
|
database_url.starts_with("postgres://") || database_url.starts_with("postgresql://");
|
||
|
|
let pool = AnyPoolOptions::new()
|
||
|
|
.max_connections(1)
|
||
|
|
.connect(database_url)
|
||
|
|
.await?;
|
||
|
|
sqlx::query(
|
||
|
|
"CREATE TABLE IF NOT EXISTS sim_state (id INTEGER PRIMARY KEY, snapshot TEXT NOT NULL)",
|
||
|
|
)
|
||
|
|
.execute(&pool)
|
||
|
|
.await?;
|
||
|
|
Ok(Self { pool, postgres })
|
||
|
|
}
|
||
|
|
|
||
|
|
pub async fn load(&self) -> Result<Option<String>, sqlx::Error> {
|
||
|
|
let row = sqlx::query("SELECT snapshot FROM sim_state WHERE id = 1")
|
||
|
|
.fetch_optional(&self.pool)
|
||
|
|
.await?;
|
||
|
|
row.map(|row| row.try_get("snapshot")).transpose()
|
||
|
|
}
|
||
|
|
|
||
|
|
pub async fn save(&self, snapshot: &str) -> Result<(), sqlx::Error> {
|
||
|
|
let query = if self.postgres {
|
||
|
|
"INSERT INTO sim_state (id, snapshot) VALUES (1, $1) ON CONFLICT (id) DO UPDATE SET snapshot = excluded.snapshot"
|
||
|
|
} else {
|
||
|
|
"INSERT INTO sim_state (id, snapshot) VALUES (1, ?) ON CONFLICT (id) DO UPDATE SET snapshot = excluded.snapshot"
|
||
|
|
};
|
||
|
|
sqlx::query(query)
|
||
|
|
.bind(snapshot)
|
||
|
|
.execute(&self.pool)
|
||
|
|
.await?;
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
#[cfg(test)]
|
||
|
|
mod tests {
|
||
|
|
use super::Storage;
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn sqlite_storage_round_trips_snapshots() {
|
||
|
|
let storage = Storage::connect("sqlite::memory:").await.unwrap();
|
||
|
|
assert_eq!(storage.load().await.unwrap(), None);
|
||
|
|
|
||
|
|
storage.save("{\"tick\":12}").await.unwrap();
|
||
|
|
assert_eq!(
|
||
|
|
storage.load().await.unwrap().as_deref(),
|
||
|
|
Some("{\"tick\":12}")
|
||
|
|
);
|
||
|
|
|
||
|
|
storage.save("{\"tick\":13}").await.unwrap();
|
||
|
|
assert_eq!(
|
||
|
|
storage.load().await.unwrap().as_deref(),
|
||
|
|
Some("{\"tick\":13}")
|
||
|
|
);
|
||
|
|
}
|
||
|
|
}
|