Skip to main content

tank_postgres/
transaction.rs

1use crate::{
2    PostgresConnection, PostgresDriver, PostgresPrepared, ValueWrap,
3    util::stream_postgres_row_to_tank_row,
4};
5use tank_core::{
6    AsQuery, Error, Executor, Query, QueryResult, RawQuery, Result, Transaction,
7    future::{Either, TryFutureExt},
8    stream::{Stream, TryStreamExt},
9};
10
11/// Postgres transaction.
12pub struct PostgresTransaction<'c>(pub(crate) tokio_postgres::Transaction<'c>);
13
14impl<'c> PostgresTransaction<'c> {
15    pub async fn new(client: &'c mut PostgresConnection) -> Result<Self> {
16        Ok(Self(client.client.transaction().await.map_err(|e| {
17            log::error!("{e:#}");
18            e
19        })?))
20    }
21}
22
23impl<'c> Executor for PostgresTransaction<'c> {
24    type Driver = PostgresDriver;
25
26    async fn do_prepare(
27        &mut self,
28        RawQuery { sql, .. }: RawQuery,
29    ) -> Result<Query<PostgresDriver>> {
30        Ok(
31            PostgresPrepared::new(self.0.prepare(&sql).await.map_err(|e| {
32                let error = Error::new(e);
33                log::error!("{error:#}");
34                error
35            })?)
36            .into(),
37        )
38    }
39    fn run<'s>(
40        &'s mut self,
41        query: impl AsQuery<PostgresDriver> + 's,
42    ) -> impl Stream<Item = Result<QueryResult>> + Send {
43        let mut query = query.as_query();
44        stream_postgres_row_to_tank_row(async move || match query.as_mut() {
45            Query::Raw(RawQuery { sql, .. }) => {
46                let stream = self
47                    .0
48                    .query_raw(sql.as_str(), Vec::<ValueWrap>::new())
49                    .await?;
50                Ok(Either::Left(stream))
51            }
52            Query::Prepared(prepared) => {
53                let params = prepared.take_params();
54                let portal = self
55                    .0
56                    .bind_raw(&prepared.statement, params.into_iter())
57                    .await?;
58                Ok(Either::Right(self.0.query_portal_raw(&portal, 0).await?))
59            }
60        })
61        .map_err(|e| {
62            log::error!("{e:#}");
63            e
64        })
65    }
66}
67
68impl<'c> Transaction<'c> for PostgresTransaction<'c> {
69    fn commit(self) -> impl Future<Output = Result<()>> + Send {
70        self.0.commit().map_err(|e| {
71            let error = Error::new(e);
72            log::error!("{error:#}");
73            error
74        })
75    }
76
77    fn rollback(self) -> impl Future<Output = Result<()>> + Send {
78        self.0.rollback().map_err(|e| {
79            let error = Error::new(e);
80            log::error!("{error:#}");
81            error
82        })
83    }
84}