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