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

tank_postgres/
connection.rs
