tank_postgres/
connection.rs

1use crate::{
2    PostgresDriver, PostgresPrepared, PostgresTransaction, ValueWrap,
3    util::{
4        postgres_type_to_value, stream_postgres_row_to_tank_row,
5        stream_postgres_simple_query_message_to_tank_query_result, value_to_postgres_type,
6    },
7};
8use async_stream::try_stream;
9use openssl::ssl::{SslConnector, SslFiletype, SslMethod, SslVerifyMode};
10use postgres_openssl::MakeTlsConnector;
11use postgres_types::ToSql;
12use std::{
13    borrow::Cow,
14    env, mem,
15    path::PathBuf,
16    pin::{Pin, pin},
17    str::FromStr,
18};
19use tank_core::{
20    AsQuery, Connection, Driver, DynQuery, Entity, Error, ErrorContext, Executor, Query,
21    QueryResult, RawQuery, Result, RowsAffected, SqlWriter, Transaction,
22    future::Either,
23    stream::{Stream, StreamExt, TryStreamExt},
24    truncate_long,
25};
26use tokio::{spawn, task::JoinHandle};
27use tokio_postgres::{NoTls, binary_copy::BinaryCopyInWriter};
28use url::Url;
29
30/// PostgreSQL connection.
31#[derive(Debug)]
32pub struct PostgresConnection {
33    pub(crate) client: tokio_postgres::Client,
34    pub(crate) handle: JoinHandle<()>,
35    pub(crate) url: Url,
36}
37
38impl Executor for PostgresConnection {
39    type Driver = PostgresDriver;
40
41    async fn do_prepare(&mut self, sql: String) -> Result<Query<PostgresDriver>> {
42        let sql = sql.as_str().trim_end().trim_end_matches(';');
43        Ok(
44            PostgresPrepared::new(self.client.prepare(&sql).await.map_err(|e| {
45                let error = Error::new(e).context(format!(
46                    "While preparing the query:\n{}",
47                    truncate_long!(sql)
48                ));
49                log::error!("{:#}", error);
50                error
51            })?)
52            .into(),
53        )
54    }
55
56    fn run<'s>(
57        &'s mut self,
58        query: impl AsQuery<PostgresDriver> + 's,
59    ) -> impl Stream<Item = Result<QueryResult>> + Send {
60        let mut query = query.as_query();
61        let context = format!("While running the query:\n{}", query.as_mut());
62        let mut owned = mem::take(query.as_mut());
63        match owned {
64            Query::Raw(raw) => Either::Left(try_stream! {
65                let sql = &raw.0;
66                {
67                    let stream = stream_postgres_simple_query_message_to_tank_query_result(
68                        async move || self.client.simple_query_raw(sql).await.map_err(Into::into),
69                    );
70                    let mut stream = pin!(stream);
71                    while let Some(value) = stream.next().await.transpose()? {
72                        yield value;
73                    }
74                }
75                *query.as_mut() = Query::Raw(raw);
76            }),
77            Query::Prepared(..) => Either::Right(try_stream! {
78                let mut transaction = self.begin().await?;
79                {
80                    let mut stream = pin!(transaction.run(&mut owned));
81                    while let Some(value) = stream.next().await.transpose()? {
82                        yield value;
83                    }
84                }
85                transaction.commit().await?;
86                *query.as_mut() = mem::take(&mut owned);
87            }),
88        }
89        .map_err(move |e: Error| {
90            let error = e.context(context.clone());
91            log::error!("{:#}", error);
92            error
93        })
94    }
95
96    fn fetch<'s>(
97        &'s mut self,
98        query: impl AsQuery<PostgresDriver> + 's,
99    ) -> impl Stream<Item = Result<tank_core::Row>> + Send {
100        let mut query = query.as_query();
101        let context = format!("While fetching the query:\n{}", query.as_mut());
102        let owned = mem::take(query.as_mut());
103        stream_postgres_row_to_tank_row(async move || {
104            let row_stream = match owned {
105                Query::Raw(RawQuery(sql)) => {
106                    let stream = self
107                        .client
108                        .query_raw(&sql, Vec::<ValueWrap>::new())
109                        .await
110                        .map_err(|e| Error::new(e).context(context.clone()))?;
111                    *query.as_mut() = Query::raw(sql);
112                    stream
113                }
114                Query::Prepared(mut prepared) => {
115                    let mut params = prepared.take_params();
116                    let types = prepared.statement.params();
117
118                    for (i, param) in params.iter_mut().enumerate() {
119                        *param = ValueWrap(Cow::Owned(
120                            mem::take(param)
121                                .0
122                                .into_owned()
123                                .try_as(&postgres_type_to_value(&types[i]))?,
124                        ));
125                    }
126                    let stream = self
127                        .client
128                        .query_raw(&prepared.statement, params)
129                        .await
130                        .map_err(|e| Error::new(e).context(context.clone()))?;
131                    *query.as_mut() = Query::Prepared(prepared);
132                    stream
133                }
134            };
135            Ok(row_stream).map_err(|e| {
136                log::error!("{:#}", e);
137                e
138            })
139        })
140    }
141
142    async fn append<'a, E, It>(&mut self, entities: It) -> Result<RowsAffected>
143    where
144        E: Entity + 'a,
145        It: IntoIterator<Item = &'a E> + Send,
146        <It as IntoIterator>::IntoIter: Send,
147    {
148        let writer = self.driver().sql_writer();
149        let context = || {
150            format!(
151                "While appending to the table `{}`",
152                E::table().full_name(writer.separator())
153            )
154        };
155        let mut query = DynQuery::default();
156        writer.write_copy::<E>(&mut query);
157        let sink = match self
158            .client
159            .copy_in(&query.as_str() as &str)
160            .await
161            .with_context(context)
162        {
163            Ok(v) => v,
164            Err(e) => {
165                log::error!("{e:#}");
166                return Err(e);
167            }
168        };
169        let types: Vec<_> = E::columns()
170            .into_iter()
171            .map(|c| value_to_postgres_type(&c.value))
172            .collect();
173        let writer = BinaryCopyInWriter::new(sink, &types);
174        let mut writer = pin!(writer);
175        let columns_len = E::columns().len();
176        let mut values = Vec::<ValueWrap>::with_capacity(columns_len);
177        let mut refs = Vec::<&(dyn ToSql + Sync)>::with_capacity(columns_len);
178        for entity in entities.into_iter() {
179            values.extend(
180                entity
181                    .row_values()
182                    .into_iter()
183                    .map(|v| ValueWrap(Cow::Owned(v))),
184            );
185            refs.extend(
186                values
187                    .iter()
188                    .map(|v| unsafe { &*(v as &(dyn ToSql + Sync) as *const _) }),
189            );
190            match Pin::as_mut(&mut writer)
191                .write(&refs)
192                .await
193                .with_context(context)
194            {
195                Ok(_) => {}
196                Err(e) => {
197                    log::error!("{e:#}");
198                    return Err(e);
199                }
200            };
201            refs.clear();
202            values.clear();
203        }
204        match writer.finish().await.with_context(context) {
205            Ok(v) => Ok(RowsAffected {
206                rows_affected: Some(v),
207                last_affected_id: None,
208            }),
209            Err(e) => {
210                log::error!("{e:#}");
211                return Err(e);
212            }
213        }
214    }
215}
216
217impl Connection for PostgresConnection {
218    async fn connect(driver: &PostgresDriver, url: Cow<'static, str>) -> Result<Self> {
219        let context = "While trying to connect to Postgres";
220        let mut url = Self::sanitize_url(driver, url)?;
221        let mut take_url_param = |key: &str, env_var: &str, remove: bool| {
222            let value = url
223                .query_pairs()
224                .find_map(|(k, v)| if k == key { Some(v) } else { None })
225                .map(|v| v.to_string());
226            if remove && let Some(..) = value {
227                let mut result = url.clone();
228                result.set_query(None);
229                result
230                    .query_pairs_mut()
231                    .extend_pairs(url.query_pairs().filter(|(k, _)| k != key));
232                url = result;
233            };
234            value.or_else(|| env::var(env_var).ok().map(Into::into))
235        };
236        let sslmode = take_url_param("sslmode", "PGSSLMODE", false).unwrap_or("disable".into());
237        let (client, handle) = if sslmode == "disable" {
238            let (client, connection) = tokio_postgres::connect(url.as_str(), NoTls).await?;
239            let handle = spawn(async move {
240                if let Err(error) = connection.await
241                    && !error.is_closed()
242                {
243                    log::error!("Postgres connection error: {:#?}", error);
244                }
245            });
246            (client, handle)
247        } else {
248            let mut builder = SslConnector::builder(SslMethod::tls())?;
249            let path = PathBuf::from_str(
250                take_url_param("sslrootcert", "PGSSLROOTCERT", true)
251                    .as_deref()
252                    .unwrap_or("~/.postgresql/root.crt"),
253            )
254            .context(context)?;
255            if path.exists() {
256                builder.set_ca_file(path)?;
257            }
258            let path = PathBuf::from_str(
259                take_url_param("sslcert", "PGSSLCERT", true)
260                    .as_deref()
261                    .unwrap_or("~/.postgresql/postgresql.crt"),
262            )
263            .context(context)?;
264            if path.exists() {
265                builder.set_certificate_chain_file(path)?;
266            }
267            let path = PathBuf::from_str(
268                take_url_param("sslkey", "PGSSLKEY", true)
269                    .as_deref()
270                    .unwrap_or("~/.postgresql/postgresql.key"),
271            )
272            .context(context)?;
273            if path.exists() {
274                builder.set_private_key_file(path, SslFiletype::PEM)?;
275            }
276            builder.set_verify(SslVerifyMode::PEER);
277            let connector = MakeTlsConnector::new(builder.build());
278            let (client, connection) = tokio_postgres::connect(url.as_str(), connector).await?;
279            let handle = spawn(async move {
280                if let Err(error) = connection.await
281                    && !error.is_closed()
282                {
283                    log::error!("Postgres connection error: {:#?}", error);
284                }
285            });
286            (client, handle)
287        };
288        Ok(Self {
289            client,
290            handle,
291            url,
292        })
293    }
294
295    fn begin(&mut self) -> impl Future<Output = Result<PostgresTransaction<'_>>> {
296        PostgresTransaction::new(self)
297    }
298
299    async fn duplicate(&self) -> Result<PostgresConnection>
300    where
301        Self: Sized,
302    {
303        Self::connect(&self.driver(), self.url.to_string().into()).await
304    }
305
306    async fn disconnect(self) -> Result<()> {
307        drop(self.client);
308        if let Err(e) = self.handle.await {
309            let error = Error::new(e).context("While disconnecting from Postgres");
310            log::error!("{:#}", error);
311            return Err(error);
312        }
313        Ok(())
314    }
315}
tank_postgres/connection.rs

tank_postgres/
connection.rs
