Skip to main content

fraiseql_wire/connection/
state.rs

1//! Connection state machine
2
3use crate::{Error, Result};
4
5/// Connection state
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum ConnectionState {
8    /// Initial state (not connected)
9    Initial,
10
11    /// Startup sent, awaiting authentication request
12    AwaitingAuth,
13
14    /// Authentication in progress
15    Authenticating,
16
17    /// Idle (ready for query)
18    Idle,
19
20    /// Query in progress
21    QueryInProgress,
22
23    /// Reading query results
24    ReadingResults,
25
26    /// Closed
27    Closed,
28}
29
30impl ConnectionState {
31    /// Check if transition is valid
32    pub fn can_transition_to(&self, next: ConnectionState) -> bool {
33        use ConnectionState::*;
34
35        matches!(
36            (self, next),
37            (Initial, AwaitingAuth)
38                | (AwaitingAuth, Authenticating)
39                | (Authenticating, Idle)
40                | (Idle, QueryInProgress)
41                | (QueryInProgress, ReadingResults)
42                | (ReadingResults, Idle)
43                | (_, Closed)
44        )
45    }
46
47    /// Transition to new state
48    pub fn transition(&mut self, next: ConnectionState) -> Result<()> {
49        if !self.can_transition_to(next) {
50            return Err(Error::InvalidState {
51                expected: format!("valid transition from {:?}", self),
52                actual: format!("{:?}", next),
53            });
54        }
55        *self = next;
56        Ok(())
57    }
58}
59
60impl std::fmt::Display for ConnectionState {
61    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62        match self {
63            Self::Initial => write!(f, "initial"),
64            Self::AwaitingAuth => write!(f, "awaiting_auth"),
65            Self::Authenticating => write!(f, "authenticating"),
66            Self::Idle => write!(f, "idle"),
67            Self::QueryInProgress => write!(f, "query_in_progress"),
68            Self::ReadingResults => write!(f, "reading_results"),
69            Self::Closed => write!(f, "closed"),
70        }
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    #[test]
79    fn test_valid_transitions() {
80        let mut state = ConnectionState::Initial;
81        assert!(state.transition(ConnectionState::AwaitingAuth).is_ok());
82        assert!(state.transition(ConnectionState::Authenticating).is_ok());
83        assert!(state.transition(ConnectionState::Idle).is_ok());
84    }
85
86    #[test]
87    fn test_invalid_transition() {
88        let mut state = ConnectionState::Initial;
89        assert!(state.transition(ConnectionState::Idle).is_err());
90    }
91
92    #[test]
93    fn test_close_from_any_state() {
94        let mut state = ConnectionState::QueryInProgress;
95        assert!(state.transition(ConnectionState::Closed).is_ok());
96    }
97}