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