systemprompt_database/services/postgres/
connection.rs1use std::future::Future;
11use std::time::Duration;
12
13use sqlx::postgres::{PgConnectOptions, PgPool, PgPoolOptions};
14
15use crate::error::DatabaseResult;
16use crate::resilience::classify::Outcome;
17use crate::resilience::config::RetryConfig;
18use crate::resilience::retry::retry_async;
19
20const RETRY_DELAYS_MS: &[u64] = &[100, 200, 400, 800, 1600];
21const MAX_ATTEMPTS: u32 = 5;
22
23#[derive(Debug, Clone, Copy)]
30pub struct PoolConfig {
31 pub max_connections: u32,
32 pub min_connections: u32,
33 pub acquire_timeout: Duration,
34 pub idle_timeout: Duration,
35 pub max_lifetime: Duration,
36}
37
38impl Default for PoolConfig {
39 fn default() -> Self {
40 Self {
41 max_connections: 50,
42 min_connections: 0,
43 acquire_timeout: Duration::from_secs(30),
44 idle_timeout: Duration::from_mins(5),
45 max_lifetime: Duration::from_mins(30),
46 }
47 }
48}
49
50#[must_use]
51pub fn build_pool_options(cfg: &PoolConfig) -> PgPoolOptions {
52 PgPoolOptions::new()
53 .max_connections(cfg.max_connections)
54 .min_connections(cfg.min_connections)
55 .max_lifetime(cfg.max_lifetime)
56 .acquire_timeout(cfg.acquire_timeout)
57 .idle_timeout(cfg.idle_timeout)
58}
59
60pub async fn connect_with_retry(
61 options: PgPoolOptions,
62 connect_options: PgConnectOptions,
63) -> DatabaseResult<PgPool> {
64 let connector = |opts: PgConnectOptions| {
65 let options = options.clone();
66 async move { options.connect_with(opts).await }
67 };
68 connect_with_retry_using(connect_options, MAX_ATTEMPTS, RETRY_DELAYS_MS, connector).await
69}
70
71pub async fn connect_with_retry_using<T, F, Fut>(
72 connect_options: PgConnectOptions,
73 max_attempts: u32,
74 delays_ms: &[u64],
75 connector: F,
76) -> DatabaseResult<T>
77where
78 T: Send,
79 F: Fn(PgConnectOptions) -> Fut + Send + Sync,
80 Fut: Future<Output = Result<T, sqlx::Error>> + Send,
81{
82 let cfg = RetryConfig {
83 max_attempts,
84 base_delay: Duration::from_millis(delays_ms.first().copied().unwrap_or(100)),
85 max_delay: Duration::from_millis(delays_ms.iter().copied().max().unwrap_or(1600)),
86 jitter: false,
87 };
88 let classify = |err: &sqlx::Error| {
89 if is_retryable(err) {
90 Outcome::Transient { retry_after: None }
91 } else {
92 Outcome::Permanent
93 }
94 };
95 retry_async(&cfg, "postgres-connect", classify, || {
96 connector(connect_options.clone())
97 })
98 .await
99 .map_err(Into::into)
100}
101
102fn is_retryable(err: &sqlx::Error) -> bool {
103 if let sqlx::Error::Io(io_err) = err
104 && io_err.kind() == std::io::ErrorKind::ConnectionRefused
105 {
106 return true;
107 }
108 let msg = err.to_string();
109 msg.contains("unexpected response from SSLRequest") || msg.contains("starting up")
110}