Correctly handle RESPONSE_TOO_LARGE in Hrana (#530)

Author: honzasp

sqld/src/hrana/batch.rs

@@ -8,7 +8,9 @@ use crate::error::Error as SqldError;
 use crate::hrana::stmt::StmtError;
 use crate::query::{Params, Query};
 use crate::query_analysis::Statement;
-use crate::query_result_builder::{QueryResultBuilder, StepResult, StepResultsBuilder};
+use crate::query_result_builder::{
+    QueryResultBuilder, QueryResultBuilderError, StepResult, StepResultsBuilder,
+};
 
 use super::result_builder::HranaBatchProtoBuilder;
 use super::stmt::{proto_stmt_to_query, stmt_error_from_sqld_error};
@@ -20,6 +22,8 @@ pub enum BatchError {
     TransactionTimeout,
     #[error("Server cannot handle additional transactions")]
     TransactionBusy,
+    #[error("Response is too large")]
+    ResponseTooLarge,
 }
 
 fn proto_cond_to_cond(cond: &proto::BatchCond, max_step_i: usize) -> Result<Cond> {
@@ -149,6 +153,9 @@ fn batch_error_from_sqld_error(sqld_error: SqldError) -> Result<BatchError, Sqld
     Ok(match sqld_error {
         SqldError::LibSqlTxTimeout => BatchError::TransactionTimeout,
         SqldError::LibSqlTxBusy => BatchError::TransactionBusy,
+        SqldError::BuilderError(QueryResultBuilderError::ResponseTooLarge(_)) => {
+            BatchError::ResponseTooLarge
+        }
         sqld_error => return Err(sqld_error),
     })
 }
@@ -158,6 +165,7 @@ impl BatchError {
         match self {
             Self::TransactionTimeout => "TRANSACTION_TIMEOUT",
             Self::TransactionBusy => "TRANSACTION_BUSY",
+            Self::ResponseTooLarge => "RESPONSE_TOO_LARGE",
         }
     }
 }

sqld/src/hrana/stmt.rs

@@ -9,7 +9,7 @@ use crate::error::Error as SqldError;
 use crate::hrana;
 use crate::query::{Params, Query, Value};
 use crate::query_analysis::Statement;
-use crate::query_result_builder::QueryResultBuilder;
+use crate::query_result_builder::{QueryResultBuilder, QueryResultBuilderError};
 
 /// An error during execution of an SQL statement.
 #[derive(thiserror::Error, Debug)]
@@ -43,6 +43,8 @@ pub enum StmtError {
 
     #[error("Operation was blocked{}", .reason.as_ref().map(|msg| format!(": {}", msg)).unwrap_or_default())]
     Blocked { reason: Option<String> },
+    #[error("Response is too large")]
+    ResponseTooLarge,
 }
 
 pub async fn execute_stmt(
@@ -196,6 +198,9 @@ pub fn stmt_error_from_sqld_error(sqld_error: SqldError) -> Result<StmtError, Sq
         SqldError::LibSqlInvalidQueryParams(source) => StmtError::ArgsInvalid { source },
         SqldError::LibSqlTxTimeout => StmtError::TransactionTimeout,
         SqldError::LibSqlTxBusy => StmtError::TransactionBusy,
+        SqldError::BuilderError(QueryResultBuilderError::ResponseTooLarge(_)) => {
+            StmtError::ResponseTooLarge
+        }
         SqldError::Blocked(reason) => StmtError::Blocked { reason },
         SqldError::RusqliteError(rusqlite_error) => match rusqlite_error {
             rusqlite::Error::SqliteFailure(sqlite_error, Some(message)) => StmtError::SqliteError {
@@ -242,6 +247,7 @@ impl StmtError {
             Self::SqliteError { source, .. } => sqlite_error_code(source.code),
             Self::SqlInputError { .. } => "SQL_INPUT_ERROR",
             Self::Blocked { .. } => "BLOCKED",
+            Self::ResponseTooLarge => "RESPONSE_TOO_LARGE",
         }
     }
 }

sqld/src/http/hrana_over_http_1.rs

@@ -138,6 +138,7 @@ fn response_error_response(err: ResponseError) -> hyper::Response<hyper::Body> {
             | StmtError::SqlManyStmts
             | StmtError::ArgsInvalid { .. }
             | StmtError::SqlInputError { .. }
+            | StmtError::ResponseTooLarge
             | StmtError::Blocked { .. } => hyper::StatusCode::BAD_REQUEST,
             StmtError::ArgsBothPositionalAndNamed => hyper::StatusCode::NOT_IMPLEMENTED,
             StmtError::TransactionTimeout | StmtError::TransactionBusy => {