pgsync_postgres/
client.rs

1//! Ordinary (non-replication) connections through tokio-postgres.
2
3use crate::errors::from_pg;
4use crate::tls;
5use crate::url::{ConnectionInfo, TlsPolicy};
6use pgsync_core::config::SafetyConfig;
7use pgsync_core::{Error, Result};
8use std::time::Duration;
9use tokio::task::JoinHandle;
10use tokio_postgres::Client;
11use tracing::debug;
12
13/// A client plus the task driving its socket. Dropping it closes the connection.
14pub struct PgClient {
15    pub client: Client,
16    task: JoinHandle<()>,
17    pub info: ConnectionInfo,
18}
19
20impl Drop for PgClient {
21    fn drop(&mut self) {
22        self.task.abort();
23    }
24}
25
26impl std::ops::Deref for PgClient {
27    type Target = Client;
28    fn deref(&self) -> &Client {
29        &self.client
30    }
31}
32
33pub async fn connect(
34    info: &ConnectionInfo,
35    safety: &SafetyConfig,
36    timeout: Duration,
37    application_name: &str,
38) -> Result<PgClient> {
39    let mut config = info.config.clone();
40    config.application_name(application_name);
41    config.connect_timeout(timeout);
42    let label = info.redacted();
43
44    let fut = async {
45        match info.tls {
46            TlsPolicy::Disabled => {
47                config.ssl_mode(tokio_postgres::config::SslMode::Disable);
48                let (client, conn) = config
49                    .connect(tokio_postgres::NoTls)
50                    .await
51                    .map_err(|e| from_pg(&label, e))?;
52                Ok::<_, Error>((client, tokio::spawn(drive(conn))))
53            }
54            TlsPolicy::Preferred { .. } | TlsPolicy::Required { .. } => {
55                config.ssl_mode(match info.tls {
56                    TlsPolicy::Required { .. } => tokio_postgres::config::SslMode::Require,
57                    _ => tokio_postgres::config::SslMode::Prefer,
58                });
59                let tls_config = tls::client_config(info.tls, safety, info.tls_ca_file.as_deref())?;
60                let connector =
61                    tokio_postgres_rustls::MakeRustlsConnect::new((*tls_config).clone());
62                let (client, conn) = config
63                    .connect(connector)
64                    .await
65                    .map_err(|e| from_pg(&label, e))?;
66                Ok((client, tokio::spawn(drive(conn))))
67            }
68        }
69    };
70    let (client, task) = tokio::time::timeout(timeout + Duration::from_secs(1), fut)
71        .await
72        .map_err(|_| {
73            Error::transient(format!("{label}: connection timed out after {timeout:?}"))
74        })??;
75    debug!(target = %label, "connected");
76    Ok(PgClient {
77        client,
78        task,
79        info: info.clone(),
80    })
81}
82
83async fn drive<S, T>(conn: tokio_postgres::Connection<S, T>)
84where
85    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
86    T: tokio_postgres::tls::TlsStream + Unpin + Send + 'static,
87{
88    if let Err(e) = conn.await {
89        debug!(error = %e, "connection task ended");
90    }
91}
pgsync_postgres/client.rs

pgsync_postgres/
client.rs
