error handling

Author: MarinPostma

sqld/src/connection/libsql.rs

@@ -858,7 +858,7 @@ mod test {
             TestBuilder::default(),
         )
         .unwrap();
-        assert_eq!(state, State::Txn);
+        assert_eq!(state, TxnStatus::Txn);
 
         tokio::time::advance(TXN_TIMEOUT * 2).await;
 
@@ -868,7 +868,7 @@ mod test {
             TestBuilder::default(),
         )
         .unwrap();
-        assert_eq!(state, State::Init);
+        assert_eq!(state, TxnStatus::Init);
         assert!(matches!(builder.into_ret()[0], Err(Error::LibSqlTxTimeout)));
     }
 
@@ -902,7 +902,7 @@ mod test {
                     TestBuilder::default(),
                 )
                 .unwrap();
-                assert_eq!(state, State::Txn);
+                assert_eq!(state, TxnStatus::Txn);
                 assert!(builder.into_ret()[0].is_ok());
             });
         }
@@ -945,7 +945,7 @@ mod test {
                     TestBuilder::default(),
                 )
                 .unwrap();
-                assert_eq!(state, State::Txn);
+                assert_eq!(state, TxnStatus::Txn);
                 assert!(builder.into_ret()[0].is_ok());
             }
         })
@@ -963,7 +963,7 @@ mod test {
                     TestBuilder::default(),
                 )
                 .unwrap();
-                assert_eq!(state, State::Txn);
+                assert_eq!(state, TxnStatus::Txn);
                 assert!(builder.into_ret()[0].is_ok());
                 before.elapsed()
             }
@@ -978,7 +978,7 @@ mod test {
                 let (builder, state) =
                     Connection::run(conn, Program::seq(&["COMMIT"]), TestBuilder::default())
                         .unwrap();
-                assert_eq!(state, State::Init);
+                assert_eq!(state, TxnStatus::Init);
                 assert!(builder.into_ret()[0].is_ok());
             }
         })

sqld/src/connection/write_proxy.rs

@@ -187,11 +187,22 @@ impl WriteProxyConnection {
     ) -> Result<(B, TxnStatus)> {
         self.stats.inc_write_requests_delegated();
         *status = TxnStatus::Invalid;
-        let (builder, new_status, new_frame_no) = self
+        let res = self
             .with_remote_conn(auth, self.builder_config, |conn| {
                 Box::pin(conn.execute(pgm, builder))
             })
-            .await?;
+            .await;
+
+        let (builder, new_status, new_frame_no) = match res {
+            Ok(res) => res,
+            Err(e @ Error::StreamDisconnect) => {
+                // drop the connection
+                self.remote_conn.lock().await.take();
+                return Err(e);
+            }
+            Err(e) => return Err(e),
+        };
+
         *status = new_status;
         if let Some(current_frame_no) = new_frame_no {
             self.update_last_write_frame_no(current_frame_no);
@@ -251,9 +262,7 @@ impl RemoteConnection {
         req.metadata_mut()
             .insert_bin(NAMESPACE_METADATA_KEY, namespace);
         auth.upgrade_grpc_request(&mut req);
-        dbg!();
         let response_stream = client.stream_exec(req).await.unwrap().into_inner();
-        dbg!();
 
         Ok(Self {
             response_stream,
@@ -276,14 +285,11 @@ impl RemoteConnection {
             request: Some(rpc::exec_req::Request::Execute(program.into())),
         };
 
-        dbg!();
         self.request_sender.send(req).await.unwrap(); // TODO: the stream was close!
-        dbg!();
         let mut txn_status = TxnStatus::Invalid;
         let mut new_frame_no = None;
 
         'outer: while let Some(resp) = self.response_stream.next().await {
-            dbg!(&resp);
             match resp {
                 Ok(resp) => {
                     if resp.request_id != request_id {
@@ -314,7 +320,6 @@ impl RemoteConnection {
                             Payload::BeginRow(_) => builder.begin_row()?,
                             Payload::AddRowValue(AddRowValue { val }) => {
                                 let value: Value = bincode::deserialize(&val.unwrap().data)
-                                    // something is wrong, better stop right here
                                     .map_err(QueryResultBuilderError::from_any)?;
                                 builder.add_row_value(ValueRef::from(&value))?;
                             }
@@ -324,7 +329,6 @@ impl RemoteConnection {
                                 txn_status = TxnStatus::from(f.state());
                                 new_frame_no = last_frame_no;
                                 builder.finish(last_frame_no, txn_status)?;
-                                dbg!();
                                 break 'outer;
                             }
                             Payload::Error(error) => {
@@ -333,7 +337,10 @@ impl RemoteConnection {
                         }
                     }
                 }
-                Err(_e) => todo!("handle stream error"),
+                Err(e) => {
+                    tracing::error!("received error from connection stream: {e}");
+                    return Err(Error::StreamDisconnect)
+                },
             }
         }
 
@@ -426,7 +433,7 @@ pub mod test {
     use rand::Fill;
 
     use super::*;
-    use crate::query_result_builder::test::test_driver;
+    use crate::{query_result_builder::test::test_driver, rpc::proxy::rpc::{ExecuteResults, query_result::RowResult}};
 
     /// generate an arbitraty rpc value. see build.rs for usage.
     pub fn arbitrary_rpc_value(u: &mut Unstructured) -> arbitrary::Result<Vec<u8>> {
@@ -442,10 +449,55 @@ pub mod test {
         Ok(v.into())
     }
 
+    fn execute_results_to_builder<B: QueryResultBuilder>(
+        execute_result: ExecuteResults,
+        mut builder: B,
+        config: &QueryBuilderConfig,
+    ) -> Result<B> {
+        builder.init(config)?;
+        for result in execute_result.results {
+            match result.row_result {
+                Some(RowResult::Row(rows)) => {
+                    builder.begin_step()?;
+                    builder.cols_description(rows.column_descriptions.iter().map(|c| Column {
+                        name: &c.name,
+                        decl_ty: c.decltype.as_deref(),
+                    }))?;
+
+                    builder.begin_rows()?;
+                    for row in rows.rows {
+                        builder.begin_row()?;
+                        for value in row.values {
+                            let value: Value = bincode::deserialize(&value.data)
+                                // something is wrong, better stop right here
+                                .map_err(QueryResultBuilderError::from_any)?;
+                            builder.add_row_value(ValueRef::from(&value))?;
+                        }
+                        builder.finish_row()?;
+                    }
+
+                    builder.finish_rows()?;
+
+                    builder.finish_step(rows.affected_row_count, rows.last_insert_rowid)?;
+                }
+                Some(RowResult::Error(err)) => {
+                    builder.begin_step()?;
+                    builder.step_error(Error::RpcQueryError(err))?;
+                    builder.finish_step(0, None)?;
+                }
+                None => (),
+            }
+        }
+
+        builder.finish(execute_result.current_frame_no, TxnStatus::Init)?;
+
+        Ok(builder)
+    }
+
     /// In this test, we generate random ExecuteResults, and ensures that the `execute_results_to_builder` drives the builder FSM correctly.
     #[test]
     fn test_execute_results_to_builder() {
-        test_driver(1000, |b| {
+        test_driver(1000, |b| -> std::result::Result<crate::query_result_builder::test::FsmQueryBuilder, Error> {
             let mut data = [0; 10_000];
             data.try_fill(&mut rand::thread_rng()).unwrap();
             let mut un = Unstructured::new(&data);

sqld/src/error.rs

@@ -79,6 +79,8 @@ pub enum Error {
     ConflictingRestoreParameters,
     #[error("failed to fork database: {0}")]
     Fork(#[from] ForkError),
+    #[error("Connection with primary broken")]
+    StreamDisconnect,
 }
 
 trait ResponseError: std::error::Error {

sqld/src/query_result_builder.rs

@@ -15,6 +15,7 @@ pub static TOTAL_RESPONSE_SIZE: AtomicUsize = AtomicUsize::new(0);
 
 #[derive(Debug)]
 pub enum QueryResultBuilderError {
+    /// The response payload is too large
     ResponseTooLarge(u64),
     Internal(anyhow::Error),
 }
@@ -616,6 +617,7 @@ pub mod test {
         fn finish(
             &mut self,
             _last_frame_no: Option<FrameNo>,
+            _txn_status: TxnStatus,
         ) -> Result<(), QueryResultBuilderError> {
             Ok(())
         }
@@ -751,7 +753,7 @@ pub mod test {
                 FinishRow => b.finish_row().unwrap(),
                 FinishRows => b.finish_rows().unwrap(),
                 Finish => {
-                    b.finish(Some(0)).unwrap();
+                    b.finish(Some(0), TxnStatus::Init).unwrap();
                     break;
                 }
                 BuilderError => return b,
@@ -899,6 +901,7 @@ pub mod test {
         fn finish(
             &mut self,
             _last_frame_no: Option<FrameNo>,
+            _txn_status: TxnStatus,
         ) -> Result<(), QueryResultBuilderError> {
             self.maybe_inject_error()?;
             self.transition(Finish)
@@ -947,7 +950,7 @@ pub mod test {
         builder.finish_rows().unwrap();
         builder.finish_step(0, None).unwrap();
 
-        builder.finish(Some(0)).unwrap();
+        builder.finish(Some(0), TxnStatus::Init).unwrap();
     }
 
     #[test]