qail_pg/driver/
builder.rs1use 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#[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 pub fn new() -> Self {
41 Self::default()
42 }
43
44 pub fn host(mut self, host: impl Into<String>) -> Self {
46 self.host = Some(host.into());
47 self
48 }
49
50 pub fn port(mut self, port: u16) -> Self {
52 self.port = Some(port);
53 self
54 }
55
56 pub fn user(mut self, user: impl Into<String>) -> Self {
58 self.user = Some(user.into());
59 self
60 }
61
62 pub fn database(mut self, database: impl Into<String>) -> Self {
64 self.database = Some(database.into());
65 self
66 }
67
68 pub fn password(mut self, password: impl Into<String>) -> Self {
70 self.password = Some(password.into());
71 self
72 }
73
74 pub fn timeout(mut self, timeout: std::time::Duration) -> Self {
76 self.timeout = Some(timeout);
77 self
78 }
79
80 pub fn tls_mode(mut self, mode: TlsMode) -> Self {
82 self.connect_options.tls_mode = mode;
83 self
84 }
85
86 pub fn gss_enc_mode(mut self, mode: GssEncMode) -> Self {
88 self.connect_options.gss_enc_mode = mode;
89 self
90 }
91
92 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 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 pub fn auth_settings(mut self, settings: AuthSettings) -> Self {
107 self.connect_options.auth = settings;
108 self
109 }
110
111 pub fn channel_binding_mode(mut self, mode: ScramChannelBindingMode) -> Self {
113 self.connect_options.auth.channel_binding = mode;
114 self
115 }
116
117 pub fn io_uring(mut self, enabled: bool) -> Self {
119 self.connect_options.io_uring = enabled;
120 self
121 }
122
123 pub fn gss_token_provider(mut self, provider: GssTokenProvider) -> Self {
125 self.connect_options.gss_token_provider = Some(provider);
126 self
127 }
128
129 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 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 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 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}