more error handling

Author: MarinPostma

sqld/src/connection/write_proxy.rs

@@ -194,7 +194,7 @@ impl WriteProxyConnection {
 
         let (builder, new_status, new_frame_no) = match res {
             Ok(res) => res,
-            Err(e @ Error::StreamDisconnect) => {
+            Err(e @ (Error::PrimaryStreamDisconnect | Error::PrimaryStreamMisuse)) => {
                 // drop the connection
                 self.remote_conn.lock().await.take();
                 return Err(e);
@@ -289,23 +289,28 @@ impl RemoteConnection {
         self.request_sender
             .send(req)
             .await
-            .map_err(|_| Error::StreamDisconnect)?;
+            .map_err(|_| Error::PrimaryStreamDisconnect)?;
 
-        'outer: while let Some(resp) = self.response_stream.next().await {
+        while let Some(resp) = self.response_stream.next().await {
             match resp {
                 Ok(resp) => {
-                    // todo: handle interuption
-                    if resp.request_id != request_id {
-                        todo!("stream misuse: connection should be serialized");
+                    // there was an interuption, and we moved to the next query
+                    if resp.request_id > request_id {
+                        return Err(Error::PrimaryStreamInterupted)
                     }
 
-                    if !response_cb(resp.response.unwrap())? {
-                        break 'outer;
+                    // we can ignore response for previously interupted requests
+                    if resp.request_id < request_id {
+                        continue;
+                    }
+
+                    if !response_cb(resp.response.ok_or(Error::PrimaryStreamMisuse)?)? {
+                        break;
                     }
                 }
                 Err(e) => {
-                    tracing::error!("received error from connection stream: {e}");
-                    return Err(Error::StreamDisconnect);
+                    tracing::error!("received an error from connection stream: {e}");
+                    return Err(Error::PrimaryStreamDisconnect);
                 }
             }
         }
@@ -325,16 +330,16 @@ impl RemoteConnection {
             match response {
                 exec_resp::Response::ProgramResp(resp) => {
                     for step in resp.steps {
-                        let Some(step) = step.step else {panic!("invalid pgm")};
+                        let Some(step) = step.step else { return Err(Error::PrimaryStreamMisuse) };
                         match step {
                             Step::Init(_) => builder.init(&builder_config)?,
                             Step::BeginStep(_) => builder.begin_step()?,
                             Step::FinishStep(FinishStep {
                                 affected_row_count,
                                 last_insert_rowid,
                             }) => builder.finish_step(affected_row_count, last_insert_rowid)?,
-                            Step::StepError(StepError { error }) => builder
-                                .step_error(crate::error::Error::RpcQueryError(error.unwrap()))?,
+                            Step::StepError(StepError { error: Some(err) }) => builder
+                                .step_error(crate::error::Error::RpcQueryError(err))?,
                             Step::ColsDescription(ColsDescription { columns }) => {
                                 let cols = columns.iter().map(|c| Column {
                                     name: &c.name,
@@ -364,12 +369,12 @@ impl RemoteConnection {
                                 builder.finish(last_frame_no, txn_status)?;
                                 return Ok(false);
                             }
-                            _ => todo!("invalid request"),
+                            _ => return Err(Error::PrimaryStreamMisuse),
                         }
                     }
                 }
-                exec_resp::Response::DescribeResp(_) => todo!("invalid resp"),
-                exec_resp::Response::Error(_) => todo!(),
+                exec_resp::Response::DescribeResp(_) => return Err(Error::PrimaryStreamMisuse),
+                exec_resp::Response::Error(e) => return Err(Error::RpcQueryError(e)),
             }
 
             Ok(true)
@@ -409,12 +414,12 @@ impl RemoteConnection {
                         is_explain: resp.is_explain,
                         is_readonly: resp.is_readonly,
                     });
+
+                    Ok(false)
                 }
-                exec_resp::Response::Error(_) => todo!(),
-                exec_resp::Response::ProgramResp(_) => todo!(),
+                exec_resp::Response::Error(e) => Err(Error::RpcQueryError(e)),
+                exec_resp::Response::ProgramResp(_) => Err(Error::PrimaryStreamMisuse),
             }
-
-            Ok(false)
         };
 
         self.make_request(
@@ -423,7 +428,7 @@ impl RemoteConnection {
         )
         .await?;
 
-        Ok(out.unwrap())
+        out.ok_or(Error::PrimaryStreamMisuse)
     }
 }
 

sqld/src/error.rs

@@ -79,8 +79,13 @@ pub enum Error {
     ConflictingRestoreParameters,
     #[error("failed to fork database: {0}")]
     Fork(#[from] ForkError),
+
     #[error("Connection with primary broken")]
-    StreamDisconnect,
+    PrimaryStreamDisconnect,
+    #[error("Proxy protocal misuse")]
+    PrimaryStreamMisuse,
+    #[error("Proxy request interupted")]
+    PrimaryStreamInterupted,
 }
 
 trait ResponseError: std::error::Error {
@@ -131,7 +136,7 @@ impl IntoResponse for Error {
             LoadDumpExistingDb => self.format_err(StatusCode::BAD_REQUEST),
             ConflictingRestoreParameters => self.format_err(StatusCode::BAD_REQUEST),
             Fork(e) => e.into_response(),
-            StreamDisconnect => self.format_err(StatusCode::INTERNAL_SERVER_ERROR),
+            PrimaryStreamDisconnect => self.format_err(StatusCode::INTERNAL_SERVER_ERROR),
         }
     }
 }