Skip to main content

tank_core/
pool.rs

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}