set connection state to unknown before sending proxy request

Author: MarinPostma

libsqlx/src/database/proxy/connection.rs

@@ -10,9 +10,27 @@ use crate::Result;
 
 use super::WaitFrameNoCb;
 
+#[derive(Debug, Default)]
+enum State {
+    Txn,
+    #[default]
+    Idle,
+    Unknown
+}
+
+impl State {
+    /// Returns `true` if the state is [`Idle`].
+    ///
+    /// [`Idle`]: State::Idle
+    #[must_use]
+    fn is_idle(&self) -> bool {
+        matches!(self, Self::Idle)
+    }
+}
+
 #[derive(Debug, Default)]
 pub(crate) struct ConnState {
-    is_txn: bool,
+    state: State,
     last_frame_no: Option<FrameNo>,
 }
 
@@ -123,6 +141,9 @@ where
                 state: self.state.clone(),
             };
 
+            // set the connection state to unknown before executing on the remote
+            self.state.lock().state = State::Unknown;
+
             self.conn
                 .execute_program(&self.pgm, Box::new(builder))
                 .unwrap();
@@ -144,7 +165,7 @@ where
         pgm: &Program,
         builder: Box<dyn ResultBuilder>,
     ) -> crate::Result<()> {
-        if !self.state.lock().is_txn && pgm.is_read_only() {
+        if self.state.lock().state.is_idle() && pgm.is_read_only() {
             if let Some(frame_no) = self.state.lock().last_frame_no {
                 (self.wait_frame_no_cb)(frame_no);
             }
@@ -162,6 +183,9 @@ where
             // rollback(&mut self.conn.read_db);
             Ok(())
         } else {
+            // we set the state to unknown because until we have received from the actual
+            // connection state from the primary.
+            self.state.lock().state = State::Unknown;
             let builder = ExtractFrameNoBuilder {
                 builder,
                 state: self.state.clone(),
@@ -243,7 +267,11 @@ impl ResultBuilder for ExtractFrameNoBuilder {
     ) -> Result<bool, QueryResultBuilderError> {
         let mut state = self.state.lock();
         state.last_frame_no = frame_no;
-        state.is_txn = is_txn;
+        if is_txn {
+            state.state = State::Txn;
+        } else {
+            state.state = State::Idle;
+        }
         self.builder.finnalize(is_txn, frame_no)
     }
 }