athena_rs 3.3.0

Database gateway API
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

//! Helpers for building and inspecting Postgres connection pools.
//!
//! Centralizes pool option construction so the rest of the codebase can request
//! pools with consistent limits and provides lightweight snapshots that can be
//! shipped to telemetry/logging.

use crate::client::config::PoolConfig;
use chrono::{DateTime, Utc};
use sqlx::postgres::{PgPool, PgPoolOptions};
use std::time::Duration;

/// Point‑in‑time view of a pool's occupancy.
#[derive(Debug, Clone)]
pub struct ConnectionPoolSnapshot {
    pub client_name: String,
    pub pool_size: u32,
    pub idle_connections: u32,
    pub active_connections: u32,
    pub max_connections: u32,
    pub recorded_at: DateTime<Utc>,
}

impl ConnectionPoolSnapshot {
    pub fn new(client_name: String, pool: &PgPool) -> Self {
        let pool_size: u32 = pool.size();
        let idle: u32 = pool.num_idle().try_into().unwrap_or_default();
        let max_connections: u32 = pool.options().get_max_connections();
        let active_connections: u32 = pool_size.saturating_sub(idle);

        Self {
            client_name,
            pool_size,
            idle_connections: idle,
            active_connections,
            max_connections,
            recorded_at: Utc::now(),
        }
    }
}

/// Constructs Postgres pools with shared defaults so we can change limits in one place.
#[derive(Clone, Debug)]
pub struct ConnectionPoolManager {
    config: PoolConfig,
    max_lifetime: Duration,
    test_before_acquire: bool,
}

impl Default for ConnectionPoolManager {
    fn default() -> Self {
        Self {
            config: PoolConfig::default(),
            max_lifetime: Duration::from_secs(1800),
            test_before_acquire: false,
        }
    }
}

impl ConnectionPoolManager {
    pub fn new(config: PoolConfig) -> Self {
        Self {
            config,
            ..Default::default()
        }
    }

    /// Customize the max lifetime used for pools created by this manager.
    pub fn with_max_lifetime(mut self, max_lifetime: Duration) -> Self {
        self.max_lifetime = max_lifetime;
        self
    }

    /// Enable/disable validating a pooled connection before handing it out.
    pub fn with_test_before_acquire(mut self, test_before_acquire: bool) -> Self {
        self.test_before_acquire = test_before_acquire;
        self
    }

    /// Create a new pool using the managed defaults.
    pub async fn connect(&self, uri: &str) -> Result<PgPool, sqlx::Error> {
        self.build_options().connect(uri).await
    }

    /// Build `PgPoolOptions` with the configured defaults.
    pub fn build_options(&self) -> PgPoolOptions {
        PgPoolOptions::new()
            .max_connections(self.config.max_connections)
            .min_connections(self.config.min_connections)
            .acquire_timeout(self.config.connection_timeout)
            .idle_timeout(Some(self.config.idle_timeout))
            .max_lifetime(Some(self.max_lifetime))
            .test_before_acquire(self.test_before_acquire)
    }
}