1use crate::{
2 AsQuery, Connection, Driver, Entity, Error, Executor, Query, QueryResult, Result, Row,
3 RowsAffected,
4};
5use deadpool::managed::{Manager, Metrics, Object, Pool, RecycleResult, Timeouts};
6use futures::{FutureExt, Stream, future::BoxFuture};
7use std::{
8 borrow::Cow,
9 fmt::Debug,
10 future,
11 ops::{Deref, DerefMut},
12 time::Duration,
13};
14
15#[derive(Debug)]
16pub struct DBConnectionManager<D: Driver> {
17 driver: D,
18 url: Cow<'static, str>,
19}
20
21impl<D: Driver> DBConnectionManager<D> {
22 pub fn new(driver: D, url: Cow<'static, str>) -> Self {
23 Self { driver, url }
24 }
25}
26
27impl<D: Driver> Manager for DBConnectionManager<D> {
28 type Type = D::Connection;
29 type Error = Error;
30 async fn create(&self) -> Result<Self::Type> {
31 Ok(D::Connection::connect(&self.driver, self.url.clone()).await?)
32 }
33 fn recycle(
34 &self,
35 _: &mut Self::Type,
36 _: &Metrics,
37 ) -> impl Future<Output = RecycleResult<Self::Error>> + Send {
38 future::ready(RecycleResult::Ok(()))
39 }
40}
41
42pub struct PooledConnection<D: Driver> {
43 pub(crate) object: Object<DBConnectionManager<D>>,
44}
45
46pub trait ConnectionPool<D: Driver>: Debug {
47 fn get<'s>(&'s self) -> BoxFuture<'s, Result<PooledConnection<D>>>;
48 fn timeout_get<'s>(&'s self, timeout: Duration) -> BoxFuture<'s, Result<PooledConnection<D>>>;
49 fn detach<'s>(&'s self) -> BoxFuture<'s, Result<D::Connection>>;
50 fn resize(&self, max_size: usize) -> Result<()>;
51 fn into_box(self) -> Box<dyn ConnectionPool<D>>
52 where
53 D: 'static;
54 fn close(self) -> BoxFuture<'static, Result<()>>;
55}
56
57impl<D: Driver> ConnectionPool<D> for Pool<DBConnectionManager<D>>
58where
59 <D as Driver>::Connection: Debug,
60{
61 fn get<'s>(&'s self) -> BoxFuture<'s, Result<PooledConnection<D>>> {
62 async move {
63 let object = Pool::<DBConnectionManager<D>>::get(self)
64 .await
65 .map_err(|e| Error::msg(format!("{e:#?}")))?;
66 Ok(PooledConnection { object })
67 }
68 .boxed()
69 }
70
71 fn timeout_get<'s>(&'s self, timeout: Duration) -> BoxFuture<'s, Result<PooledConnection<D>>> {
72 async move {
73 let object = Pool::<DBConnectionManager<D>>::timeout_get(
74 self,
75 &Timeouts::wait_millis(timeout.as_millis() as u64),
76 )
77 .await
78 .map_err(|e| Error::msg(format!("{e:#?}")))?;
79 Ok(PooledConnection::<D> { object })
80 }
81 .boxed()
82 }
83
84 fn detach<'s>(&'s self) -> BoxFuture<'s, Result<D::Connection>>
85 where
86 Self: Sized,
87 {
88 async {
89 let v = Pool::<DBConnectionManager<D>>::get(self)
90 .await
91 .map_err(|e| Error::msg(format!("{e:#?}")))?;
92 Ok(Object::<DBConnectionManager<D>>::take(v))
93 }
94 .boxed()
95 }
96
97 fn resize(&self, max_size: usize) -> Result<()> {
98 Ok(self.resize(max_size))
99 }
100
101 fn into_box(self) -> Box<dyn ConnectionPool<D>>
102 where
103 D: 'static,
104 {
105 Box::new(self)
106 }
107
108 fn close(self) -> BoxFuture<'static, Result<()>> {
109 Self::close(&self);
110 future::ready(Ok(())).boxed()
111 }
112}
113
114impl<D: Driver> Executor for PooledConnection<D> {
115 type Driver = D;
116
117 fn accepts_multiple_statements(&self) -> bool {
118 self.object.accepts_multiple_statements()
119 }
120
121 fn driver(&self) -> D {
122 self.object.driver()
123 }
124
125 fn prepare<'s>(
126 &'s mut self,
127 query: impl AsQuery<D> + 's,
128 ) -> impl Future<Output = Result<Query<D>>> + Send {
129 self.object.prepare(query)
130 }
131
132 fn do_prepare(&mut self, sql: String) -> impl Future<Output = Result<Query<D>>> + Send {
133 self.object.do_prepare(sql)
134 }
135
136 fn run<'s>(
137 &'s mut self,
138 query: impl AsQuery<D> + 's,
139 ) -> impl Stream<Item = Result<QueryResult>> + Send {
140 self.object.run(query)
141 }
142
143 fn fetch<'s>(
144 &'s mut self,
145 query: impl AsQuery<D> + 's,
146 ) -> impl Stream<Item = Result<Row>> + Send {
147 self.object.fetch(query)
148 }
149
150 fn execute<'s>(
151 &'s mut self,
152 query: impl AsQuery<D> + 's,
153 ) -> impl Future<Output = Result<RowsAffected>> + Send {
154 self.object.execute(query)
155 }
156
157 fn append<'a, E, It>(
158 &mut self,
159 entities: It,
160 ) -> impl Future<Output = Result<RowsAffected>> + Send
161 where
162 E: Entity + 'a,
163 It: IntoIterator<Item = &'a E> + Send,
164 <It as IntoIterator>::IntoIter: Send,
165 {
166 self.object.append(entities)
167 }
168}
169
170impl<D: Driver> Connection for PooledConnection<D> {
171 fn connect(
172 _driver: &D,
173 _url: Cow<'static, str>,
174 ) -> impl Future<Output = Result<<D as Driver>::Connection>> + Send
175 where
176 Self: Sized,
177 {
178 future::ready(Err(Error::msg(
179 "Cannot connect using a PooledConnection, such object must be obtained from a connection pool",
180 )))
181 }
182
183 fn begin(&mut self) -> impl Future<Output = Result<<D as Driver>::Transaction<'_>>> + Send {
184 self.object.begin()
185 }
186}
187
188impl<D: Driver> Deref for PooledConnection<D> {
189 type Target = D::Connection;
190 fn deref(&self) -> &Self::Target {
191 &self.object
192 }
193}
194
195impl<D: Driver> DerefMut for PooledConnection<D> {
196 fn deref_mut(&mut self) -> &mut Self::Target {
197 &mut self.object
198 }
199}
200
201impl<D: Driver> AsRef<D::Connection> for PooledConnection<D> {
202 fn as_ref(&self) -> &D::Connection {
203 &self.object
204 }
205}
206
207impl<D: Driver> AsMut<D::Connection> for PooledConnection<D> {
208 fn as_mut(&mut self) -> &mut D::Connection {
209 &mut self.object
210 }
211}