Skip to main content

fraiseql_wire/connection/
state.rs

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