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

tank_postgres/
connection.rs
