Add multi-database data layer

Introduce aether-data-schema and driver-specific schema generation for Postgres, MySQL, and SQLite.

Split data backends, lifecycle, repositories, and gateway runtime integration across database drivers.

Verified with cargo fmt --all --check, cargo clippy --workspace --all-targets -- -D warnings, and cargo test --workspace.
This commit is contained in:
fawney19
2026-05-05 18:27:36 +08:00
parent 099653f732
commit fce7e959e5
372 changed files with 86217 additions and 21160 deletions

View File

@@ -0,0 +1,10 @@
//! Low-level data driver primitives.
//!
//! These modules own pools, transactions, leases, and Redis client helpers.
//! Domain repository logic belongs in `repository/*`, and application-facing
//! composition belongs in `backend`.
pub mod mysql;
pub mod postgres;
pub mod redis;
pub mod sqlite;

View File

@@ -0,0 +1,3 @@
mod pool;
pub use pool::{MysqlPool, MysqlPoolConfig, MysqlPoolFactory};

View File

@@ -0,0 +1,94 @@
use std::str::FromStr;
use std::time::Duration;
use crate::database::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig};
use crate::DataLayerError;
use sqlx::mysql::{MySqlConnectOptions, MySqlPoolOptions, MySqlSslMode};
use sqlx::MySqlPool as SqlxMysqlPool;
pub type MysqlPool = SqlxMysqlPool;
pub type MysqlPoolConfig = SqlDatabaseConfig;
#[derive(Debug, Clone)]
pub struct MysqlPoolFactory {
config: MysqlPoolConfig,
}
impl MysqlPoolFactory {
pub fn new(config: MysqlPoolConfig) -> Result<Self, DataLayerError> {
if config.driver != DatabaseDriver::Mysql {
return Err(DataLayerError::InvalidConfiguration(format!(
"mysql pool requires mysql driver, got {}",
config.driver
)));
}
config.validate()?;
Ok(Self { config })
}
pub fn config(&self) -> &MysqlPoolConfig {
&self.config
}
pub fn connect_options(&self) -> Result<MySqlConnectOptions, DataLayerError> {
let ssl_mode = if self.config.pool.require_ssl {
MySqlSslMode::Required
} else {
MySqlSslMode::Preferred
};
MySqlConnectOptions::from_str(self.config.url.trim())
.map(|options| {
options
.ssl_mode(ssl_mode)
.statement_cache_capacity(self.config.pool.statement_cache_capacity)
})
.map_err(|err| {
DataLayerError::InvalidConfiguration(format!("invalid mysql database url: {err}"))
})
}
pub fn connect_lazy(&self) -> Result<MysqlPool, DataLayerError> {
let SqlPoolConfig {
min_connections,
max_connections,
acquire_timeout_ms,
idle_timeout_ms,
max_lifetime_ms,
..
} = self.config.pool;
Ok(MySqlPoolOptions::new()
.min_connections(min_connections)
.max_connections(max_connections)
.acquire_timeout(Duration::from_millis(acquire_timeout_ms))
.idle_timeout(Duration::from_millis(idle_timeout_ms))
.max_lifetime(Duration::from_millis(max_lifetime_ms))
.connect_lazy_with(self.connect_options()?))
}
}
#[cfg(test)]
mod tests {
use super::MysqlPoolFactory;
use crate::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig};
#[tokio::test]
async fn factory_builds_lazy_pool_from_valid_config() {
let config = SqlDatabaseConfig {
driver: DatabaseDriver::Mysql,
url: "mysql://user:pass@localhost:3306/aether".to_string(),
pool: SqlPoolConfig {
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
},
};
let factory = MysqlPoolFactory::new(config).expect("factory should build");
let _pool = factory.connect_lazy().expect("lazy pool should build");
}
}

View File

@@ -0,0 +1,431 @@
use crate::driver::postgres::{
DatabaseRecordId, PostgresTransactionOptions, PostgresTransactionRunner,
};
use crate::error::SqlxResultExt;
use crate::DataLayerError;
use futures_util::{FutureExt, TryStreamExt};
use sqlx::query_scalar;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PostgresLeaseClaimOptions {
pub batch_size: usize,
pub lease_ms: u64,
}
impl PostgresLeaseClaimOptions {
pub fn validate(&self) -> Result<(), DataLayerError> {
if self.batch_size == 0 {
return Err(DataLayerError::InvalidConfiguration(
"postgres lease batch_size must be positive".to_string(),
));
}
if self.lease_ms == 0 {
return Err(DataLayerError::InvalidConfiguration(
"postgres lease lease_ms must be positive".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PostgresLeaseClaimSpec {
pub table: &'static str,
pub id_column: &'static str,
pub lease_owner_column: &'static str,
pub lease_expires_at_column: &'static str,
pub eligibility_predicate_sql: &'static str,
pub order_by_sql: &'static str,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct PostgresLeaseRunnerConfig {
pub statement_timeout_ms: Option<u64>,
pub lock_timeout_ms: Option<u64>,
}
impl PostgresLeaseRunnerConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if matches!(self.statement_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"postgres lease statement_timeout_ms must be positive".to_string(),
));
}
if matches!(self.lock_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"postgres lease lock_timeout_ms must be positive".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct PostgresLeaseRunner {
transaction_runner: PostgresTransactionRunner,
config: PostgresLeaseRunnerConfig,
}
impl PostgresLeaseRunner {
pub fn new(
transaction_runner: PostgresTransactionRunner,
config: PostgresLeaseRunnerConfig,
) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self {
transaction_runner,
config,
})
}
pub fn config(&self) -> PostgresLeaseRunnerConfig {
self.config
}
pub fn transaction_runner(&self) -> &PostgresTransactionRunner {
&self.transaction_runner
}
pub async fn claim_ids(
&self,
spec: &PostgresLeaseClaimSpec,
options: PostgresLeaseClaimOptions,
owner: &str,
) -> Result<Vec<DatabaseRecordId>, DataLayerError> {
validate_lease_owner(owner)?;
let sql = build_postgres_lease_claim_sql(spec, options)?;
let owner = owner.trim().to_string();
let lease_ms = i64::try_from(options.lease_ms).map_err(|_| {
DataLayerError::InvalidInput("postgres lease lease_ms exceeds i64 range".to_string())
})?;
let tx_options = PostgresTransactionOptions {
statement_timeout_ms: self.config.statement_timeout_ms,
lock_timeout_ms: self.config.lock_timeout_ms,
..PostgresTransactionOptions::read_write()
};
self.transaction_runner
.run(tx_options, |tx| {
async move {
let mut rows = query_scalar::<_, String>(&sql)
.bind(owner)
.bind(lease_ms)
.fetch(&mut **tx);
let mut ids = Vec::new();
while let Some(id) = rows.try_next().await.map_postgres_err()? {
ids.push(DatabaseRecordId(id));
}
Ok(ids)
}
.boxed()
})
.await
}
pub async fn release_ids(
&self,
spec: &PostgresLeaseClaimSpec,
ids: &[DatabaseRecordId],
owner: &str,
) -> Result<Vec<DatabaseRecordId>, DataLayerError> {
validate_lease_owner(owner)?;
if ids.is_empty() {
return Ok(Vec::new());
}
let sql = build_postgres_lease_release_sql(spec)?;
let owner = owner.trim().to_string();
let ids = ids.iter().map(|id| id.0.clone()).collect::<Vec<_>>();
let tx_options = PostgresTransactionOptions {
statement_timeout_ms: self.config.statement_timeout_ms,
lock_timeout_ms: self.config.lock_timeout_ms,
..PostgresTransactionOptions::read_write()
};
self.transaction_runner
.run(tx_options, |tx| {
async move {
let mut rows = query_scalar::<_, String>(&sql)
.bind(ids)
.bind(owner)
.fetch(&mut **tx);
let mut released = Vec::new();
while let Some(id) = rows.try_next().await.map_postgres_err()? {
released.push(DatabaseRecordId(id));
}
Ok(released)
}
.boxed()
})
.await
}
pub async fn renew_ids(
&self,
spec: &PostgresLeaseClaimSpec,
ids: &[DatabaseRecordId],
owner: &str,
lease_ms: u64,
) -> Result<Vec<DatabaseRecordId>, DataLayerError> {
validate_lease_owner(owner)?;
if lease_ms == 0 {
return Err(DataLayerError::InvalidInput(
"postgres lease lease_ms must be positive".to_string(),
));
}
if ids.is_empty() {
return Ok(Vec::new());
}
let sql = build_postgres_lease_renew_sql(spec)?;
let owner = owner.trim().to_string();
let lease_ms = i64::try_from(lease_ms).map_err(|_| {
DataLayerError::InvalidInput("postgres lease lease_ms exceeds i64 range".to_string())
})?;
let ids = ids.iter().map(|id| id.0.clone()).collect::<Vec<_>>();
let tx_options = PostgresTransactionOptions {
statement_timeout_ms: self.config.statement_timeout_ms,
lock_timeout_ms: self.config.lock_timeout_ms,
..PostgresTransactionOptions::read_write()
};
self.transaction_runner
.run(tx_options, |tx| {
async move {
let mut rows = query_scalar::<_, String>(&sql)
.bind(ids)
.bind(owner)
.bind(lease_ms)
.fetch(&mut **tx);
let mut renewed = Vec::new();
while let Some(id) = rows.try_next().await.map_postgres_err()? {
renewed.push(DatabaseRecordId(id));
}
Ok(renewed)
}
.boxed()
})
.await
}
}
pub fn build_postgres_lease_claim_sql(
spec: &PostgresLeaseClaimSpec,
options: PostgresLeaseClaimOptions,
) -> Result<String, DataLayerError> {
options.validate()?;
validate_lease_spec(spec)?;
Ok(format!(
"WITH claimable AS (\
SELECT {id_column} \
FROM {table} \
WHERE ({eligibility_predicate_sql}) \
AND ({lease_expires_at_column} IS NULL OR {lease_expires_at_column} <= NOW()) \
ORDER BY {order_by_sql} \
FOR UPDATE SKIP LOCKED \
LIMIT {batch_size}\
) \
UPDATE {table} AS target \
SET {lease_owner_column} = $1, \
{lease_expires_at_column} = NOW() + ($2::bigint * INTERVAL '1 millisecond') \
FROM claimable \
WHERE target.{id_column} = claimable.{id_column} \
RETURNING target.{id_column}",
id_column = spec.id_column,
table = spec.table,
eligibility_predicate_sql = spec.eligibility_predicate_sql,
lease_expires_at_column = spec.lease_expires_at_column,
order_by_sql = spec.order_by_sql,
batch_size = options.batch_size,
lease_owner_column = spec.lease_owner_column,
))
}
pub fn build_postgres_lease_release_sql(
spec: &PostgresLeaseClaimSpec,
) -> Result<String, DataLayerError> {
validate_lease_spec(spec)?;
Ok(format!(
"UPDATE {table} \
SET {lease_owner_column} = NULL, \
{lease_expires_at_column} = NULL \
WHERE {id_column} = ANY($1) \
AND {lease_owner_column} = $2 \
RETURNING {id_column}",
table = spec.table,
id_column = spec.id_column,
lease_owner_column = spec.lease_owner_column,
lease_expires_at_column = spec.lease_expires_at_column,
))
}
pub fn build_postgres_lease_renew_sql(
spec: &PostgresLeaseClaimSpec,
) -> Result<String, DataLayerError> {
validate_lease_spec(spec)?;
Ok(format!(
"UPDATE {table} \
SET {lease_expires_at_column} = NOW() + ($3::bigint * INTERVAL '1 millisecond') \
WHERE {id_column} = ANY($1) \
AND {lease_owner_column} = $2 \
RETURNING {id_column}",
table = spec.table,
id_column = spec.id_column,
lease_owner_column = spec.lease_owner_column,
lease_expires_at_column = spec.lease_expires_at_column,
))
}
fn validate_lease_spec(spec: &PostgresLeaseClaimSpec) -> Result<(), DataLayerError> {
for (field, value) in [
("table", spec.table),
("id_column", spec.id_column),
("lease_owner_column", spec.lease_owner_column),
("lease_expires_at_column", spec.lease_expires_at_column),
("eligibility_predicate_sql", spec.eligibility_predicate_sql),
("order_by_sql", spec.order_by_sql),
] {
if value.trim().is_empty() {
return Err(DataLayerError::InvalidConfiguration(format!(
"postgres lease {field} cannot be empty"
)));
}
}
Ok(())
}
fn validate_lease_owner(owner: &str) -> Result<(), DataLayerError> {
if owner.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"postgres lease owner cannot be empty".to_string(),
));
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::{
build_postgres_lease_claim_sql, build_postgres_lease_release_sql,
build_postgres_lease_renew_sql, PostgresLeaseClaimOptions, PostgresLeaseClaimSpec,
PostgresLeaseRunner, PostgresLeaseRunnerConfig,
};
use crate::driver::postgres::{
PostgresPoolConfig, PostgresPoolFactory, PostgresTransactionRunner,
};
fn sample_spec() -> PostgresLeaseClaimSpec {
PostgresLeaseClaimSpec {
table: "video_tasks",
id_column: "id",
lease_owner_column: "lease_owner",
lease_expires_at_column: "lease_expires_at",
eligibility_predicate_sql: "status IN ('submitted', 'processing')",
order_by_sql: "updated_at ASC",
}
}
#[test]
fn builds_skip_locked_claim_sql() {
let sql = build_postgres_lease_claim_sql(
&sample_spec(),
PostgresLeaseClaimOptions {
batch_size: 16,
lease_ms: 15_000,
},
)
.expect("claim SQL should build");
assert!(sql.contains("FOR UPDATE SKIP LOCKED"));
assert!(sql.contains("LIMIT 16"));
assert!(sql.contains("lease_owner = $1"));
assert!(sql.contains("NOW() + ($2::bigint * INTERVAL '1 millisecond')"));
}
#[test]
fn builds_release_sql() {
let sql =
build_postgres_lease_release_sql(&sample_spec()).expect("release SQL should build");
assert!(sql.contains("id = ANY($1)"));
assert!(sql.contains("lease_owner = $2"));
assert!(sql.contains("lease_expires_at = NULL"));
}
#[test]
fn builds_renew_sql() {
let sql = build_postgres_lease_renew_sql(&sample_spec()).expect("renew SQL should build");
assert!(sql.contains("id = ANY($1)"));
assert!(sql.contains("lease_owner = $2"));
assert!(sql.contains("NOW() + ($3::bigint * INTERVAL '1 millisecond')"));
}
#[tokio::test]
async fn lease_runner_reuses_transaction_runner() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let pool = factory.connect_lazy().expect("lazy pool should build");
let transaction_runner = PostgresTransactionRunner::new(pool);
let lease_runner = PostgresLeaseRunner::new(
transaction_runner.clone(),
PostgresLeaseRunnerConfig {
statement_timeout_ms: Some(2_000),
lock_timeout_ms: Some(500),
},
)
.expect("lease runner should build");
assert_eq!(
lease_runner.config(),
PostgresLeaseRunnerConfig {
statement_timeout_ms: Some(2_000),
lock_timeout_ms: Some(500),
}
);
let _runner_ref = lease_runner.transaction_runner();
}
#[tokio::test]
async fn release_and_renew_empty_ids_short_circuit() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let pool = factory.connect_lazy().expect("lazy pool should build");
let runner = PostgresLeaseRunner::new(
PostgresTransactionRunner::new(pool),
PostgresLeaseRunnerConfig::default(),
)
.expect("lease runner should build");
assert!(runner
.release_ids(&sample_spec(), &[], "worker-1")
.await
.expect("empty release should succeed")
.is_empty());
assert!(runner
.renew_ids(&sample_spec(), &[], "worker-1", 5_000)
.await
.expect("empty renew should succeed")
.is_empty());
}
}

View File

@@ -0,0 +1,15 @@
mod lease;
mod pool;
mod tx;
mod types;
pub use lease::{
build_postgres_lease_claim_sql, build_postgres_lease_release_sql,
build_postgres_lease_renew_sql, PostgresLeaseClaimOptions, PostgresLeaseClaimSpec,
PostgresLeaseRunner, PostgresLeaseRunnerConfig,
};
pub use pool::{PostgresPool, PostgresPoolConfig, PostgresPoolFactory};
pub use tx::{
PostgresTransaction, PostgresTransactionOptions, PostgresTransactionRunner, TransactionMode,
};
pub use types::DatabaseRecordId;

View File

@@ -0,0 +1,122 @@
use crate::DataLayerError;
use sqlx::postgres::{PgConnectOptions, PgPoolOptions, PgSslMode};
use sqlx::PgPool;
use std::str::FromStr;
use std::time::Duration;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub struct PostgresPoolConfig {
pub database_url: String,
pub min_connections: u32,
pub max_connections: u32,
pub acquire_timeout_ms: u64,
pub idle_timeout_ms: u64,
pub max_lifetime_ms: u64,
pub statement_cache_capacity: usize,
pub require_ssl: bool,
}
impl Default for PostgresPoolConfig {
fn default() -> Self {
Self {
database_url: String::new(),
min_connections: 1,
max_connections: 20,
acquire_timeout_ms: 10_000,
idle_timeout_ms: 30_000,
max_lifetime_ms: 30 * 60_000,
statement_cache_capacity: 100,
require_ssl: false,
}
}
}
impl PostgresPoolConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if self.database_url.trim().is_empty() {
return Err(DataLayerError::InvalidConfiguration(
"postgres database_url cannot be empty".to_string(),
));
}
if self.min_connections > self.max_connections {
return Err(DataLayerError::InvalidConfiguration(
"postgres min_connections cannot exceed max_connections".to_string(),
));
}
if self.statement_cache_capacity == 0 {
return Err(DataLayerError::InvalidConfiguration(
"postgres statement_cache_capacity must be positive".to_string(),
));
}
Ok(())
}
pub fn connect_options(&self) -> Result<PgConnectOptions, DataLayerError> {
self.validate()?;
let ssl_mode = if self.require_ssl {
PgSslMode::Require
} else {
PgSslMode::Prefer
};
let options = PgConnectOptions::from_str(self.database_url.trim()).map_err(|err| {
DataLayerError::InvalidConfiguration(format!("invalid postgres database_url: {err}"))
})?;
Ok(options
.ssl_mode(ssl_mode)
.statement_cache_capacity(self.statement_cache_capacity))
}
}
pub type PostgresPool = PgPool;
#[derive(Debug, Clone)]
pub struct PostgresPoolFactory {
config: PostgresPoolConfig,
}
impl PostgresPoolFactory {
pub fn new(config: PostgresPoolConfig) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self { config })
}
pub fn config(&self) -> &PostgresPoolConfig {
&self.config
}
pub fn connect_lazy(&self) -> Result<PostgresPool, DataLayerError> {
let options = self.config.connect_options()?;
Ok(PgPoolOptions::new()
.min_connections(self.config.min_connections)
.max_connections(self.config.max_connections)
.acquire_timeout(Duration::from_millis(self.config.acquire_timeout_ms))
.idle_timeout(Duration::from_millis(self.config.idle_timeout_ms))
.max_lifetime(Duration::from_millis(self.config.max_lifetime_ms))
.connect_lazy_with(options))
}
}
#[cfg(test)]
mod tests {
use super::{PostgresPoolConfig, PostgresPoolFactory};
#[tokio::test]
async fn factory_builds_lazy_pool_from_valid_config() {
let config = PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
};
let factory = PostgresPoolFactory::new(config).expect("factory should build");
let _pool = factory.connect_lazy().expect("lazy pool should build");
}
}

View File

@@ -0,0 +1,206 @@
use futures_util::future::BoxFuture;
use sqlx::{Postgres, Transaction};
use crate::driver::postgres::PostgresPool;
use crate::error::{postgres_error, SqlxResultExt};
use crate::DataLayerError;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum TransactionMode {
ReadOnly,
#[default]
ReadWrite,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct PostgresTransactionOptions {
pub mode: TransactionMode,
pub statement_timeout_ms: Option<u64>,
pub lock_timeout_ms: Option<u64>,
}
impl PostgresTransactionOptions {
pub fn read_only() -> Self {
Self {
mode: TransactionMode::ReadOnly,
..Self::default()
}
}
pub fn read_write() -> Self {
Self {
mode: TransactionMode::ReadWrite,
..Self::default()
}
}
pub fn validate(&self) -> Result<(), DataLayerError> {
if matches!(self.statement_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"postgres statement_timeout_ms must be positive".to_string(),
));
}
if matches!(self.lock_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"postgres lock_timeout_ms must be positive".to_string(),
));
}
Ok(())
}
}
pub type PostgresTransaction = Transaction<'static, Postgres>;
#[derive(Debug, Clone)]
pub struct PostgresTransactionRunner {
pool: PostgresPool,
}
impl PostgresTransactionRunner {
pub fn new(pool: PostgresPool) -> Self {
Self { pool }
}
pub fn pool(&self) -> &PostgresPool {
&self.pool
}
pub async fn begin(
&self,
options: PostgresTransactionOptions,
) -> Result<PostgresTransaction, DataLayerError> {
options.validate()?;
let mut tx = self.pool.begin().await.map_postgres_err()?;
for statement in build_transaction_setup_statements(options) {
sqlx::query(statement.as_str())
.execute(&mut *tx)
.await
.map_postgres_err()?;
}
Ok(tx)
}
pub async fn run<T, F>(
&self,
options: PostgresTransactionOptions,
f: F,
) -> Result<T, DataLayerError>
where
F: for<'tx> FnOnce(
&'tx mut PostgresTransaction,
) -> BoxFuture<'tx, Result<T, DataLayerError>>,
{
let mut tx = self.begin(options).await?;
match f(&mut tx).await {
Ok(value) => {
tx.commit().await.map_err(postgres_error)?;
Ok(value)
}
Err(err) => {
let _ = tx.rollback().await;
Err(err)
}
}
}
pub async fn run_read_only<T, F>(&self, f: F) -> Result<T, DataLayerError>
where
F: for<'tx> FnOnce(
&'tx mut PostgresTransaction,
) -> BoxFuture<'tx, Result<T, DataLayerError>>,
{
self.run(PostgresTransactionOptions::read_only(), f).await
}
pub async fn run_read_write<T, F>(&self, f: F) -> Result<T, DataLayerError>
where
F: for<'tx> FnOnce(
&'tx mut PostgresTransaction,
) -> BoxFuture<'tx, Result<T, DataLayerError>>,
{
self.run(PostgresTransactionOptions::read_write(), f).await
}
}
pub(crate) fn build_transaction_setup_statements(
options: PostgresTransactionOptions,
) -> Vec<String> {
let mut statements = Vec::new();
if options.mode == TransactionMode::ReadOnly {
statements.push("SET TRANSACTION READ ONLY".to_string());
}
if let Some(statement_timeout_ms) = options.statement_timeout_ms {
statements.push(format!(
"SET LOCAL statement_timeout = {}",
statement_timeout_ms
));
}
if let Some(lock_timeout_ms) = options.lock_timeout_ms {
statements.push(format!("SET LOCAL lock_timeout = {}", lock_timeout_ms));
}
statements
}
#[cfg(test)]
mod tests {
use super::{
build_transaction_setup_statements, PostgresTransactionOptions, PostgresTransactionRunner,
TransactionMode,
};
use crate::driver::postgres::{PostgresPoolConfig, PostgresPoolFactory};
#[test]
fn validates_transaction_options() {
assert!(PostgresTransactionOptions {
statement_timeout_ms: Some(0),
..PostgresTransactionOptions::default()
}
.validate()
.is_err());
assert!(PostgresTransactionOptions {
lock_timeout_ms: Some(0),
..PostgresTransactionOptions::default()
}
.validate()
.is_err());
}
#[test]
fn builds_expected_setup_statements() {
let statements = build_transaction_setup_statements(PostgresTransactionOptions {
mode: TransactionMode::ReadOnly,
statement_timeout_ms: Some(1_500),
lock_timeout_ms: Some(750),
});
assert_eq!(
statements,
vec![
"SET TRANSACTION READ ONLY".to_string(),
"SET LOCAL statement_timeout = 1500".to_string(),
"SET LOCAL lock_timeout = 750".to_string(),
]
);
}
#[tokio::test]
async fn runner_reuses_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let pool = factory.connect_lazy().expect("lazy pool should build");
let runner = PostgresTransactionRunner::new(pool.clone());
let _pool_ref = runner.pool();
}
}

View File

@@ -0,0 +1,2 @@
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct DatabaseRecordId(pub String);

View File

@@ -0,0 +1,69 @@
use crate::driver::redis::RedisKeyspace;
use crate::error::RedisResultExt;
use crate::DataLayerError;
pub type RedisClient = redis::Client;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub struct RedisClientConfig {
pub url: String,
pub key_prefix: Option<String>,
}
impl RedisClientConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
let raw = self.url.trim();
if raw.is_empty() {
return Err(DataLayerError::InvalidConfiguration(
"redis url cannot be empty".to_string(),
));
}
url::Url::parse(raw).map_err(|err| {
DataLayerError::InvalidConfiguration(format!("invalid redis url: {err}"))
})?;
Ok(())
}
pub fn keyspace(&self) -> RedisKeyspace {
RedisKeyspace::new(self.key_prefix.as_deref())
}
}
#[derive(Debug, Clone)]
pub struct RedisClientFactory {
config: RedisClientConfig,
}
impl RedisClientFactory {
pub fn new(config: RedisClientConfig) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self { config })
}
pub fn config(&self) -> &RedisClientConfig {
&self.config
}
pub fn connect_lazy(&self) -> Result<RedisClient, DataLayerError> {
RedisClient::open(self.config.url.clone()).map_redis_err()
}
}
#[cfg(test)]
mod tests {
use super::{RedisClientConfig, RedisClientFactory};
#[test]
fn factory_builds_lazy_client_from_valid_config() {
let config = RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
};
let factory = RedisClientFactory::new(config.clone()).expect("factory should build");
assert_eq!(factory.config(), &config);
let _client = factory
.connect_lazy()
.expect("lazy redis client should build");
}
}

View File

@@ -0,0 +1,182 @@
use std::future::Future;
use std::time::Duration;
use crate::driver::redis::{RedisClient, RedisKeyspace};
use crate::error::RedisResultExt;
use crate::DataLayerError;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RedisKvRunnerConfig {
pub command_timeout_ms: Option<u64>,
pub default_ttl_seconds: u64,
}
impl Default for RedisKvRunnerConfig {
fn default() -> Self {
Self {
command_timeout_ms: Some(1_000),
default_ttl_seconds: 300,
}
}
}
impl RedisKvRunnerConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if let Some(timeout) = self.command_timeout_ms {
if timeout == 0 {
return Err(DataLayerError::InvalidConfiguration(
"redis kv command_timeout_ms must be positive".to_string(),
));
}
}
if self.default_ttl_seconds == 0 {
return Err(DataLayerError::InvalidConfiguration(
"redis kv default_ttl_seconds must be positive".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct RedisKvRunner {
client: RedisClient,
keyspace: RedisKeyspace,
config: RedisKvRunnerConfig,
}
impl RedisKvRunner {
pub fn new(
client: RedisClient,
keyspace: RedisKeyspace,
config: RedisKvRunnerConfig,
) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self {
client,
keyspace,
config,
})
}
pub fn client(&self) -> &RedisClient {
&self.client
}
pub fn keyspace(&self) -> &RedisKeyspace {
&self.keyspace
}
pub fn config(&self) -> RedisKvRunnerConfig {
self.config
}
pub async fn setex(
&self,
key: &str,
value: &str,
ttl_seconds: Option<u64>,
) -> Result<String, DataLayerError> {
let resolved_ttl = ttl_seconds.unwrap_or(self.config.default_ttl_seconds);
let namespaced_key = self.keyspace.key(key);
self.run_with_timeout("redis kv setex", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
redis::cmd("SETEX")
.arg(&namespaced_key)
.arg(resolved_ttl)
.arg(value)
.query_async(&mut connection)
.await
.map_redis_err()
})
.await
}
pub async fn del(&self, key: &str) -> Result<i64, DataLayerError> {
let namespaced_key = self.keyspace.key(key);
self.run_with_timeout("redis kv del", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
redis::cmd("DEL")
.arg(&namespaced_key)
.query_async(&mut connection)
.await
.map_redis_err()
})
.await
}
async fn run_with_timeout<T, F>(
&self,
operation: &'static str,
future: F,
) -> Result<T, DataLayerError>
where
F: Future<Output = Result<T, DataLayerError>>,
{
if let Some(timeout_ms) = self.config.command_timeout_ms {
tokio::time::timeout(Duration::from_millis(timeout_ms), future)
.await
.map_err(|_| {
DataLayerError::TimedOut(format!("{operation} exceeded {timeout_ms}ms timeout"))
})?
} else {
future.await
}
}
}
#[cfg(test)]
mod tests {
use super::{RedisKvRunner, RedisKvRunnerConfig};
use crate::driver::redis::{RedisClientConfig, RedisClientFactory, RedisKeyspace};
fn build_runner() -> RedisKvRunner {
let config = RedisClientConfig {
url: "redis://localhost/0".to_string(),
key_prefix: Some("aether-test".to_string()),
};
let factory = RedisClientFactory::new(config).expect("redis factory");
let client = factory.connect_lazy().expect("connect");
let keyspace = factory.config().keyspace();
RedisKvRunner::new(client, keyspace, RedisKvRunnerConfig::default()).expect("runner build")
}
#[test]
fn runner_reuses_client_keyspace_and_config() {
let runner = build_runner();
assert_eq!(
runner.keyspace().key("kv:setex:1"),
"aether-test:kv:setex:1"
);
assert_eq!(runner.config(), RedisKvRunnerConfig::default());
let _client = runner.client();
}
#[test]
fn rejects_zero_default_ttl() {
let config = RedisKvRunnerConfig {
command_timeout_ms: Some(100),
default_ttl_seconds: 0,
};
assert!(RedisKvRunner::new(
RedisClientFactory::new(RedisClientConfig {
url: "redis://localhost/0".to_string(),
key_prefix: Some("aether-test".to_string()),
})
.expect("redis factory")
.connect_lazy()
.expect("redis client"),
RedisKeyspace::new(Some("aether-test")),
config,
)
.is_err());
}
}

View File

@@ -0,0 +1,323 @@
use std::future::Future;
use std::time::Duration;
use crate::driver::redis::{RedisClient, RedisKeyspace};
use crate::error::RedisResultExt;
use crate::DataLayerError;
use uuid::Uuid;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RedisLockKey(pub String);
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisLockLease {
pub key: RedisLockKey,
pub owner: String,
pub token: String,
pub ttl_ms: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RedisLockRunnerConfig {
pub command_timeout_ms: Option<u64>,
pub default_ttl_ms: u64,
}
impl Default for RedisLockRunnerConfig {
fn default() -> Self {
Self {
command_timeout_ms: Some(1_000),
default_ttl_ms: 15_000,
}
}
}
impl RedisLockRunnerConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if matches!(self.command_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"redis lock command_timeout_ms must be positive".to_string(),
));
}
if self.default_ttl_ms == 0 {
return Err(DataLayerError::InvalidConfiguration(
"redis lock default_ttl_ms must be positive".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct RedisLockRunner {
client: RedisClient,
keyspace: RedisKeyspace,
config: RedisLockRunnerConfig,
}
impl RedisLockRunner {
pub fn new(
client: RedisClient,
keyspace: RedisKeyspace,
config: RedisLockRunnerConfig,
) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self {
client,
keyspace,
config,
})
}
pub fn client(&self) -> &RedisClient {
&self.client
}
pub fn keyspace(&self) -> &RedisKeyspace {
&self.keyspace
}
pub fn config(&self) -> RedisLockRunnerConfig {
self.config
}
pub async fn try_acquire(
&self,
key: &RedisLockKey,
owner: &str,
ttl_ms: Option<u64>,
) -> Result<Option<RedisLockLease>, DataLayerError> {
validate_owner(owner)?;
validate_key(key)?;
let ttl_ms = self.resolve_ttl_ms(ttl_ms)?;
let token = format!("{owner}:{}", Uuid::new_v4());
self.run_with_timeout("redis lock acquire", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
let status = redis::cmd("SET")
.arg(&key.0)
.arg(&token)
.arg("NX")
.arg("PX")
.arg(ttl_ms)
.query_async::<Option<String>>(&mut connection)
.await
.map_redis_err()?;
Ok(status.map(|_| RedisLockLease {
key: key.clone(),
owner: owner.to_string(),
token,
ttl_ms,
}))
})
.await
}
pub async fn release(&self, lease: &RedisLockLease) -> Result<bool, DataLayerError> {
validate_lease(lease)?;
self.run_with_timeout("redis lock release", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
let deleted = redis::Script::new(
"if redis.call('get', KEYS[1]) == ARGV[1] then \
return redis.call('del', KEYS[1]) \
else \
return 0 \
end",
)
.key(&lease.key.0)
.arg(&lease.token)
.invoke_async::<i32>(&mut connection)
.await
.map_redis_err()?;
Ok(deleted > 0)
})
.await
}
pub async fn renew(
&self,
lease: &RedisLockLease,
ttl_ms: Option<u64>,
) -> Result<bool, DataLayerError> {
validate_lease(lease)?;
let ttl_ms = self.resolve_ttl_ms(ttl_ms)?;
self.run_with_timeout("redis lock renew", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
let renewed = redis::Script::new(
"if redis.call('get', KEYS[1]) == ARGV[1] then \
return redis.call('pexpire', KEYS[1], ARGV[2]) \
else \
return 0 \
end",
)
.key(&lease.key.0)
.arg(&lease.token)
.arg(ttl_ms)
.invoke_async::<i32>(&mut connection)
.await
.map_redis_err()?;
Ok(renewed > 0)
})
.await
}
async fn run_with_timeout<T, F>(
&self,
operation: &'static str,
future: F,
) -> Result<T, DataLayerError>
where
F: Future<Output = Result<T, DataLayerError>>,
{
if let Some(timeout_ms) = self.config.command_timeout_ms {
tokio::time::timeout(Duration::from_millis(timeout_ms), future)
.await
.map_err(|_| {
DataLayerError::TimedOut(format!("{operation} exceeded {timeout_ms}ms timeout"))
})?
} else {
future.await
}
}
fn resolve_ttl_ms(&self, ttl_ms: Option<u64>) -> Result<u64, DataLayerError> {
let ttl_ms = ttl_ms.unwrap_or(self.config.default_ttl_ms);
if ttl_ms == 0 {
return Err(DataLayerError::InvalidInput(
"redis lock ttl_ms must be positive".to_string(),
));
}
Ok(ttl_ms)
}
}
fn validate_owner(owner: &str) -> Result<(), DataLayerError> {
if owner.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis lock owner cannot be empty".to_string(),
));
}
Ok(())
}
fn validate_key(key: &RedisLockKey) -> Result<(), DataLayerError> {
if key.0.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis lock key cannot be empty".to_string(),
));
}
Ok(())
}
fn validate_lease(lease: &RedisLockLease) -> Result<(), DataLayerError> {
validate_key(&lease.key)?;
validate_owner(&lease.owner)?;
if lease.token.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis lock token cannot be empty".to_string(),
));
}
if lease.ttl_ms == 0 {
return Err(DataLayerError::InvalidInput(
"redis lock ttl_ms must be positive".to_string(),
));
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::{RedisLockKey, RedisLockLease, RedisLockRunner, RedisLockRunnerConfig};
use crate::driver::redis::{RedisClientConfig, RedisClientFactory};
fn sample_runner() -> RedisLockRunner {
let client = RedisClientFactory::new(RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
})
.expect("factory should build")
.connect_lazy()
.expect("client should build");
RedisLockRunner::new(
client,
RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
}
.keyspace(),
RedisLockRunnerConfig::default(),
)
.expect("runner should build")
}
#[test]
fn validates_runner_config() {
assert!(RedisLockRunnerConfig {
command_timeout_ms: Some(0),
..RedisLockRunnerConfig::default()
}
.validate()
.is_err());
assert!(RedisLockRunnerConfig {
default_ttl_ms: 0,
..RedisLockRunnerConfig::default()
}
.validate()
.is_err());
}
#[test]
fn runner_reuses_client_and_keyspace() {
let runner = sample_runner();
assert_eq!(runner.config(), RedisLockRunnerConfig::default());
assert_eq!(runner.keyspace().lock_key("poller").0, "aether:lock:poller");
let _client_ref = runner.client();
}
#[tokio::test]
async fn rejects_invalid_owner_or_lease_before_network() {
let runner = sample_runner();
assert!(runner
.try_acquire(&RedisLockKey("aether:lock:poller".to_string()), "", None)
.await
.is_err());
assert!(runner
.release(&RedisLockLease {
key: RedisLockKey("aether:lock:poller".to_string()),
owner: "worker-1".to_string(),
token: String::new(),
ttl_ms: 1_000,
})
.await
.is_err());
assert!(runner
.renew(
&RedisLockLease {
key: RedisLockKey("aether:lock:poller".to_string()),
owner: "worker-1".to_string(),
token: "token-1".to_string(),
ttl_ms: 1_000,
},
Some(0),
)
.await
.is_err());
}
}

View File

@@ -0,0 +1,14 @@
mod client;
mod kv;
mod lock;
mod namespace;
mod stream;
pub use client::{RedisClient, RedisClientConfig, RedisClientFactory};
pub use kv::{RedisKvRunner, RedisKvRunnerConfig};
pub use lock::{RedisLockKey, RedisLockLease, RedisLockRunner, RedisLockRunnerConfig};
pub use namespace::RedisKeyspace;
pub use stream::{
RedisConsumerGroup, RedisConsumerName, RedisStreamEntry, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunner, RedisStreamRunnerConfig,
};

View File

@@ -0,0 +1,43 @@
use aether_cache::CacheKeyNamespace;
use crate::driver::redis::{RedisLockKey, RedisStreamName};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisKeyspace {
namespace: CacheKeyNamespace,
}
impl RedisKeyspace {
pub fn new(prefix: Option<&str>) -> Self {
let normalized = prefix.unwrap_or_default().trim().trim_matches(':');
Self {
namespace: CacheKeyNamespace::new(normalized),
}
}
pub fn key(&self, raw_key: &str) -> String {
self.namespace.key(raw_key)
}
pub fn lock_key(&self, raw_key: &str) -> RedisLockKey {
RedisLockKey(self.namespace.child("lock").key(raw_key))
}
pub fn stream_name(&self, raw_name: &str) -> RedisStreamName {
RedisStreamName(self.namespace.child("stream").key(raw_name))
}
}
#[cfg(test)]
mod tests {
use super::RedisKeyspace;
#[test]
fn composes_prefixed_lock_and_stream_names() {
let keyspace = RedisKeyspace::new(Some("aether"));
assert_eq!(keyspace.key("auth:user"), "aether:auth:user");
assert_eq!(keyspace.lock_key("poller").0, "aether:lock:poller");
assert_eq!(keyspace.stream_name("audit").0, "aether:stream:audit");
}
}

View File

@@ -0,0 +1,747 @@
use std::collections::BTreeMap;
use std::future::Future;
use std::time::Duration;
use redis::from_redis_value;
use redis::streams::StreamReadReply;
use redis::Value as RedisValue;
use crate::driver::redis::{RedisClient, RedisKeyspace};
use crate::error::{redis_error, RedisResultExt};
use crate::DataLayerError;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RedisStreamName(pub String);
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RedisConsumerGroup(pub String);
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RedisConsumerName(pub String);
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisStreamEntry {
pub id: String,
pub fields: BTreeMap<String, String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisStreamReclaimResult {
pub next_start_id: String,
pub entries: Vec<RedisStreamEntry>,
pub deleted_ids: Vec<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RedisStreamReclaimConfig {
pub min_idle_ms: u64,
pub count: usize,
}
impl Default for RedisStreamReclaimConfig {
fn default() -> Self {
Self {
min_idle_ms: 60_000,
count: 32,
}
}
}
impl RedisStreamReclaimConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if self.min_idle_ms == 0 {
return Err(DataLayerError::InvalidConfiguration(
"redis stream reclaim min_idle_ms must be positive".to_string(),
));
}
if self.count == 0 {
return Err(DataLayerError::InvalidConfiguration(
"redis stream reclaim count must be positive".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RedisStreamRunnerConfig {
pub command_timeout_ms: Option<u64>,
pub read_block_ms: Option<u64>,
pub read_count: usize,
}
impl Default for RedisStreamRunnerConfig {
fn default() -> Self {
Self {
command_timeout_ms: Some(2_000),
read_block_ms: Some(1_000),
read_count: 32,
}
}
}
impl RedisStreamRunnerConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if matches!(self.command_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"redis stream command_timeout_ms must be positive".to_string(),
));
}
if matches!(self.read_block_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"redis stream read_block_ms must be positive".to_string(),
));
}
if let (Some(command_timeout_ms), Some(read_block_ms)) =
(self.command_timeout_ms, self.read_block_ms)
{
if command_timeout_ms <= read_block_ms {
return Err(DataLayerError::InvalidConfiguration(
"redis stream command_timeout_ms must be greater than read_block_ms"
.to_string(),
));
}
}
if self.read_count == 0 {
return Err(DataLayerError::InvalidConfiguration(
"redis stream read_count must be positive".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct RedisStreamRunner {
client: RedisClient,
keyspace: RedisKeyspace,
config: RedisStreamRunnerConfig,
}
impl RedisStreamRunner {
pub fn new(
client: RedisClient,
keyspace: RedisKeyspace,
config: RedisStreamRunnerConfig,
) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self {
client,
keyspace,
config,
})
}
pub fn client(&self) -> &RedisClient {
&self.client
}
pub fn keyspace(&self) -> &RedisKeyspace {
&self.keyspace
}
pub fn config(&self) -> RedisStreamRunnerConfig {
self.config
}
pub async fn ensure_consumer_group(
&self,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
start_id: &str,
) -> Result<(), DataLayerError> {
validate_stream_name(stream)?;
validate_group(group)?;
validate_stream_position(start_id)?;
self.run_with_timeout("redis stream ensure consumer group", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
let result = redis::cmd("XGROUP")
.arg("CREATE")
.arg(&stream.0)
.arg(&group.0)
.arg(start_id)
.arg("MKSTREAM")
.query_async::<String>(&mut connection)
.await;
match result {
Ok(_) => Ok(()),
Err(err) if err.code() == Some("BUSYGROUP") => Ok(()),
Err(err) => Err(redis_error(err)),
}
})
.await
}
pub async fn append_fields(
&self,
stream: &RedisStreamName,
fields: &BTreeMap<String, String>,
) -> Result<String, DataLayerError> {
self.append_fields_with_maxlen(stream, fields, None).await
}
pub async fn append_fields_with_maxlen(
&self,
stream: &RedisStreamName,
fields: &BTreeMap<String, String>,
maxlen: Option<usize>,
) -> Result<String, DataLayerError> {
validate_stream_name(stream)?;
if fields.is_empty() {
return Err(DataLayerError::InvalidInput(
"redis stream fields cannot be empty".to_string(),
));
}
self.run_with_timeout("redis stream append", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
let mut command = redis::cmd("XADD");
command.arg(&stream.0);
if let Some(maxlen) = maxlen.filter(|value| *value > 0) {
command.arg("MAXLEN").arg("~").arg(maxlen);
}
command.arg("*");
for (key, value) in fields {
command.arg(key).arg(value);
}
command
.query_async::<String>(&mut connection)
.await
.map_redis_err()
})
.await
}
pub async fn append_json(
&self,
stream: &RedisStreamName,
field: &str,
payload: &serde_json::Value,
) -> Result<String, DataLayerError> {
if field.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis stream json field cannot be empty".to_string(),
));
}
let mut fields = BTreeMap::new();
fields.insert(
field.to_string(),
serde_json::to_string(payload).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"failed to serialize redis stream payload: {err}"
))
})?,
);
self.append_fields(stream, &fields).await
}
pub async fn read_group(
&self,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
consumer: &RedisConsumerName,
) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
validate_stream_name(stream)?;
validate_group(group)?;
validate_consumer(consumer)?;
self.run_with_timeout("redis stream read group", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
let mut command = redis::cmd("XREADGROUP");
command
.arg("GROUP")
.arg(&group.0)
.arg(&consumer.0)
.arg("COUNT")
.arg(self.config.read_count);
if let Some(block_ms) = self.config.read_block_ms {
command.arg("BLOCK").arg(block_ms);
}
command.arg("STREAMS").arg(&stream.0).arg(">");
let reply = command
.query_async::<StreamReadReply>(&mut connection)
.await
.map_redis_err()?;
Ok(reply
.keys
.into_iter()
.flat_map(|key| key.ids.into_iter())
.map(|id| RedisStreamEntry {
id: id.id,
fields: id
.map
.into_iter()
.filter_map(|(field, value)| {
redis::from_redis_value::<String>(&value)
.ok()
.map(|text| (field, text))
})
.collect(),
})
.collect())
})
.await
}
pub async fn ack(
&self,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
ids: &[String],
) -> Result<usize, DataLayerError> {
validate_stream_name(stream)?;
validate_group(group)?;
if ids.is_empty() {
return Ok(0);
}
self.run_with_timeout("redis stream ack", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
let mut command = redis::cmd("XACK");
command.arg(&stream.0).arg(&group.0);
for id in ids {
command.arg(id);
}
command
.query_async::<usize>(&mut connection)
.await
.map_redis_err()
})
.await
}
pub async fn delete(
&self,
stream: &RedisStreamName,
ids: &[String],
) -> Result<usize, DataLayerError> {
validate_stream_name(stream)?;
if ids.is_empty() {
return Ok(0);
}
self.run_with_timeout("redis stream delete", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
let mut command = redis::cmd("XDEL");
command.arg(&stream.0);
for id in ids {
command.arg(id);
}
command
.query_async::<usize>(&mut connection)
.await
.map_redis_err()
})
.await
}
pub async fn claim_stale(
&self,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
consumer: &RedisConsumerName,
start_id: &str,
config: RedisStreamReclaimConfig,
) -> Result<RedisStreamReclaimResult, DataLayerError> {
validate_stream_name(stream)?;
validate_group(group)?;
validate_consumer(consumer)?;
validate_stream_position(start_id)?;
config.validate()?;
self.run_with_timeout("redis stream reclaim", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
let reply = redis::cmd("XAUTOCLAIM")
.arg(&stream.0)
.arg(&group.0)
.arg(&consumer.0)
.arg(config.min_idle_ms)
.arg(start_id)
.arg("COUNT")
.arg(config.count)
.query_async::<RedisValue>(&mut connection)
.await
.map_redis_err()?;
parse_reclaim_result(reply)
})
.await
}
async fn run_with_timeout<T, F>(
&self,
operation: &'static str,
future: F,
) -> Result<T, DataLayerError>
where
F: Future<Output = Result<T, DataLayerError>>,
{
if let Some(timeout_ms) = self.config.command_timeout_ms {
tokio::time::timeout(Duration::from_millis(timeout_ms), future)
.await
.map_err(|_| {
DataLayerError::TimedOut(format!("{operation} exceeded {timeout_ms}ms timeout"))
})?
} else {
future.await
}
}
}
fn validate_stream_name(stream: &RedisStreamName) -> Result<(), DataLayerError> {
if stream.0.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis stream name cannot be empty".to_string(),
));
}
Ok(())
}
fn validate_group(group: &RedisConsumerGroup) -> Result<(), DataLayerError> {
if group.0.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis consumer group cannot be empty".to_string(),
));
}
Ok(())
}
fn validate_consumer(consumer: &RedisConsumerName) -> Result<(), DataLayerError> {
if consumer.0.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis consumer name cannot be empty".to_string(),
));
}
Ok(())
}
fn validate_stream_position(position: &str) -> Result<(), DataLayerError> {
if position.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis stream position cannot be empty".to_string(),
));
}
Ok(())
}
fn parse_reclaim_result(value: RedisValue) -> Result<RedisStreamReclaimResult, DataLayerError> {
let RedisValue::Array(parts) = value else {
return Err(DataLayerError::UnexpectedValue(
"redis xautoclaim returned non-array payload".to_string(),
));
};
if parts.len() < 2 || parts.len() > 3 {
return Err(DataLayerError::UnexpectedValue(format!(
"redis xautoclaim returned {} top-level fields, expected 2 or 3",
parts.len()
)));
}
let next_start_id = parse_string_value(&parts[0], "redis xautoclaim next_start_id")?;
let entries = parse_reclaim_entries(&parts[1])?;
let deleted_ids = match parts.get(2) {
Some(value) => parse_string_array(value, "redis xautoclaim deleted_ids")?,
None => Vec::new(),
};
Ok(RedisStreamReclaimResult {
next_start_id,
entries,
deleted_ids,
})
}
fn parse_reclaim_entries(value: &RedisValue) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
match value {
RedisValue::Array(entries) => entries.iter().map(parse_reclaim_entry).collect(),
RedisValue::Nil => Ok(Vec::new()),
_ => Err(DataLayerError::UnexpectedValue(
"redis xautoclaim entries payload was not an array".to_string(),
)),
}
}
fn parse_reclaim_entry(value: &RedisValue) -> Result<RedisStreamEntry, DataLayerError> {
let RedisValue::Array(parts) = value else {
return Err(DataLayerError::UnexpectedValue(
"redis xautoclaim entry was not an array".to_string(),
));
};
if parts.len() != 2 {
return Err(DataLayerError::UnexpectedValue(format!(
"redis xautoclaim entry had {} fields, expected 2",
parts.len()
)));
}
let id = parse_string_value(&parts[0], "redis xautoclaim entry id")?;
let fields = parse_string_map(&parts[1], "redis xautoclaim entry fields")?;
Ok(RedisStreamEntry { id, fields })
}
fn parse_string_map(
value: &RedisValue,
context: &str,
) -> Result<BTreeMap<String, String>, DataLayerError> {
match value {
RedisValue::Array(values) => {
if values.len() % 2 != 0 {
return Err(DataLayerError::UnexpectedValue(format!(
"{context} expected an even number of field elements, got {}",
values.len()
)));
}
let mut fields = BTreeMap::new();
for pair in values.chunks(2) {
let key = parse_string_value(&pair[0], context)?;
let value = parse_string_value(&pair[1], context)?;
fields.insert(key, value);
}
Ok(fields)
}
RedisValue::Map(entries) => entries
.iter()
.map(|(key, value)| {
Ok((
parse_string_value(key, context)?,
parse_string_value(value, context)?,
))
})
.collect(),
RedisValue::Nil => Ok(BTreeMap::new()),
_ => Err(DataLayerError::UnexpectedValue(format!(
"{context} expected a redis array/map payload"
))),
}
}
fn parse_string_array(value: &RedisValue, context: &str) -> Result<Vec<String>, DataLayerError> {
match value {
RedisValue::Array(values) => values
.iter()
.map(|value| parse_string_value(value, context))
.collect(),
RedisValue::Nil => Ok(Vec::new()),
_ => Err(DataLayerError::UnexpectedValue(format!(
"{context} expected a redis array payload"
))),
}
}
fn parse_string_value(value: &RedisValue, context: &str) -> Result<String, DataLayerError> {
from_redis_value::<String>(value).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"{context} was not a string-compatible redis value: {err}"
))
})
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use super::{
parse_reclaim_result, RedisConsumerGroup, RedisConsumerName, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunner,
RedisStreamRunnerConfig,
};
use crate::driver::redis::{RedisClientConfig, RedisClientFactory};
use redis::Value as RedisValue;
fn sample_runner() -> RedisStreamRunner {
let config = RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
};
let client = RedisClientFactory::new(config.clone())
.expect("factory should build")
.connect_lazy()
.expect("client should build");
RedisStreamRunner::new(
client,
config.keyspace(),
RedisStreamRunnerConfig::default(),
)
.expect("runner should build")
}
#[test]
fn validates_stream_runner_config() {
assert!(RedisStreamRunnerConfig {
command_timeout_ms: Some(0),
..RedisStreamRunnerConfig::default()
}
.validate()
.is_err());
assert!(RedisStreamRunnerConfig {
read_block_ms: Some(0),
..RedisStreamRunnerConfig::default()
}
.validate()
.is_err());
assert!(RedisStreamRunnerConfig {
command_timeout_ms: Some(1_000),
read_block_ms: Some(1_000),
..RedisStreamRunnerConfig::default()
}
.validate()
.is_err());
assert!(RedisStreamRunnerConfig {
read_count: 0,
..RedisStreamRunnerConfig::default()
}
.validate()
.is_err());
}
#[test]
fn validates_reclaim_config() {
assert!(RedisStreamReclaimConfig {
min_idle_ms: 0,
..RedisStreamReclaimConfig::default()
}
.validate()
.is_err());
assert!(RedisStreamReclaimConfig {
count: 0,
..RedisStreamReclaimConfig::default()
}
.validate()
.is_err());
}
#[test]
fn runner_reuses_client_and_keyspace() {
let runner = sample_runner();
assert_eq!(runner.config(), RedisStreamRunnerConfig::default());
assert_eq!(
runner.keyspace().stream_name("audit").0,
"aether:stream:audit"
);
let _client_ref = runner.client();
}
#[tokio::test]
async fn rejects_invalid_inputs_before_network() {
let runner = sample_runner();
let stream = RedisStreamName("aether:stream:audit".to_string());
let group = RedisConsumerGroup("audit-workers".to_string());
let consumer = RedisConsumerName("worker-1".to_string());
assert!(runner
.ensure_consumer_group(&stream, &group, "")
.await
.is_err());
assert!(runner
.append_fields(&stream, &BTreeMap::new())
.await
.is_err());
assert!(runner
.append_json(&stream, "", &serde_json::json!({"ok": true}))
.await
.is_err());
assert!(runner
.read_group(&stream, &group, &RedisConsumerName(String::new()))
.await
.is_err());
assert_eq!(
runner.ack(&stream, &group, &[]).await.expect("empty ack"),
0
);
assert_eq!(runner.delete(&stream, &[]).await.expect("empty delete"), 0);
assert!(runner
.claim_stale(
&stream,
&group,
&consumer,
"",
RedisStreamReclaimConfig::default()
)
.await
.is_err());
let _ = consumer;
}
#[test]
fn parses_reclaim_result_with_deleted_ids() {
let parsed = parse_reclaim_result(RedisValue::Array(vec![
RedisValue::BulkString(b"0-0".to_vec()),
RedisValue::Array(vec![RedisValue::Array(vec![
RedisValue::BulkString(b"1710000000000-0".to_vec()),
RedisValue::Array(vec![
RedisValue::BulkString(b"payload".to_vec()),
RedisValue::BulkString(br#"{"ok":true}"#.to_vec()),
RedisValue::BulkString(b"kind".to_vec()),
RedisValue::BulkString(b"shadow".to_vec()),
]),
])]),
RedisValue::Array(vec![RedisValue::BulkString(b"1709999999999-0".to_vec())]),
]))
.expect("reclaim result should parse");
assert_eq!(
parsed,
RedisStreamReclaimResult {
next_start_id: "0-0".to_string(),
entries: vec![super::RedisStreamEntry {
id: "1710000000000-0".to_string(),
fields: BTreeMap::from([
("kind".to_string(), "shadow".to_string()),
("payload".to_string(), r#"{"ok":true}"#.to_string()),
]),
}],
deleted_ids: vec!["1709999999999-0".to_string()],
}
);
}
#[test]
fn parses_reclaim_result_without_deleted_ids() {
let parsed = parse_reclaim_result(RedisValue::Array(vec![
RedisValue::BulkString(b"0-0".to_vec()),
RedisValue::Array(vec![]),
]))
.expect("reclaim result should parse");
assert_eq!(
parsed,
RedisStreamReclaimResult {
next_start_id: "0-0".to_string(),
entries: Vec::new(),
deleted_ids: Vec::new(),
}
);
}
}

View File

@@ -0,0 +1,3 @@
mod pool;
pub use pool::{SqlitePool, SqlitePoolConfig, SqlitePoolFactory};

View File

@@ -0,0 +1,160 @@
use std::path::PathBuf;
use std::str::FromStr;
use std::time::Duration;
use crate::database::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig};
use crate::DataLayerError;
use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions};
use sqlx::SqlitePool as SqlxSqlitePool;
pub type SqlitePool = SqlxSqlitePool;
pub type SqlitePoolConfig = SqlDatabaseConfig;
#[derive(Debug, Clone)]
pub struct SqlitePoolFactory {
config: SqlitePoolConfig,
}
impl SqlitePoolFactory {
pub fn new(config: SqlitePoolConfig) -> Result<Self, DataLayerError> {
if config.driver != DatabaseDriver::Sqlite {
return Err(DataLayerError::InvalidConfiguration(format!(
"sqlite pool requires sqlite driver, got {}",
config.driver
)));
}
config.validate()?;
Ok(Self { config })
}
pub fn config(&self) -> &SqlitePoolConfig {
&self.config
}
pub fn connect_options(&self) -> Result<SqliteConnectOptions, DataLayerError> {
ensure_sqlite_parent_dir(self.config.url.trim())?;
let is_memory = is_sqlite_memory_url(self.config.url.trim());
SqliteConnectOptions::from_str(self.config.url.trim())
.map(|options| {
let options = options
.create_if_missing(true)
.foreign_keys(true)
.statement_cache_capacity(self.config.pool.statement_cache_capacity);
if is_memory {
options
} else {
options.journal_mode(SqliteJournalMode::Wal)
}
})
.map_err(|err| {
DataLayerError::InvalidConfiguration(format!("invalid sqlite database url: {err}"))
})
}
pub fn connect_lazy(&self) -> Result<SqlitePool, DataLayerError> {
let SqlPoolConfig {
min_connections,
max_connections,
acquire_timeout_ms,
idle_timeout_ms,
max_lifetime_ms,
..
} = self.config.pool;
Ok(SqlitePoolOptions::new()
.min_connections(min_connections)
.max_connections(max_connections)
.acquire_timeout(Duration::from_millis(acquire_timeout_ms))
.idle_timeout(Duration::from_millis(idle_timeout_ms))
.max_lifetime(Duration::from_millis(max_lifetime_ms))
.connect_lazy_with(self.connect_options()?))
}
}
fn is_sqlite_memory_url(url: &str) -> bool {
matches!(url.trim(), "sqlite::memory:" | "sqlite://:memory:")
}
fn sqlite_file_path_from_url(url: &str) -> Option<PathBuf> {
let url = url.trim();
if is_sqlite_memory_url(url) {
return None;
}
let path = url
.strip_prefix("sqlite://")
.or_else(|| url.strip_prefix("sqlite:"))?;
if path.is_empty() {
return None;
}
Some(PathBuf::from(path))
}
fn ensure_sqlite_parent_dir(url: &str) -> Result<(), DataLayerError> {
let Some(path) = sqlite_file_path_from_url(url) else {
return Ok(());
};
let Some(parent) = path
.parent()
.filter(|parent| !parent.as_os_str().is_empty())
else {
return Ok(());
};
std::fs::create_dir_all(parent).map_err(|err| {
DataLayerError::InvalidConfiguration(format!(
"failed to create sqlite database parent directory '{}': {err}",
parent.display()
))
})
}
#[cfg(test)]
mod tests {
use super::SqlitePoolFactory;
use crate::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig};
use std::path::PathBuf;
#[tokio::test]
async fn factory_builds_lazy_pool_from_valid_config() {
let config = SqlDatabaseConfig {
driver: DatabaseDriver::Sqlite,
url: "sqlite://./data/aether.db".to_string(),
pool: SqlPoolConfig {
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
},
};
let factory = SqlitePoolFactory::new(config).expect("factory should build");
let _pool = factory.connect_lazy().expect("lazy pool should build");
}
#[tokio::test]
async fn factory_creates_parent_directory_for_file_database() {
let db_path = unique_temp_db_path();
let parent = db_path.parent().expect("temp db path should have parent");
let _ = std::fs::remove_dir_all(parent);
let config = SqlDatabaseConfig {
driver: DatabaseDriver::Sqlite,
url: format!("sqlite://{}", db_path.display()),
pool: SqlPoolConfig::default(),
};
let factory = SqlitePoolFactory::new(config).expect("factory should build");
let _pool = factory.connect_lazy().expect("lazy pool should build");
assert!(parent.exists());
let _ = std::fs::remove_dir_all(parent);
}
fn unique_temp_db_path() -> PathBuf {
std::env::temp_dir()
.join(format!("aether-sqlite-{}", uuid::Uuid::new_v4()))
.join("nested")
.join("aether.db")
}
}