heliosdb_proxy/pool/
session.rs

1//! Session Mode Handler
2//!
3//! Implements session pooling mode where connections are held
4//! for the entire client session lifetime.
5
6use super::lease::{ClientId, ConnectionLease, LeaseAction};
7use super::mode::PoolingMode;
8use crate::connection_pool::PooledConnection;
9
10/// Session mode handler
11///
12/// In session mode, a connection is held for the entire client session.
13/// This provides 1:1 client-to-backend mapping, which is:
14/// - Safest for all PostgreSQL features
15/// - Compatible with server-side prepared statements
16/// - Compatible with LISTEN/NOTIFY
17/// - Compatible with temp tables and session variables
18///
19/// The downside is less connection sharing between clients.
20pub struct SessionModeHandler {
21    /// Whether to track prepared statements
22    track_prepared_statements: bool,
23}
24
25impl Default for SessionModeHandler {
26    fn default() -> Self {
27        Self::new()
28    }
29}
30
31impl SessionModeHandler {
32    /// Create a new session mode handler
33    pub fn new() -> Self {
34        Self {
35            track_prepared_statements: false,
36        }
37    }
38
39    /// Create with prepared statement tracking enabled
40    pub fn with_prepared_tracking() -> Self {
41        Self {
42            track_prepared_statements: true,
43        }
44    }
45
46    /// Create a lease for this mode
47    pub fn create_lease(&self, connection: PooledConnection, client_id: ClientId) -> ConnectionLease {
48        ConnectionLease::new(connection, PoolingMode::Session, client_id)
49    }
50
51    /// Process a statement and determine action
52    ///
53    /// Session mode always holds the connection.
54    pub fn on_statement_complete(&self, _lease: &mut ConnectionLease, _sql: &str) -> LeaseAction {
55        LeaseAction::Hold
56    }
57
58    /// Process transaction end
59    ///
60    /// Session mode always holds the connection.
61    pub fn on_transaction_end(&self, _lease: &mut ConnectionLease) -> LeaseAction {
62        LeaseAction::Hold
63    }
64
65    /// Check if connection should be released
66    ///
67    /// Session mode never releases until client disconnects.
68    pub fn should_release(&self, _lease: &ConnectionLease) -> bool {
69        false
70    }
71
72    /// Called when client disconnects
73    pub fn on_client_disconnect(&self, _lease: ConnectionLease) -> LeaseAction {
74        // When client disconnects, reset and return to pool
75        LeaseAction::Reset
76    }
77
78    /// Get the pooling mode
79    pub fn mode(&self) -> PoolingMode {
80        PoolingMode::Session
81    }
82
83    /// Check if prepared statement tracking is enabled
84    pub fn tracks_prepared_statements(&self) -> bool {
85        self.track_prepared_statements
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92    use crate::connection_pool::ConnectionState;
93    use crate::NodeId;
94    use uuid::Uuid;
95
96    fn create_test_connection() -> PooledConnection {
97        PooledConnection {
98            id: Uuid::new_v4(),
99            node_id: NodeId::new(),
100            created_at: chrono::Utc::now(),
101            last_used: chrono::Utc::now(),
102            state: ConnectionState::InUse,
103            use_count: 1,
104            permit: None,
105            client: None,
106        }
107    }
108
109    #[test]
110    fn test_session_mode_always_holds() {
111        let handler = SessionModeHandler::new();
112        let conn = create_test_connection();
113        let mut lease = handler.create_lease(conn, ClientId::new());
114
115        // Any statement should hold
116        assert_eq!(
117            handler.on_statement_complete(&mut lease, "SELECT 1"),
118            LeaseAction::Hold
119        );
120
121        // Transaction statements should hold
122        assert_eq!(
123            handler.on_statement_complete(&mut lease, "BEGIN"),
124            LeaseAction::Hold
125        );
126        assert_eq!(
127            handler.on_statement_complete(&mut lease, "COMMIT"),
128            LeaseAction::Hold
129        );
130
131        // Transaction end should hold
132        assert_eq!(handler.on_transaction_end(&mut lease), LeaseAction::Hold);
133    }
134
135    #[test]
136    fn test_session_mode_never_releases() {
137        let handler = SessionModeHandler::new();
138        let conn = create_test_connection();
139        let lease = handler.create_lease(conn, ClientId::new());
140
141        assert!(!handler.should_release(&lease));
142    }
143
144    #[test]
145    fn test_session_mode_disconnect_resets() {
146        let handler = SessionModeHandler::new();
147        let conn = create_test_connection();
148        let lease = handler.create_lease(conn, ClientId::new());
149
150        assert_eq!(handler.on_client_disconnect(lease), LeaseAction::Reset);
151    }
152
153    #[test]
154    fn test_mode() {
155        let handler = SessionModeHandler::new();
156        assert_eq!(handler.mode(), PoolingMode::Session);
157    }
158}
heliosdb_proxy/pool/session.rs

heliosdb_proxy/pool/
session.rs
