Skip to main content

qail_pg/driver/
builder.rs

1//! PgDriverBuilder — ergonomic builder pattern for PgDriver connections.
2
3use super::auth_types::{
4    AuthSettings, ConnectOptions, GssEncMode, GssTokenProvider, ScramChannelBindingMode, TlsMode,
5};
6use super::core::PgDriver;
7use super::types::{PgError, PgResult};
8use crate::driver::connection::TlsConfig;
9
10// ============================================================================
11// Connection Builder
12// ============================================================================
13
14/// Builder for creating PgDriver connections with named parameters.
15/// # Example
16/// ```ignore
17/// let driver = PgDriver::builder()
18///     .host("localhost")
19///     .port(5432)
20///     .user("admin")
21///     .database("mydb")
22///     .password("secret")
23///     .connect()
24///     .await?;
25/// ```
26#[derive(Default)]
27pub struct PgDriverBuilder {
28    host: Option<String>,
29    port: Option<u16>,
30    user: Option<String>,
31    database: Option<String>,
32    password: Option<String>,
33    timeout: Option<std::time::Duration>,
34    pub(crate) connect_options: ConnectOptions,
35}
36
37impl PgDriverBuilder {
38    /// Create a new builder with default values.
39    pub fn new() -> Self {
40        Self::default()
41    }
42
43    /// Set the host (default: "127.0.0.1").
44    pub fn host(mut self, host: impl Into<String>) -> Self {
45        self.host = Some(host.into());
46        self
47    }
48
49    /// Set the port (default: 5432).
50    pub fn port(mut self, port: u16) -> Self {
51        self.port = Some(port);
52        self
53    }
54
55    /// Set the username (required).
56    pub fn user(mut self, user: impl Into<String>) -> Self {
57        self.user = Some(user.into());
58        self
59    }
60
61    /// Set the database name (required).
62    pub fn database(mut self, database: impl Into<String>) -> Self {
63        self.database = Some(database.into());
64        self
65    }
66
67    /// Set the password (optional, for cleartext/MD5/SCRAM-SHA-256 auth).
68    pub fn password(mut self, password: impl Into<String>) -> Self {
69        self.password = Some(password.into());
70        self
71    }
72
73    /// Set connection timeout (optional).
74    pub fn timeout(mut self, timeout: std::time::Duration) -> Self {
75        self.timeout = Some(timeout);
76        self
77    }
78
79    /// Set TLS policy (`disable`, `prefer`, `require`).
80    pub fn tls_mode(mut self, mode: TlsMode) -> Self {
81        self.connect_options.tls_mode = mode;
82        self
83    }
84
85    /// Set GSSAPI session encryption mode (`disable`, `prefer`, `require`).
86    pub fn gss_enc_mode(mut self, mode: GssEncMode) -> Self {
87        self.connect_options.gss_enc_mode = mode;
88        self
89    }
90
91    /// Set custom CA bundle PEM for TLS validation.
92    pub fn tls_ca_cert_pem(mut self, ca_pem: Vec<u8>) -> Self {
93        self.connect_options.tls_ca_cert_pem = Some(ca_pem);
94        self
95    }
96
97    /// Enable mTLS using client certificate/key config.
98    pub fn mtls(mut self, config: TlsConfig) -> Self {
99        self.connect_options.mtls = Some(config);
100        self.connect_options.tls_mode = TlsMode::Require;
101        self
102    }
103
104    /// Override password-auth policy.
105    pub fn auth_settings(mut self, settings: AuthSettings) -> Self {
106        self.connect_options.auth = settings;
107        self
108    }
109
110    /// Set SCRAM channel-binding mode.
111    pub fn channel_binding_mode(mut self, mode: ScramChannelBindingMode) -> Self {
112        self.connect_options.auth.channel_binding = mode;
113        self
114    }
115
116    /// Opt into Linux io_uring for plain TCP transport.
117    pub fn io_uring(mut self, enabled: bool) -> Self {
118        self.connect_options.io_uring = enabled;
119        self
120    }
121
122    /// Set a stateful Kerberos/GSS/SSPI token provider.
123    pub fn gss_token_provider(mut self, provider: GssTokenProvider) -> Self {
124        self.connect_options.gss_token_provider = Some(provider);
125        self
126    }
127
128    /// Add a custom StartupMessage parameter.
129    ///
130    /// Example: `.startup_param("application_name", "qail-replica")`
131    pub fn startup_param(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
132        let key = key.into();
133        let value = value.into();
134        self.connect_options
135            .startup_params
136            .retain(|(existing, _)| !existing.eq_ignore_ascii_case(&key));
137        self.connect_options.startup_params.push((key, value));
138        self
139    }
140
141    /// Enable logical replication startup mode (`replication=database`).
142    ///
143    /// This is required before issuing commands like `IDENTIFY_SYSTEM` or
144    /// `CREATE_REPLICATION_SLOT` on a replication connection.
145    pub fn logical_replication(mut self) -> Self {
146        self.connect_options
147            .startup_params
148            .retain(|(k, _)| !k.eq_ignore_ascii_case("replication"));
149        self.connect_options
150            .startup_params
151            .push(("replication".to_string(), "database".to_string()));
152        self
153    }
154
155    /// Connect to PostgreSQL using the configured parameters.
156    pub async fn connect(self) -> PgResult<PgDriver> {
157        let host = self.host.unwrap_or_else(|| "127.0.0.1".to_string());
158        let port = self.port.unwrap_or(5432);
159        let user = self
160            .user
161            .ok_or_else(|| PgError::Connection("User is required".to_string()))?;
162        let database = self
163            .database
164            .ok_or_else(|| PgError::Connection("Database is required".to_string()))?;
165
166        let password = self.password;
167        let options = self.connect_options;
168
169        if let Some(timeout) = self.timeout {
170            let options = options.clone();
171            tokio::time::timeout(
172                timeout,
173                PgDriver::connect_with_options(
174                    &host,
175                    port,
176                    &user,
177                    &database,
178                    password.as_deref(),
179                    options,
180                ),
181            )
182            .await
183            .map_err(|_| PgError::Timeout(format!("connection after {:?}", timeout)))?
184        } else {
185            PgDriver::connect_with_options(
186                &host,
187                port,
188                &user,
189                &database,
190                password.as_deref(),
191                options,
192            )
193            .await
194        }
195    }
196}