use actix_web::{
dev::{forward_ready, Service, ServiceRequest, ServiceResponse, Transform},
web, FromRequest, HttpMessage, HttpRequest, HttpResponse, Responder,
};
use serde::Serialize;
use std::{
future::{ready, Ready},
rc::Rc,
sync::Arc,
};
use sz_orm_core::{Pool, PooledConnection, QueryRows, Value};
use tokio::sync::{Mutex, MutexGuard};
#[derive(Clone)]
pub struct PoolState {
pool: Arc<Pool>,
}
impl PoolState {
pub fn new(pool: Pool) -> Self {
Self {
pool: Arc::new(pool),
}
}
pub fn from_arc(pool: Arc<Pool>) -> Self {
Self { pool }
}
pub fn pool(&self) -> &Pool {
&self.pool
}
}
impl FromRequest for PoolState {
type Error = actix_web::Error;
type Future = Ready<Result<Self, Self::Error>>;
fn from_request(req: &HttpRequest, _: &mut actix_web::dev::Payload) -> Self::Future {
if let Some(state) = req.app_data::<web::Data<PoolState>>() {
return ready(Ok(state.get_ref().clone()));
}
if let Some(pool) = req.app_data::<web::Data<Arc<Pool>>>() {
return ready(Ok(PoolState::from_arc(pool.get_ref().clone())));
}
ready(Err(actix_web::error::ErrorInternalServerError(
"PoolState not found in app data",
)))
}
}
pub struct JsonRows(pub QueryRows);
impl Responder for JsonRows {
type Body = actix_web::body::BoxBody;
fn respond_to(self, _: &HttpRequest) -> HttpResponse {
let json: Vec<serde_json::Value> = self
.0
.iter()
.map(|row| {
let mut map = serde_json::Map::new();
for (k, v) in row {
map.insert(k.clone(), value_to_json(v));
}
serde_json::Value::Object(map)
})
.collect();
HttpResponse::Ok().json(json)
}
}
fn value_to_json(v: &Value) -> serde_json::Value {
match v {
Value::Null => serde_json::Value::Null,
Value::Bool(b) => serde_json::Value::Bool(*b),
Value::I8(n) => (*n).into(),
Value::I16(n) => (*n).into(),
Value::I32(n) => (*n).into(),
Value::I64(n) => (*n).into(),
Value::U8(n) => (*n).into(),
Value::U16(n) => (*n).into(),
Value::U32(n) => (*n).into(),
Value::U64(n) => (*n).into(),
Value::F32(f) => serde_json::Value::from(*f),
Value::F64(f) => serde_json::Value::from(*f),
Value::String(s) => serde_json::Value::String(s.clone()),
Value::Decimal(s) => serde_json::Value::String(s.clone()),
Value::Bytes(b) => {
serde_json::Value::String(b.iter().map(|byte| format!("{:02x}", byte)).collect())
}
Value::Date(s) | Value::DateTime(s) | Value::Time(s) => {
serde_json::Value::String(s.clone())
}
Value::Json(s) => {
serde_json::from_str(s).unwrap_or_else(|_| serde_json::Value::String(s.clone()))
}
Value::Uuid(s) => serde_json::Value::String(s.clone()),
Value::Array(arr) => serde_json::Value::Array(arr.iter().map(value_to_json).collect()),
Value::Object(map) => {
let mut obj = serde_json::Map::new();
for (k, v) in map {
obj.insert(k.clone(), value_to_json(v));
}
serde_json::Value::Object(obj)
}
_ => serde_json::Value::Null,
}
}
pub struct JsonResp<T: Serialize>(pub T);
impl<T: Serialize> Responder for JsonResp<T> {
type Body = actix_web::body::BoxBody;
fn respond_to(self, _: &HttpRequest) -> HttpResponse {
match serde_json::to_value(&self.0) {
Ok(v) => HttpResponse::Ok().json(v),
Err(e) => HttpResponse::InternalServerError().body(format!("JSON 序列化失败: {}", e)),
}
}
}
pub struct TransactionConn {
inner: Arc<Mutex<Option<PooledConnection>>>,
}
impl TransactionConn {
fn new(conn: PooledConnection) -> Self {
Self {
inner: Arc::new(Mutex::new(Some(conn))),
}
}
pub async fn conn(&self) -> Option<MutexGuard<'_, Option<PooledConnection>>> {
let guard = self.inner.lock().await;
if guard.is_none() {
return None;
}
Some(guard)
}
}
impl Clone for TransactionConn {
fn clone(&self) -> Self {
Self {
inner: Arc::clone(&self.inner),
}
}
}
impl FromRequest for TransactionConn {
type Error = actix_web::Error;
type Future = Ready<Result<Self, Self::Error>>;
fn from_request(req: &HttpRequest, _: &mut actix_web::dev::Payload) -> Self::Future {
if let Some(tx) = req.extensions().get::<TransactionConn>() {
return ready(Ok(tx.clone()));
}
ready(Err(actix_web::error::ErrorInternalServerError(
"TransactionConn not found in request extensions. \
Is TransactionMiddleware registered?",
)))
}
}
pub struct TransactionMiddleware;
impl<S, B> Transform<S, ServiceRequest> for TransactionMiddleware
where
S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error> + 'static,
S::Future: 'static,
B: 'static,
{
type Response = ServiceResponse<B>;
type Error = actix_web::Error;
type Transform = TransactionMiddlewareService<S>;
type InitError = ();
type Future = Ready<Result<Self::Transform, Self::InitError>>;
fn new_transform(&self, service: S) -> Self::Future {
ready(Ok(TransactionMiddlewareService {
service: Rc::new(service),
}))
}
}
pub struct TransactionMiddlewareService<S> {
service: Rc<S>,
}
impl<S, B> Service<ServiceRequest> for TransactionMiddlewareService<S>
where
S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error> + 'static,
S::Future: 'static,
B: 'static,
{
type Response = ServiceResponse<B>;
type Error = actix_web::Error;
type Future =
std::pin::Pin<Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>>>>;
forward_ready!(service);
fn call(&self, req: ServiceRequest) -> Self::Future {
let pool_state = req
.app_data::<web::Data<PoolState>>()
.map(|s| s.get_ref().clone());
let svc = Rc::clone(&self.service);
Box::pin(async move {
let pool_state = match pool_state {
Some(state) => state,
None => {
tracing::warn!(
target: "sz_orm_actix::transaction",
"PoolState not found in app_data, TransactionMiddleware degrades to passthrough"
);
return svc.call(req).await;
}
};
let mut conn = match pool_state.pool().acquire().await {
Ok(c) => c,
Err(e) => {
tracing::warn!(
target: "sz_orm_actix::transaction",
error = %e,
"acquire connection failed, TransactionMiddleware degrades to passthrough"
);
return svc.call(req).await;
}
};
if let Err(e) = conn.begin_transaction().await {
tracing::warn!(
target: "sz_orm_actix::transaction",
error = %e,
"begin_transaction failed, TransactionMiddleware degrades to passthrough"
);
return svc.call(req).await;
}
let tx_conn = TransactionConn::new(conn);
let tx_clone = tx_conn.clone(); req.extensions_mut().insert(tx_conn);
let res = svc.call(req).await?;
let mut guard = tx_clone.inner.lock().await;
if let Some(mut conn) = guard.take() {
let tx_result = if res.status().is_success() {
conn.commit().await
} else {
conn.rollback().await
};
if let Err(e) = tx_result {
tracing::error!(
target: "sz_orm_actix::transaction",
error = %e,
status = %res.status(),
"transaction commit/rollback failed"
);
}
} else {
tracing::debug!(
target: "sz_orm_actix::transaction",
"TransactionConn was None after service call (handler may have dropped it)"
);
}
Ok(res)
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
#[test]
fn test_pool_state_clone() {
fn _assert_clone<T: Clone>() {}
_assert_clone::<PoolState>();
}
#[test]
fn test_value_to_json_variants() {
assert_eq!(value_to_json(&Value::Null), serde_json::Value::Null);
assert_eq!(
value_to_json(&Value::Bool(true)),
serde_json::Value::Bool(true)
);
assert_eq!(value_to_json(&Value::I64(42)), serde_json::json!(42));
assert_eq!(
value_to_json(&Value::String("hi".into())),
serde_json::json!("hi")
);
assert_eq!(
value_to_json(&Value::Bytes(vec![0x1a, 0x2b])),
serde_json::json!("1a2b")
);
assert_eq!(
value_to_json(&Value::Json("{\"k\":1}".into())),
serde_json::json!({"k": 1})
);
}
#[test]
fn test_json_rows_responder() {
let mut row = HashMap::new();
row.insert("id".to_string(), Value::I64(1));
row.insert("name".to_string(), Value::String("Alice".into()));
let rows: QueryRows = vec![row];
let req = actix_web::test::TestRequest::default().to_http_request();
let resp = JsonRows(rows).respond_to(&req);
assert_eq!(resp.status(), actix_web::http::StatusCode::OK);
}
#[test]
fn test_json_resp_responder() {
#[derive(Serialize)]
struct User {
id: i64,
name: String,
}
let user = User {
id: 1,
name: "Bob".into(),
};
let req = actix_web::test::TestRequest::default().to_http_request();
let resp = JsonResp(user).respond_to(&req);
assert_eq!(resp.status(), actix_web::http::StatusCode::OK);
}
}