|
1 | 1 | use anyhow::anyhow;
|
| 2 | +use native_tls::TlsConnector; |
| 3 | +use postgres_native_tls::MakeTlsConnector; |
2 | 4 | use spin_core::HostComponent;
|
3 | 5 | use std::collections::HashMap;
|
4 | 6 | use tokio_postgres::{
|
| 7 | + config::SslMode, |
5 | 8 | types::{ToSql, Type},
|
6 |
| - Client, NoTls, Row, |
| 9 | + Client, NoTls, Row, Socket, |
7 | 10 | };
|
8 | 11 | use wit_bindgen_wasmtime::async_trait;
|
9 | 12 |
|
@@ -237,15 +240,42 @@ impl OutboundPg {
|
237 | 240 | }
|
238 | 241 |
|
239 | 242 | async fn build_client(address: &str) -> anyhow::Result<Client> {
|
| 243 | + let config = address.parse::<tokio_postgres::Config>()?; |
| 244 | + |
240 | 245 | tracing::log::debug!("Build new connection: {}", address);
|
241 | 246 |
|
242 |
| - let (client, connection) = tokio_postgres::connect(address, NoTls).await?; |
| 247 | + if config.get_ssl_mode() == SslMode::Disable { |
| 248 | + connect(config).await |
| 249 | + } else { |
| 250 | + connect_tls(config).await |
| 251 | + } |
| 252 | +} |
| 253 | + |
| 254 | +async fn connect(config: tokio_postgres::Config) -> anyhow::Result<Client> { |
| 255 | + let (client, connection) = config.connect(NoTls).await?; |
| 256 | + |
| 257 | + spawn(connection); |
| 258 | + |
| 259 | + Ok(client) |
| 260 | +} |
243 | 261 |
|
| 262 | +async fn connect_tls(config: tokio_postgres::Config) -> anyhow::Result<Client> { |
| 263 | + let builder = TlsConnector::builder(); |
| 264 | + let connector = MakeTlsConnector::new(builder.build()?); |
| 265 | + let (client, connection) = config.connect(connector).await?; |
| 266 | + |
| 267 | + spawn(connection); |
| 268 | + |
| 269 | + Ok(client) |
| 270 | +} |
| 271 | + |
| 272 | +fn spawn<T>(connection: tokio_postgres::Connection<Socket, T>) |
| 273 | +where |
| 274 | + T: tokio_postgres::tls::TlsStream + std::marker::Unpin + std::marker::Send + 'static, |
| 275 | +{ |
244 | 276 | tokio::spawn(async move {
|
245 | 277 | if let Err(e) = connection.await {
|
246 | 278 | tracing::warn!("Postgres connection error: {}", e);
|
247 | 279 | }
|
248 | 280 | });
|
249 |
| - |
250 |
| - Ok(client) |
251 | 281 | }
|
0 commit comments