tank_postgres/
connection.rs

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

tank_postgres/
connection.rs
