systemprompt_database/services/postgres/
connection.rs

1//! Initial-connect retry policy for `PostgresProvider`.
2//!
3//! Wraps the first `PgPool` connect in a bounded exponential backoff so
4//! transient startup races (Postgres still booting, SSL handshake racing
5//! the TCP listener) recover without surfacing as user-visible failures.
6//! The retry loop intentionally targets a narrow set of error shapes so
7//! permanent failures (auth, missing database, bad URL) fail fast. The
8//! backoff itself runs on [`crate::resilience::retry::retry_async`].
9//!
10//! Copyright (c) systemprompt.io — Business Source License 1.1.
11//! See <https://systemprompt.io> for licensing details.
12
13use std::future::Future;
14use std::str::FromStr;
15use std::time::Duration;
16
17use sqlx::postgres::{PgConnectOptions, PgPool, PgPoolOptions};
18
19use crate::error::DatabaseResult;
20use crate::resilience::classify::Outcome;
21use crate::resilience::config::RetryConfig;
22use crate::resilience::retry::retry_async;
23
24const RETRY_DELAYS_MS: &[u64] = &[100, 200, 400, 800, 1600];
25const MAX_ATTEMPTS: u32 = 5;
26
27/// Operator-tunable connection-pool sizing for a `PostgresProvider`.
28///
29/// Only the sizing/lifetime knobs an operator needs to fit the pool to their
30/// Postgres `max_connections` and replica count are exposed; the connect, SSL
31/// and retry behaviour is fixed.
32#[derive(Debug, Clone, Copy)]
33pub struct PoolConfig {
34    pub max_connections: u32,
35    pub min_connections: u32,
36    pub acquire_timeout: Duration,
37    pub idle_timeout: Duration,
38    pub max_lifetime: Duration,
39}
40
41impl Default for PoolConfig {
42    fn default() -> Self {
43        Self {
44            max_connections: 50,
45            min_connections: 0,
46            acquire_timeout: Duration::from_secs(30),
47            idle_timeout: Duration::from_mins(5),
48            max_lifetime: Duration::from_mins(30),
49        }
50    }
51}
52
53#[must_use]
54pub fn build_pool_options(cfg: &PoolConfig) -> PgPoolOptions {
55    PgPoolOptions::new()
56        .max_connections(cfg.max_connections)
57        .min_connections(cfg.min_connections)
58        .max_lifetime(cfg.max_lifetime)
59        .acquire_timeout(cfg.acquire_timeout)
60        .idle_timeout(cfg.idle_timeout)
61}
62
63pub fn connect_options(database_url: &str) -> DatabaseResult<PgConnectOptions> {
64    let options = PgConnectOptions::from_str(database_url)?
65        .application_name("systemprompt")
66        // Why: migrations run DDL on the serving pool and sqlx never invalidates
67        // a connection's prepared statements, so a cached plan would fail with
68        // SQLSTATE 0A000 ("cached plan must not change result type") after an
69        // ALTER TABLE.
70        .statement_cache_capacity(0)
71        .options([("client_min_messages", "warning")]);
72    Ok(options)
73}
74
75pub async fn connect_with_retry(
76    options: PgPoolOptions,
77    connect_options: PgConnectOptions,
78) -> DatabaseResult<PgPool> {
79    let connector = |opts: PgConnectOptions| {
80        let options = options.clone();
81        async move { options.connect_with(opts).await }
82    };
83    connect_with_retry_using(connect_options, MAX_ATTEMPTS, RETRY_DELAYS_MS, connector).await
84}
85
86pub async fn connect_with_retry_using<T, F, Fut>(
87    connect_options: PgConnectOptions,
88    max_attempts: u32,
89    delays_ms: &[u64],
90    connector: F,
91) -> DatabaseResult<T>
92where
93    T: Send,
94    F: Fn(PgConnectOptions) -> Fut + Send + Sync,
95    Fut: Future<Output = Result<T, sqlx::Error>> + Send,
96{
97    let cfg = RetryConfig {
98        max_attempts,
99        base_delay: Duration::from_millis(delays_ms.first().copied().unwrap_or(100)),
100        max_delay: Duration::from_millis(delays_ms.iter().copied().max().unwrap_or(1600)),
101        jitter: false,
102    };
103    let classify = |err: &sqlx::Error| {
104        if is_retryable(err) {
105            Outcome::Transient { retry_after: None }
106        } else {
107            Outcome::Permanent
108        }
109    };
110    retry_async(&cfg, "postgres-connect", classify, || {
111        connector(connect_options.clone())
112    })
113    .await
114    .map_err(Into::into)
115}
116
117fn is_retryable(err: &sqlx::Error) -> bool {
118    if let sqlx::Error::Io(io_err) = err
119        && io_err.kind() == std::io::ErrorKind::ConnectionRefused
120    {
121        return true;
122    }
123    let msg = err.to_string();
124    msg.contains("unexpected response from SSLRequest") || msg.contains("starting up")
125}
systemprompt_database/services/postgres/connection.rs

systemprompt_database/services/postgres/
connection.rs
