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