1use actix_web::{
14 dev::{forward_ready, Service, ServiceRequest, ServiceResponse, Transform},
15 web, FromRequest, HttpMessage, HttpRequest, HttpResponse, Responder,
16};
17use serde::Serialize;
18use std::{
19 future::{ready, Ready},
20 rc::Rc,
21 sync::Arc,
22};
23use sz_orm_core::{Pool, PooledConnection, QueryRows, Value};
24use tokio::sync::{Mutex, MutexGuard};
25
26#[derive(Clone)]
35pub struct PoolState {
36 pool: Arc<Pool>,
37}
38
39impl PoolState {
40 pub fn new(pool: Pool) -> Self {
42 Self {
43 pool: Arc::new(pool),
44 }
45 }
46
47 pub fn from_arc(pool: Arc<Pool>) -> Self {
49 Self { pool }
50 }
51
52 pub fn pool(&self) -> &Pool {
54 &self.pool
55 }
56}
57
58impl FromRequest for PoolState {
59 type Error = actix_web::Error;
60 type Future = Ready<Result<Self, Self::Error>>;
61
62 fn from_request(req: &HttpRequest, _: &mut actix_web::dev::Payload) -> Self::Future {
63 if let Some(state) = req.app_data::<web::Data<PoolState>>() {
67 return ready(Ok(state.get_ref().clone()));
68 }
69 if let Some(pool) = req.app_data::<web::Data<Arc<Pool>>>() {
71 return ready(Ok(PoolState::from_arc(pool.get_ref().clone())));
72 }
73 ready(Err(actix_web::error::ErrorInternalServerError(
74 "PoolState not found in app data",
75 )))
76 }
77}
78
79pub struct JsonRows(pub QueryRows);
87
88impl Responder for JsonRows {
89 type Body = actix_web::body::BoxBody;
90
91 fn respond_to(self, _: &HttpRequest) -> HttpResponse {
92 let json: Vec<serde_json::Value> = self
93 .0
94 .iter()
95 .map(|row| {
96 let mut map = serde_json::Map::new();
97 for (k, v) in row {
98 map.insert(k.clone(), value_to_json(v));
99 }
100 serde_json::Value::Object(map)
101 })
102 .collect();
103 HttpResponse::Ok().json(json)
104 }
105}
106
107fn value_to_json(v: &Value) -> serde_json::Value {
112 match v {
113 Value::Null => serde_json::Value::Null,
114 Value::Bool(b) => serde_json::Value::Bool(*b),
115 Value::I8(n) => (*n).into(),
116 Value::I16(n) => (*n).into(),
117 Value::I32(n) => (*n).into(),
118 Value::I64(n) => (*n).into(),
119 Value::U8(n) => (*n).into(),
120 Value::U16(n) => (*n).into(),
121 Value::U32(n) => (*n).into(),
122 Value::U64(n) => (*n).into(),
123 Value::F32(f) => serde_json::Value::from(*f),
124 Value::F64(f) => serde_json::Value::from(*f),
125 Value::String(s) => serde_json::Value::String(s.clone()),
126 Value::Decimal(s) => serde_json::Value::String(s.clone()),
127 Value::Bytes(b) => {
129 serde_json::Value::String(b.iter().map(|byte| format!("{:02x}", byte)).collect())
130 }
131 Value::Date(s) | Value::DateTime(s) | Value::Time(s) => {
132 serde_json::Value::String(s.clone())
133 }
134 Value::Json(s) => {
135 serde_json::from_str(s).unwrap_or_else(|_| serde_json::Value::String(s.clone()))
136 }
137 Value::Uuid(s) => serde_json::Value::String(s.clone()),
138 Value::Array(arr) => serde_json::Value::Array(arr.iter().map(value_to_json).collect()),
139 Value::Object(map) => {
140 let mut obj = serde_json::Map::new();
141 for (k, v) in map {
142 obj.insert(k.clone(), value_to_json(v));
143 }
144 serde_json::Value::Object(obj)
145 }
146 _ => serde_json::Value::Null,
148 }
149}
150
151pub struct JsonResp<T: Serialize>(pub T);
159
160impl<T: Serialize> Responder for JsonResp<T> {
161 type Body = actix_web::body::BoxBody;
162
163 fn respond_to(self, _: &HttpRequest) -> HttpResponse {
164 match serde_json::to_value(&self.0) {
165 Ok(v) => HttpResponse::Ok().json(v),
166 Err(e) => HttpResponse::InternalServerError().body(format!("JSON 序列化失败: {}", e)),
167 }
168 }
169}
170
171pub struct TransactionConn {
200 inner: Arc<Mutex<Option<PooledConnection>>>,
201}
202
203impl TransactionConn {
204 fn new(conn: PooledConnection) -> Self {
206 Self {
207 inner: Arc::new(Mutex::new(Some(conn))),
208 }
209 }
210
211 pub async fn conn(&self) -> Option<MutexGuard<'_, Option<PooledConnection>>> {
219 let guard = self.inner.lock().await;
220 if guard.is_none() {
221 return None;
222 }
223 Some(guard)
224 }
225}
226
227impl Clone for TransactionConn {
228 fn clone(&self) -> Self {
229 Self {
230 inner: Arc::clone(&self.inner),
231 }
232 }
233}
234
235impl FromRequest for TransactionConn {
236 type Error = actix_web::Error;
237 type Future = Ready<Result<Self, Self::Error>>;
238
239 fn from_request(req: &HttpRequest, _: &mut actix_web::dev::Payload) -> Self::Future {
240 if let Some(tx) = req.extensions().get::<TransactionConn>() {
241 return ready(Ok(tx.clone()));
242 }
243 ready(Err(actix_web::error::ErrorInternalServerError(
244 "TransactionConn not found in request extensions. \
245 Is TransactionMiddleware registered?",
246 )))
247 }
248}
249
250pub struct TransactionMiddleware;
271
272impl<S, B> Transform<S, ServiceRequest> for TransactionMiddleware
273where
274 S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error> + 'static,
275 S::Future: 'static,
276 B: 'static,
277{
278 type Response = ServiceResponse<B>;
279 type Error = actix_web::Error;
280 type Transform = TransactionMiddlewareService<S>;
281 type InitError = ();
282 type Future = Ready<Result<Self::Transform, Self::InitError>>;
283
284 fn new_transform(&self, service: S) -> Self::Future {
285 ready(Ok(TransactionMiddlewareService {
286 service: Rc::new(service),
287 }))
288 }
289}
290
291pub struct TransactionMiddlewareService<S> {
296 service: Rc<S>,
297}
298
299impl<S, B> Service<ServiceRequest> for TransactionMiddlewareService<S>
300where
301 S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error> + 'static,
302 S::Future: 'static,
303 B: 'static,
304{
305 type Response = ServiceResponse<B>;
306 type Error = actix_web::Error;
307 type Future =
308 std::pin::Pin<Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>>>>;
309
310 forward_ready!(service);
311
312 fn call(&self, req: ServiceRequest) -> Self::Future {
313 let pool_state = req
315 .app_data::<web::Data<PoolState>>()
316 .map(|s| s.get_ref().clone());
317
318 let svc = Rc::clone(&self.service);
320
321 Box::pin(async move {
322 let pool_state = match pool_state {
324 Some(state) => state,
325 None => {
326 tracing::warn!(
327 target: "sz_orm_actix::transaction",
328 "PoolState not found in app_data, TransactionMiddleware degrades to passthrough"
329 );
330 return svc.call(req).await;
331 }
332 };
333
334 let mut conn = match pool_state.pool().acquire().await {
336 Ok(c) => c,
337 Err(e) => {
338 tracing::warn!(
339 target: "sz_orm_actix::transaction",
340 error = %e,
341 "acquire connection failed, TransactionMiddleware degrades to passthrough"
342 );
343 return svc.call(req).await;
344 }
345 };
346
347 if let Err(e) = conn.begin_transaction().await {
349 tracing::warn!(
350 target: "sz_orm_actix::transaction",
351 error = %e,
352 "begin_transaction failed, TransactionMiddleware degrades to passthrough"
353 );
354 return svc.call(req).await;
356 }
357
358 let tx_conn = TransactionConn::new(conn);
360 let tx_clone = tx_conn.clone(); req.extensions_mut().insert(tx_conn);
362
363 let res = svc.call(req).await?;
365
366 let mut guard = tx_clone.inner.lock().await;
368 if let Some(mut conn) = guard.take() {
369 let tx_result = if res.status().is_success() {
370 conn.commit().await
371 } else {
372 conn.rollback().await
373 };
374 if let Err(e) = tx_result {
375 tracing::error!(
376 target: "sz_orm_actix::transaction",
377 error = %e,
378 status = %res.status(),
379 "transaction commit/rollback failed"
380 );
381 }
383 } else {
384 tracing::debug!(
385 target: "sz_orm_actix::transaction",
386 "TransactionConn was None after service call (handler may have dropped it)"
387 );
388 }
389
390 Ok(res)
391 })
392 }
393}
394
395#[cfg(test)]
396mod tests {
397 use super::*;
398 use std::collections::HashMap;
399
400 #[test]
401 fn test_pool_state_clone() {
402 fn _assert_clone<T: Clone>() {}
403 _assert_clone::<PoolState>();
404 }
405
406 #[test]
407 fn test_value_to_json_variants() {
408 assert_eq!(value_to_json(&Value::Null), serde_json::Value::Null);
409 assert_eq!(
410 value_to_json(&Value::Bool(true)),
411 serde_json::Value::Bool(true)
412 );
413 assert_eq!(value_to_json(&Value::I64(42)), serde_json::json!(42));
414 assert_eq!(
415 value_to_json(&Value::String("hi".into())),
416 serde_json::json!("hi")
417 );
418 assert_eq!(
419 value_to_json(&Value::Bytes(vec![0x1a, 0x2b])),
420 serde_json::json!("1a2b")
421 );
422 assert_eq!(
424 value_to_json(&Value::Json("{\"k\":1}".into())),
425 serde_json::json!({"k": 1})
426 );
427 }
428
429 #[test]
430 fn test_json_rows_responder() {
431 let mut row = HashMap::new();
432 row.insert("id".to_string(), Value::I64(1));
433 row.insert("name".to_string(), Value::String("Alice".into()));
434 let rows: QueryRows = vec![row];
435 let req = actix_web::test::TestRequest::default().to_http_request();
438 let resp = JsonRows(rows).respond_to(&req);
439 assert_eq!(resp.status(), actix_web::http::StatusCode::OK);
440 }
441
442 #[test]
443 fn test_json_resp_responder() {
444 #[derive(Serialize)]
445 struct User {
446 id: i64,
447 name: String,
448 }
449 let user = User {
450 id: 1,
451 name: "Bob".into(),
452 };
453 let req = actix_web::test::TestRequest::default().to_http_request();
454 let resp = JsonResp(user).respond_to(&req);
455 assert_eq!(resp.status(), actix_web::http::StatusCode::OK);
456 }
457}