sova_db/
handle.rs

1use crate::DbError;
2use sea_orm::{
3    ConnectionTrait, DatabaseConnection, DatabaseTransaction, DbBackend, DbErr, ExecResult,
4    QueryResult, Statement,
5};
6use sova_core::Request;
7use std::sync::{Arc, RwLock};
8
9/// Shared pool handle filled during `on_startup`.
10#[derive(Clone, Default)]
11pub struct DbPool {
12    inner: Arc<RwLock<Option<DatabaseConnection>>>,
13}
14
15impl DbPool {
16    pub fn new() -> Self {
17        Self::default()
18    }
19
20    pub fn set(&self, conn: DatabaseConnection) {
21        *self.inner.write().unwrap() = Some(conn);
22    }
23
24    pub fn get(&self) -> Result<DatabaseConnection, DbError> {
25        self.inner
26            .read()
27            .unwrap()
28            .clone()
29            .ok_or_else(|| DbError(DbErr::Custom("database not connected".into())))
30    }
31
32    pub fn clear(&self) {
33        let _ = self.inner.write().unwrap().take();
34    }
35}
36
37/// Request-scoped DB handle: pool connection or open transaction.
38#[derive(Clone)]
39pub enum DbHandle {
40    Conn(DatabaseConnection),
41    Tx(Arc<DatabaseTransaction>),
42}
43
44impl DbHandle {
45    pub fn as_conn(&self) -> Option<&DatabaseConnection> {
46        match self {
47            Self::Conn(c) => Some(c),
48            Self::Tx(_) => None,
49        }
50    }
51}
52
53#[async_trait::async_trait]
54impl ConnectionTrait for DbHandle {
55    fn get_database_backend(&self) -> DbBackend {
56        match self {
57            Self::Conn(c) => c.get_database_backend(),
58            Self::Tx(t) => t.get_database_backend(),
59        }
60    }
61
62    async fn execute_raw(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
63        let sql = stmt.sql.clone();
64        let started = std::time::Instant::now();
65        let result = match self {
66            Self::Conn(c) => c.execute_raw(stmt).await,
67            Self::Tx(t) => t.execute_raw(stmt).await,
68        };
69        crate::trace::log_query(&sql, started, result.is_ok());
70        result
71    }
72
73    async fn execute_unprepared(&self, sql: &str) -> Result<ExecResult, DbErr> {
74        let started = std::time::Instant::now();
75        let result = match self {
76            Self::Conn(c) => c.execute_unprepared(sql).await,
77            Self::Tx(t) => t.execute_unprepared(sql).await,
78        };
79        crate::trace::log_query(sql, started, result.is_ok());
80        result
81    }
82
83    async fn query_one_raw(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
84        let sql = stmt.sql.clone();
85        let started = std::time::Instant::now();
86        let result = match self {
87            Self::Conn(c) => c.query_one_raw(stmt).await,
88            Self::Tx(t) => t.query_one_raw(stmt).await,
89        };
90        crate::trace::log_query(&sql, started, result.is_ok());
91        result
92    }
93
94    async fn query_all_raw(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
95        let sql = stmt.sql.clone();
96        let started = std::time::Instant::now();
97        let result = match self {
98            Self::Conn(c) => c.query_all_raw(stmt).await,
99            Self::Tx(t) => t.query_all_raw(stmt).await,
100        };
101        crate::trace::log_query(&sql, started, result.is_ok());
102        result
103    }
104
105    fn support_returning(&self) -> bool {
106        match self {
107            Self::Conn(c) => c.support_returning(),
108            Self::Tx(t) => t.support_returning(),
109        }
110    }
111}
112
113/// Convenient access to the request [`DbHandle`].
114pub trait DbExt {
115    fn db(&self) -> &DbHandle;
116}
117
118impl DbExt for Request {
119    fn db(&self) -> &DbHandle {
120        self.get::<DbHandle>()
121            .expect("Db plugin is not installed (missing req.db())")
122    }
123}
sova_db/handle.rs

sova_db/
handle.rs
