Skip to main content

sova_db/
handle.rs

1use crate::DbError;
2use sova_core::Request;
3use sea_orm::{
4    ConnectionTrait, DatabaseConnection, DatabaseTransaction, DbBackend, DbErr, ExecResult,
5    QueryResult, Statement,
6};
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        match self {
64            Self::Conn(c) => c.execute_raw(stmt).await,
65            Self::Tx(t) => t.execute_raw(stmt).await,
66        }
67    }
68
69    async fn execute_unprepared(&self, sql: &str) -> Result<ExecResult, DbErr> {
70        match self {
71            Self::Conn(c) => c.execute_unprepared(sql).await,
72            Self::Tx(t) => t.execute_unprepared(sql).await,
73        }
74    }
75
76    async fn query_one_raw(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
77        match self {
78            Self::Conn(c) => c.query_one_raw(stmt).await,
79            Self::Tx(t) => t.query_one_raw(stmt).await,
80        }
81    }
82
83    async fn query_all_raw(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
84        match self {
85            Self::Conn(c) => c.query_all_raw(stmt).await,
86            Self::Tx(t) => t.query_all_raw(stmt).await,
87        }
88    }
89
90    fn support_returning(&self) -> bool {
91        match self {
92            Self::Conn(c) => c.support_returning(),
93            Self::Tx(t) => t.support_returning(),
94        }
95    }
96}
97
98/// Convenient access to the request [`DbHandle`].
99pub trait DbExt {
100    fn db(&self) -> &DbHandle;
101}
102
103impl DbExt for Request {
104    fn db(&self) -> &DbHandle {
105        self.get::<DbHandle>()
106            .expect("Db plugin is not installed (missing req.db())")
107    }
108}