spin_factor_outbound_pg/
client.rs

1#![allow(clippy::result_large_err)]
2
3use std::sync::Arc;
4
5use anyhow::{Context, Result};
6use futures::stream::TryStreamExt as _;
7use native_tls::TlsConnector;
8use postgres_native_tls::MakeTlsConnector;
9use spin_world::async_trait;
10use spin_world::spin::postgres4_2_0::postgres::{
11    self as v4, Column, DbValue, ParameterValue, RowSet,
12};
13use tokio_postgres::config::SslMode;
14use tokio_postgres::types::ToSql;
15use tokio_postgres::{NoTls, Row};
16
17use crate::types::{convert_data_type, convert_entry, to_sql_parameter, to_sql_parameters};
18
19/// Max connections in a given address' connection pool
20const CONNECTION_POOL_SIZE: usize = 64;
21/// Max addresses for which to keep pools in cache.
22const CONNECTION_POOL_CACHE_CAPACITY: u64 = 16;
23
24/// A factory object for Postgres clients. This abstracts
25/// details of client creation such as pooling.
26#[async_trait]
27pub trait ClientFactory: Default + Send + Sync + 'static {
28    /// The type of client produced by `get_client`.
29    type Client: Client;
30    /// Gets a client from the factory.
31    async fn get_client(
32        &self,
33        address: &str,
34        root_ca: Option<HashableCertificate>,
35    ) -> Result<Self::Client>;
36}
37
38#[derive(Clone)]
39pub struct HashableCertificate {
40    certificate: native_tls::Certificate,
41    hash: String,
42}
43
44impl HashableCertificate {
45    pub fn from_pem(text: &str) -> anyhow::Result<Self> {
46        let cert_bytes = text.as_bytes();
47        let hash = spin_common::sha256::hex_digest_from_bytes(cert_bytes);
48        let certificate =
49            native_tls::Certificate::from_pem(cert_bytes).context("invalid root certificate")?;
50        Ok(Self { certificate, hash })
51    }
52}
53
54/// A `ClientFactory` that uses a connection pool per address.
55pub struct PooledTokioClientFactory {
56    pools: moka::sync::Cache<(String, Option<String>), deadpool_postgres::Pool>,
57}
58
59impl Default for PooledTokioClientFactory {
60    fn default() -> Self {
61        Self {
62            pools: moka::sync::Cache::new(CONNECTION_POOL_CACHE_CAPACITY),
63        }
64    }
65}
66
67#[derive(Clone)]
68pub struct PooledTokioClient(Arc<deadpool_postgres::Object>);
69
70impl AsRef<deadpool_postgres::Object> for PooledTokioClient {
71    fn as_ref(&self) -> &deadpool_postgres::Object {
72        self.0.as_ref()
73    }
74}
75
76#[async_trait]
77impl ClientFactory for PooledTokioClientFactory {
78    type Client = PooledTokioClient;
79
80    async fn get_client(
81        &self,
82        address: &str,
83        root_ca: Option<HashableCertificate>,
84    ) -> Result<Self::Client> {
85        let (root_ca, root_ca_hash) = match root_ca {
86            None => (None, None),
87            Some(HashableCertificate { certificate, hash }) => (Some(certificate), Some(hash)),
88        };
89        let pool_key = (address.to_string(), root_ca_hash);
90        let pool = self
91            .pools
92            .try_get_with_by_ref(&pool_key, || create_connection_pool(address, root_ca))
93            .map_err(ArcError)
94            .context("establishing PostgreSQL connection pool")?;
95
96        Ok(PooledTokioClient(Arc::new(pool.get().await?)))
97    }
98}
99
100/// Creates a Postgres connection pool for the given address.
101fn create_connection_pool(
102    address: &str,
103    root_ca: Option<native_tls::Certificate>,
104) -> Result<deadpool_postgres::Pool> {
105    let config = address
106        .parse::<tokio_postgres::Config>()
107        .context("parsing Postgres connection string")?;
108
109    // The debug form isn't as easy to read as logging the address,
110    // but it redacts the password!
111    tracing::debug!("Build new connection: {config:?}");
112
113    let mgr_config = deadpool_postgres::ManagerConfig {
114        recycling_method: deadpool_postgres::RecyclingMethod::Clean,
115    };
116
117    let mgr = if config.get_ssl_mode() == SslMode::Disable {
118        deadpool_postgres::Manager::from_config(config, NoTls, mgr_config)
119    } else {
120        let mut builder = TlsConnector::builder();
121        if let Some(root_ca) = root_ca {
122            builder.add_root_certificate(root_ca);
123        }
124        let connector = MakeTlsConnector::new(builder.build()?);
125        deadpool_postgres::Manager::from_config(config, connector, mgr_config)
126    };
127
128    // TODO: what is our max size heuristic?  Should this be passed in so that different
129    // hosts can manage it according to their needs?  Will a plain number suffice for
130    // sophisticated hosts anyway?
131    let pool = deadpool_postgres::Pool::builder(mgr)
132        .max_size(CONNECTION_POOL_SIZE)
133        .build()
134        .context("building Postgres connection pool")?;
135
136    Ok(pool)
137}
138
139#[async_trait]
140pub trait Client: Clone + Send + Sync + 'static {
141    async fn execute(
142        &self,
143        statement: String,
144        params: Vec<ParameterValue>,
145    ) -> Result<u64, v4::Error>;
146
147    async fn query(
148        &self,
149        statement: String,
150        params: Vec<ParameterValue>,
151        max_result_bytes: usize,
152    ) -> Result<RowSet, v4::Error>;
153
154    async fn query_async(
155        &self,
156        statement: String,
157        params: Vec<ParameterValue>,
158        max_result_bytes: usize,
159    ) -> Result<QueryAsyncResult, v4::Error>;
160}
161
162pub struct QueryAsyncResult {
163    pub columns: Vec<v4::Column>,
164    pub rows: tokio::sync::mpsc::Receiver<v4::Row>,
165    pub error: tokio::sync::oneshot::Receiver<Result<(), v4::Error>>,
166}
167
168/// Extract weak-typed error data for WIT purposes
169fn pg_extras(dbe: &tokio_postgres::error::DbError) -> Vec<(String, String)> {
170    let mut extras = vec![];
171
172    macro_rules! pg_extra {
173        ( $n:ident ) => {
174            if let Some(value) = dbe.$n() {
175                extras.push((stringify!($n).to_owned(), value.to_string()));
176            }
177        };
178    }
179
180    pg_extra!(column);
181    pg_extra!(constraint);
182    pg_extra!(routine);
183    pg_extra!(hint);
184    pg_extra!(table);
185    pg_extra!(datatype);
186    pg_extra!(schema);
187    pg_extra!(file);
188    pg_extra!(line);
189    pg_extra!(where_);
190
191    extras
192}
193
194fn query_failed(e: tokio_postgres::error::Error) -> v4::Error {
195    let flattened = format!("{e:?}");
196    let query_error = match e.as_db_error() {
197        None => v4::QueryError::Text(flattened),
198        Some(dbe) => v4::QueryError::DbError(v4::DbError {
199            as_text: flattened,
200            severity: dbe.severity().to_owned(),
201            code: dbe.code().code().to_owned(),
202            message: dbe.message().to_owned(),
203            detail: dbe.detail().map(|s| s.to_owned()),
204            extras: pg_extras(dbe),
205        }),
206    };
207    v4::Error::QueryFailed(query_error)
208}
209
210fn query_failed_anyhow(e: anyhow::Error) -> v4::Error {
211    let text = format!("{e:?}");
212    v4::Error::QueryFailed(v4::QueryError::Text(text))
213}
214
215#[async_trait]
216impl Client for PooledTokioClient {
217    async fn execute(
218        &self,
219        statement: String,
220        params: Vec<ParameterValue>,
221    ) -> Result<u64, v4::Error> {
222        let params = params
223            .iter()
224            .map(to_sql_parameter)
225            .collect::<Result<Vec<_>>>()
226            .map_err(|e| v4::Error::ValueConversionFailed(format!("{e:?}")))?;
227
228        let params_refs: Vec<&(dyn ToSql + Sync)> = params
229            .iter()
230            .map(|b| b.as_ref() as &(dyn ToSql + Sync))
231            .collect();
232
233        self.as_ref()
234            .execute(&statement, params_refs.as_slice())
235            .await
236            .map_err(query_failed)
237    }
238
239    async fn query(
240        &self,
241        statement: String,
242        params: Vec<ParameterValue>,
243        max_result_bytes: usize,
244    ) -> Result<RowSet, v4::Error> {
245        let (columns, mut results) = self.query_stream(statement, params).await?;
246
247        let mut byte_count = std::mem::size_of::<RowSet>();
248        let mut rows = Vec::new();
249
250        async {
251            while let Some(row) = results.try_next().await? {
252                byte_count += row.iter().map(|v| v.memory_size()).sum::<usize>();
253                if byte_count > max_result_bytes {
254                    anyhow::bail!("query result exceeds limit of {max_result_bytes} bytes")
255                }
256                rows.push(row);
257            }
258            Ok(())
259        }
260        .await
261        .map_err(|e| v4::Error::QueryFailed(v4::QueryError::Text(format!("{e:?}"))))?;
262
263        Ok(RowSet { columns, rows })
264    }
265
266    async fn query_async(
267        &self,
268        statement: String,
269        params: Vec<ParameterValue>,
270        max_result_bytes: usize,
271    ) -> Result<QueryAsyncResult, v4::Error> {
272        use futures::StreamExt;
273
274        let (columns, mut rows) = self.query_stream(statement, params).await?;
275
276        let (rows_tx, rows_rx) = tokio::sync::mpsc::channel(4);
277        let (err_tx, err_rx) = tokio::sync::oneshot::channel();
278
279        tokio::spawn(async move {
280            loop {
281                tokio::select! {
282                    biased;
283                    // Exit immediately if the consumer (rows_rx) has been dropped,
284                    // preventing this task from leaking.
285                    _ = rows_tx.closed() => {
286                        _ = err_tx.send(Ok(()));
287                        return;
288                    }
289                    row = rows.next() => {
290                        let Some(row) = row else {
291                            _ = err_tx.send(Ok(()));
292                            return;
293                        };
294                        match row {
295                            Ok(row) => {
296                                let byte_count = row.iter().map(|v| v.memory_size()).sum::<usize>();
297                                if byte_count > max_result_bytes {
298                                    _ = err_tx.send(Err(v4::Error::QueryFailed(v4::QueryError::Text(
299                                        format!("query result exceeds limit of {max_result_bytes} bytes"),
300                                    ))));
301                                    return;
302                                }
303
304                                if let Err(e) = rows_tx.send(row).await {
305                                    _ = err_tx.send(Err(v4::Error::QueryFailed(v4::QueryError::Text(
306                                        format!("async error: {e}"),
307                                    ))));
308                                    return;
309                                }
310                            }
311                            Err(e) => {
312                                _ = err_tx.send(Err(e));
313                                return;
314                            }
315                        }
316                    }
317                }
318            }
319        });
320
321        Ok(QueryAsyncResult {
322            columns,
323            rows: rows_rx,
324            error: err_rx,
325        })
326    }
327}
328
329impl PooledTokioClient {
330    async fn query_stream(
331        &self,
332        statement: String,
333        params: Vec<ParameterValue>,
334    ) -> Result<
335        (
336            Vec<v4::Column>,
337            impl futures::Stream<Item = Result<Vec<DbValue>, v4::Error>> + 'static,
338        ),
339        v4::Error,
340    > {
341        use futures::StreamExt;
342
343        let params = to_sql_parameters(params)?;
344        let stmt = self
345            .as_ref()
346            .prepare(&statement)
347            .await
348            .map_err(query_failed)?;
349        let columns = stmt.columns().iter().map(infer_column).collect();
350
351        let results = Box::pin(
352            self.as_ref()
353                .query_raw(&stmt, params)
354                .await
355                .map_err(query_failed)?,
356        );
357
358        let row_stm = results.map(move |row_res| {
359            let row_res = row_res.map_err(query_failed);
360            row_res.and_then(|r| convert_row(&r).map_err(query_failed_anyhow))
361        });
362
363        Ok((columns, Box::pin(row_stm)))
364    }
365}
366
367fn infer_column(column: &tokio_postgres::Column) -> Column {
368    let name = column.name().to_owned();
369    let data_type = convert_data_type(column.type_());
370    Column { name, data_type }
371}
372
373fn convert_row(row: &Row) -> anyhow::Result<Vec<DbValue>> {
374    let mut result = Vec::with_capacity(row.len());
375    for index in 0..row.len() {
376        result.push(convert_entry(row, index)?);
377    }
378    Ok(result)
379}
380
381/// Workaround for moka returning Arc<Error> which, although
382/// necessary for concurrency, does not play well with others.
383struct ArcError(std::sync::Arc<anyhow::Error>);
384
385impl std::error::Error for ArcError {
386    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
387        self.0.source()
388    }
389}
390
391impl std::fmt::Debug for ArcError {
392    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
393        std::fmt::Debug::fmt(&self.0, f)
394    }
395}
396
397impl std::fmt::Display for ArcError {
398    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
399        std::fmt::Display::fmt(&self.0, f)
400    }
401}
spin_factor_outbound_pg/client.rs

spin_factor_outbound_pg/
client.rs
