use anyhow::Result;
use sqlx::AnyPool;
mod clients;
mod context;
mod credentials;
mod messages;
mod migrations;
mod sessions;
pub use clients::{ClientRow, HUB_DEFAULT_SENTINEL};
pub use messages::{MessageRow, RecentOutboundRow, SessionStatusEntry};
pub use sessions::BackendSessionRow;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DatabaseKind {
Sqlite,
Postgres,
MySql,
}
impl DatabaseKind {
fn from_url(url: &str) -> Result<Self> {
if url.starts_with("postgres:") || url.starts_with("postgresql:") {
#[cfg(feature = "postgres")]
return Ok(DatabaseKind::Postgres);
#[cfg(not(feature = "postgres"))]
anyhow::bail!(
"PostgreSQL support requires the `postgres` feature flag to be enabled at compile \
time; rebuild with --features postgres"
);
} else if url.starts_with("mysql:") || url.starts_with("mariadb:") {
#[cfg(feature = "mysql")]
return Ok(DatabaseKind::MySql);
#[cfg(not(feature = "mysql"))]
anyhow::bail!(
"MySQL support requires the `mysql` feature flag to be enabled at compile time; \
rebuild with `--features mysql`"
);
} else if url.starts_with("sqlite:") || url.is_empty() {
Ok(DatabaseKind::Sqlite)
} else {
if (url.starts_with('/') || url.starts_with("./") || url.starts_with("~/"))
&& !url.contains("://")
{
anyhow::bail!(
"DATABASE_URL {:?} looks like a file path; use the sqlite: scheme instead, e.g. `sqlite:{}`",
url, url
);
}
anyhow::bail!(
"unsupported DATABASE_URL scheme in {:?}; \
supported schemes: sqlite:, postgres://, postgresql://, mysql://, mariadb://",
url
)
}
}
}
pub struct Store {
pool: AnyPool,
rpool: AnyPool,
kind: DatabaseKind,
master_key: std::sync::OnceLock<std::sync::Arc<ring::aead::LessSafeKey>>,
}
impl Store {
pub async fn connect(url: &str) -> Result<Self> {
#[cfg(test)]
{
static TEST_INIT_KEY: std::sync::Once = std::sync::Once::new();
TEST_INIT_KEY.call_once(|| {
if std::env::var("ILINK_HUB_MASTER_KEY").is_err() {
std::env::set_var(
"ILINK_HUB_MASTER_KEY",
"000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f",
);
}
});
}
sqlx::any::install_default_drivers();
let kind = DatabaseKind::from_url(url)?;
let is_sqlite = kind == DatabaseKind::Sqlite;
if is_sqlite {
let url_owned = url.to_string();
tokio::time::timeout(
std::time::Duration::from_secs(10),
tokio::task::spawn_blocking(move || Self::ensure_sqlite_file(&url_owned)),
)
.await
.map_err(|_| anyhow::anyhow!("ensure_sqlite_file timed out after 10s"))?
.map_err(|e| anyhow::anyhow!("spawn_blocking failed: {e}"))??;
}
let is_memory_sqlite = is_sqlite && (url.contains(":memory:") || url.ends_with("sqlite:"));
let (pool, rpool) = if is_sqlite {
let wpool = sqlx::pool::PoolOptions::<sqlx::Any>::new()
.max_connections(1)
.after_connect(|conn, _meta| {
Box::pin(async move {
sqlx::query("PRAGMA busy_timeout=5000; PRAGMA journal_mode=WAL")
.execute(&mut *conn)
.await?;
Ok(())
})
})
.connect(url)
.await?;
let rpool = if is_memory_sqlite {
wpool.clone()
} else {
sqlx::pool::PoolOptions::<sqlx::Any>::new()
.max_connections(4)
.after_connect(|conn, _meta| {
Box::pin(async move {
sqlx::query("PRAGMA busy_timeout=5000; PRAGMA journal_mode=WAL")
.execute(&mut *conn)
.await?;
Ok(())
})
})
.connect(url)
.await?
};
(wpool, rpool)
} else {
let pool = AnyPool::connect(url).await?;
(pool.clone(), pool)
};
let store = Self {
pool,
rpool,
kind,
master_key: std::sync::OnceLock::new(),
};
store.run_migrations().await?;
Ok(store)
}
pub fn set_master_key(
&self,
key: std::sync::Arc<ring::aead::LessSafeKey>,
) -> Result<(), std::sync::Arc<ring::aead::LessSafeKey>> {
self.master_key.set(key)
}
pub fn master_key(&self) -> Option<&std::sync::Arc<ring::aead::LessSafeKey>> {
self.master_key.get()
}
#[cfg(test)]
pub fn pool(&self) -> &sqlx::AnyPool {
&self.pool
}
fn ensure_sqlite_file(url: &str) -> Result<()> {
let path_part = url
.strip_prefix("sqlite:///")
.or_else(|| url.strip_prefix("sqlite://"))
.or_else(|| url.strip_prefix("sqlite:"))
.unwrap_or("");
let path_str = path_part.split('?').next().unwrap_or("").trim();
if path_str.is_empty() || path_str == ":memory:" {
return Ok(());
}
let path = std::path::Path::new(path_str);
if let Some(parent) = path.parent() {
if !parent.as_os_str().is_empty() {
std::fs::create_dir_all(parent)?;
}
}
if !path.exists() {
match std::fs::OpenOptions::new()
.create_new(true)
.write(true)
.open(path)
{
Ok(_) => {}
Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
}
Err(e) => return Err(e.into()),
}
}
Ok(())
}
}
#[cfg(test)]
mod store_tests;