feat: support for cancel requests in handshake

This commit is contained in:
Jindřich Moravec 2024-01-22 21:59:39 +01:00
parent f9fb8f0670
commit ebabf50291
2 changed files with 9 additions and 1 deletions

View file

@ -8,6 +8,7 @@ use crate::writer::backend::BackendProtoWriter;
use crate::writer::protowriter::ProtoFlush;
/// Performs server-side handshake with the client until ending it with `ReadyForQuery` message.
/// Client can send `CancelRequest` message instead of `StartupMessage` to cancel the request.
/// For more info visit the [`55.2.1. Start-up`](https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-START-UP)
pub async fn do_server_handshake(
writer: &mut (impl BackendProtoWriter + ProtoFlush),
@ -27,12 +28,16 @@ pub async fn do_server_handshake(
}
}
// Wait for mandatory StartupMessage
// Wait for mandatory StartupMessage or CancelRequest
let startup_message = match &reader.peek_special_message().await? {
Some(msg @ SpecialMessage::StartupMessage(data)) => {
reader.consume_special_message(msg).await?;
data.clone()
}
Some(msg @ SpecialMessage::CancelRequest(data)) => {
reader.consume_special_message(msg).await?;
return Err(ServerHandshakeError::IsCancelRequest(data.clone()));
}
_ => {
return Err(ServerHandshakeError::MissingStartupMessage);
}