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#[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#[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
92pub 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}