mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
feat: 引入 aether-runtime/cache/data/http/testkit 基础 crate，完善并发门控与审计系统
新增 crate： - aether-runtime: 服务运行时基础设施（并发门控、分布式并发、指标、队列、优雅关闭、tracing） - aether-cache: 通用 TTL 缓存与命名空间抽象 - aether-data: 数据访问层（PostgreSQL/Redis 后端、repository 模式） - aether-http: HTTP 客户端封装（重试、配置） - aether-testkit: 集成测试工具集（gateway/executor/hub/proxy fixture、等待、负载测试） gateway 扩展： - 引入 audit 模块（shadow 执行审计、决策链路追踪、请求审计 bundle） - 引入 cache 模块（AuthContext 缓存、direct-plan bypass 缓存） - 引入 data 模块（auth/candidates/config/usage/video_tasks 数据访问） - 集成 ConcurrencyGate/DistributedConcurrencyGate 请求门控 - 新增本地 auth 拒绝、过载响应构建器 - 补充 control/auth_cache/video/concurrency 集成测试 aether-proxy 扩展： - AppState 集成 stream_gate / distributed_stream_gate 并发门控 - 新增 ProxyAdmissionError 及准入拒绝流程 - stream_handler 补充门控饱和/不可用场景测试 - 配置与注册客户端逻辑完善 aether-hub 扩展： - main.rs 引入运行时初始化、指标端点、健康检查 - local_relay 重构为 lib.rs 暴露公共接口
This commit is contained in:
fawney19
417
crates/aether-data/src/postgres/lease.rs
Normal file
417
crates/aether-data/src/postgres/lease.rs
Normal file
@@ -0,0 +1,417 @@
|
||||
use crate::postgres::{DatabaseRecordId, PostgresTransactionOptions, PostgresTransactionRunner};
|
||||
use crate::DataLayerError;
|
||||
use futures_util::FutureExt;
|
||||
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 rows = query_scalar::<_, String>(&sql)
|
||||
.bind(owner)
|
||||
.bind(lease_ms)
|
||||
.fetch_all(&mut **tx)
|
||||
.await?;
|
||||
Ok(rows.into_iter().map(DatabaseRecordId).collect())
|
||||
}
|
||||
.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 rows = query_scalar::<_, String>(&sql)
|
||||
.bind(ids)
|
||||
.bind(owner)
|
||||
.fetch_all(&mut **tx)
|
||||
.await?;
|
||||
Ok(rows.into_iter().map(DatabaseRecordId).collect())
|
||||
}
|
||||
.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 rows = query_scalar::<_, String>(&sql)
|
||||
.bind(ids)
|
||||
.bind(owner)
|
||||
.bind(lease_ms)
|
||||
.fetch_all(&mut **tx)
|
||||
.await?;
|
||||
Ok(rows.into_iter().map(DatabaseRecordId).collect())
|
||||
}
|
||||
.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::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());
|
||||
}
|
||||
}
|
||||
15
crates/aether-data/src/postgres/mod.rs
Normal file
15
crates/aether-data/src/postgres/mod.rs
Normal 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;
|
||||
122
crates/aether-data/src/postgres/pool.rs
Normal file
122
crates/aether-data/src/postgres/pool.rs
Normal 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: 5_000,
|
||||
idle_timeout_ms: 60_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");
|
||||
}
|
||||
}
|
||||
202
crates/aether-data/src/postgres/tx.rs
Normal file
202
crates/aether-data/src/postgres/tx.rs
Normal file
@@ -0,0 +1,202 @@
|
||||
use futures_util::future::BoxFuture;
|
||||
use sqlx::{Postgres, Transaction};
|
||||
|
||||
use crate::postgres::PostgresPool;
|
||||
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?;
|
||||
for statement in build_transaction_setup_statements(options) {
|
||||
sqlx::query(statement.as_str()).execute(&mut *tx).await?;
|
||||
}
|
||||
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?;
|
||||
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::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();
|
||||
}
|
||||
}
|
||||
2
crates/aether-data/src/postgres/types.rs
Normal file
2
crates/aether-data/src/postgres/types.rs
Normal file
@@ -0,0 +1,2 @@
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
pub struct DatabaseRecordId(pub String);
|
||||
Reference in New Issue
Block a user