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(
48        &self,
49        connection: PooledConnection,
50        client_id: ClientId,
51    ) -> ConnectionLease {
52        ConnectionLease::new(connection, PoolingMode::Session, client_id)
53    }
54
55    /// Process a statement and determine action
56    ///
57    /// Session mode always holds the connection.
58    pub fn on_statement_complete(&self, _lease: &mut ConnectionLease, _sql: &str) -> LeaseAction {
59        LeaseAction::Hold
60    }
61
62    /// Process transaction end
63    ///
64    /// Session mode always holds the connection.
65    pub fn on_transaction_end(&self, _lease: &mut ConnectionLease) -> LeaseAction {
66        LeaseAction::Hold
67    }
68
69    /// Check if connection should be released
70    ///
71    /// Session mode never releases until client disconnects.
72    pub fn should_release(&self, _lease: &ConnectionLease) -> bool {
73        false
74    }
75
76    /// Called when client disconnects
77    pub fn on_client_disconnect(&self, _lease: ConnectionLease) -> LeaseAction {
78        // When client disconnects, reset and return to pool
79        LeaseAction::Reset
80    }
81
82    /// Get the pooling mode
83    pub fn mode(&self) -> PoolingMode {
84        PoolingMode::Session
85    }
86
87    /// Check if prepared statement tracking is enabled
88    pub fn tracks_prepared_statements(&self) -> bool {
89        self.track_prepared_statements
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use crate::connection_pool::ConnectionState;
97    use crate::NodeId;
98    use uuid::Uuid;
99
100    fn create_test_connection() -> PooledConnection {
101        PooledConnection {
102            id: Uuid::new_v4(),
103            node_id: NodeId::new(),
104            created_at: chrono::Utc::now(),
105            last_used: chrono::Utc::now(),
106            state: ConnectionState::InUse,
107            use_count: 1,
108            permit: None,
109            client: None,
110        }
111    }
112
113    #[test]
114    fn test_session_mode_always_holds() {
115        let handler = SessionModeHandler::new();
116        let conn = create_test_connection();
117        let mut lease = handler.create_lease(conn, ClientId::new());
118
119        // Any statement should hold
120        assert_eq!(
121            handler.on_statement_complete(&mut lease, "SELECT 1"),
122            LeaseAction::Hold
123        );
124
125        // Transaction statements should hold
126        assert_eq!(
127            handler.on_statement_complete(&mut lease, "BEGIN"),
128            LeaseAction::Hold
129        );
130        assert_eq!(
131            handler.on_statement_complete(&mut lease, "COMMIT"),
132            LeaseAction::Hold
133        );
134
135        // Transaction end should hold
136        assert_eq!(handler.on_transaction_end(&mut lease), LeaseAction::Hold);
137    }
138
139    #[test]
140    fn test_session_mode_never_releases() {
141        let handler = SessionModeHandler::new();
142        let conn = create_test_connection();
143        let lease = handler.create_lease(conn, ClientId::new());
144
145        assert!(!handler.should_release(&lease));
146    }
147
148    #[test]
149    fn test_session_mode_disconnect_resets() {
150        let handler = SessionModeHandler::new();
151        let conn = create_test_connection();
152        let lease = handler.create_lease(conn, ClientId::new());
153
154        assert_eq!(handler.on_client_disconnect(lease), LeaseAction::Reset);
155    }
156
157    #[test]
158    fn test_mode() {
159        let handler = SessionModeHandler::new();
160        assert_eq!(handler.mode(), PoolingMode::Session);
161    }
162}
heliosdb_proxy/pool/session.rs

heliosdb_proxy/pool/
session.rs
