use std::io;
use bytes::{BufMut, Bytes, BytesMut};
use crate::{
Conn, Dirty, Pristine,
auth::Ready,
codec::{
BackendMessage, Bind, Close, CopyResponse, Describe, DiagnosticResponse, Execute, Frame,
FunctionCall, Parse, TransactionStatus,
},
demux::SessionItem,
grammar::frontend,
pre_startup::Terminated,
replication::{BackendReplication, FrontendReplication},
};
#[derive(Debug)]
pub enum SimpleQuery {}
#[derive(Debug)]
pub enum FunctionCalling {}
#[derive(Debug)]
pub enum Building {}
#[derive(Debug)]
pub enum BoundBuilding {}
#[derive(Debug)]
pub enum AwaitingReady {}
#[derive(Debug)]
pub enum CopyIn {}
#[derive(Debug)]
pub enum CopyOut {}
#[derive(Debug)]
pub enum CopyBoth {}
#[derive(Debug)]
pub enum CopyBothClientDone {}
#[derive(Debug)]
pub enum CopyBothServerDone {}
#[derive(Debug)]
pub enum Draining {}
#[derive(Debug)]
pub enum Resetting {}
#[derive(Debug)]
pub enum ResetComplete {}
pub type ErrorResponse = DiagnosticResponse;
pub type Fallible<T, S, C = Pristine> = Result<T, (Conn<S, Draining, C>, ErrorResponse)>;
#[derive(Debug)]
pub enum SimpleTransition<S, C> {
Continue(Conn<S, SimpleQuery, C>, SessionItem),
CopyIn(Conn<S, CopyIn, C>, CopyResponse),
CopyOut(Conn<S, CopyOut, C>, CopyResponse),
CopyBoth(Conn<S, CopyBoth, C>, CopyResponse),
Ready(ReadyState<S, C>),
Error(Conn<S, Draining, C>, ErrorResponse),
}
#[derive(Debug)]
pub enum AwaitingReadyTransition<S, C> {
Continue(Conn<S, AwaitingReady, C>, SessionItem),
Ready(ReadyState<S, C>),
Error(Conn<S, Draining, C>, ErrorResponse),
}
#[derive(Debug)]
pub enum FunctionCallTransition<S, C> {
Response(Conn<S, AwaitingReady, C>, Bytes),
Error(Conn<S, Draining, C>, ErrorResponse),
}
#[derive(Debug)]
pub enum DrainingTransition<S, C> {
Continue(Conn<S, Draining, C>, SessionItem),
Ready(ReadyState<S, C>),
}
#[derive(Debug)]
pub enum CopyOutTransition<S, C> {
Data(Conn<S, CopyOut, C>, Bytes),
Done(Conn<S, AwaitingReady, C>),
Error(Conn<S, Draining, C>, ErrorResponse),
}
#[derive(Debug)]
pub enum CopyInTransition<S, C> {
Error(Conn<S, Draining, C>, ErrorResponse),
}
#[derive(Debug)]
pub enum CopyBothReceive<S, C> {
Data(Conn<S, CopyBoth, C>, Bytes),
Done(Conn<S, CopyBothServerDone, C>),
Error(Conn<S, Draining, C>, ErrorResponse),
}
#[derive(Debug)]
pub enum CopyBothClientDoneReceive<S, C> {
Data(Conn<S, CopyBothClientDone, C>, Bytes),
Done(Conn<S, AwaitingReady, C>),
Error(Conn<S, Draining, C>, ErrorResponse),
}
#[derive(Debug)]
pub enum ReplicationReceive<S, C> {
Message(Conn<S, CopyBoth, C>, BackendReplication),
Done(Conn<S, CopyBothServerDone, C>),
Error(Conn<S, Draining, C>, ErrorResponse),
}
#[derive(Debug)]
pub enum ReplicationClientDoneReceive<S, C> {
Message(Conn<S, CopyBothClientDone, C>, BackendReplication),
Done(Conn<S, AwaitingReady, C>),
Error(Conn<S, Draining, C>, ErrorResponse),
}
pub type ReplicationProjection<S, C> =
Result<ReplicationReceive<S, C>, (Conn<S, CopyBoth, C>, io::Error)>;
pub type ReplicationClientDoneProjection<S, C> =
Result<ReplicationClientDoneReceive<S, C>, (Conn<S, CopyBothClientDone, C>, io::Error)>;
#[derive(Debug)]
pub enum ReadyState<S, C> {
Clean(Conn<S, Ready, C>),
Dirty {
conn: Conn<S, Ready, Dirty>,
status: TransactionStatus,
parameters_changed: bool,
},
}
#[derive(Debug)]
pub enum ResettingTransition<S> {
Continue(Conn<S, Resetting, Dirty>, SessionItem),
Complete(Conn<S, ResetComplete, Dirty>),
Error(Conn<S, Draining, Dirty>, ErrorResponse),
}
#[derive(Debug)]
pub enum ResetCompleteTransition<S> {
Continue(Conn<S, ResetComplete, Dirty>, SessionItem),
Ready(Conn<S, Ready, Pristine>),
Dirty {
conn: Conn<S, Ready, Dirty>,
status: TransactionStatus,
parameters_changed: bool,
},
Error(Conn<S, Draining, Dirty>, ErrorResponse),
}
impl<S, C> Conn<S, Ready, C> {
pub fn push_terminate(self) -> (Conn<S, Terminated, C>, Frame) {
(self.transition(), empty_frame(b'X'))
}
pub fn push_query(self, query: &[u8]) -> io::Result<(Conn<S, SimpleQuery, Dirty>, Frame)> {
Ok((self.transition(), cstr_frame(b'Q', query)?))
}
pub fn push_stateless_query(
self,
query: &[u8],
) -> io::Result<(Conn<S, SimpleQuery, C>, Frame)> {
Ok((self.transition(), cstr_frame(b'Q', query)?))
}
pub fn push_function_call(
self,
message: &FunctionCall,
) -> io::Result<(Conn<S, FunctionCalling, Dirty>, Frame)> {
Ok((self.transition(), message.to_frame()?))
}
pub fn push_stateless_function_call(
self,
message: &FunctionCall,
) -> io::Result<(Conn<S, FunctionCalling, C>, Frame)> {
Ok((self.transition(), message.to_frame()?))
}
pub fn begin_extended(self) -> Conn<S, Building, C> {
self.transition()
}
}
impl<S, C> Conn<S, FunctionCalling, C> {
pub fn offer(
self,
message: BackendMessage,
) -> Result<FunctionCallTransition<S, C>, (Self, BackendMessage)> {
match (
frontend::project_external(frontend::RuntimeState::FunctionCalling, &message),
message,
) {
(
Some(frontend::Event::FunctionResponse),
BackendMessage::FunctionCallResponse(value),
) => Ok(FunctionCallTransition::Response(self.transition(), value)),
(Some(frontend::Event::Error), BackendMessage::ErrorResponse(error)) => {
Ok(FunctionCallTransition::Error(self.transition(), error))
}
(_, other) => Err((self, other)),
}
}
}
impl<S> Conn<S, Ready, Pristine> {
pub fn release(self) -> S {
self.into_transport()
}
}
impl<S> Conn<S, Ready, Dirty> {
pub fn begin_reset(self) -> io::Result<(Conn<S, Resetting, Dirty>, Frame)> {
Ok((
self.transition(),
cstr_frame(b'Q', b"ROLLBACK; DISCARD ALL")?,
))
}
}
impl<S, C> Conn<S, Building, C> {
pub fn push_parse(self, message: &Parse) -> io::Result<(Conn<S, Building, Dirty>, Frame)> {
Ok((self.transition(), message.to_frame()?))
}
pub fn push_describe(self, message: &Describe) -> io::Result<(Self, Frame)> {
Ok((self, message.to_frame()?))
}
pub fn push_bind(self, message: &Bind) -> io::Result<(Conn<S, BoundBuilding, Dirty>, Frame)> {
Ok((self.transition(), message.to_frame()?))
}
pub fn push_close(self, message: &Close) -> io::Result<(Self, Frame)> {
Ok((self, message.to_frame()?))
}
pub fn push_flush(self) -> (Self, Frame) {
(self, empty_frame(b'H'))
}
pub fn push_sync(self) -> (Conn<S, AwaitingReady, C>, Frame) {
(self.transition(), empty_frame(b'S'))
}
}
impl<S, C> Conn<S, BoundBuilding, C> {
pub fn push_parse(self, message: &Parse) -> io::Result<(Conn<S, BoundBuilding, Dirty>, Frame)> {
Ok((self.transition(), message.to_frame()?))
}
pub fn push_bind(self, message: &Bind) -> io::Result<(Conn<S, BoundBuilding, Dirty>, Frame)> {
Ok((self.transition(), message.to_frame()?))
}
pub fn push_describe(self, message: &Describe) -> io::Result<(Self, Frame)> {
Ok((self, message.to_frame()?))
}
pub fn push_execute(self, message: &Execute) -> io::Result<(Self, Frame)> {
Ok((self, message.to_frame()?))
}
pub fn push_close(self, message: &Close) -> io::Result<(Self, Frame)> {
Ok((self, message.to_frame()?))
}
pub fn push_flush(self) -> (Self, Frame) {
(self, empty_frame(b'H'))
}
pub fn push_sync(self) -> (Conn<S, AwaitingReady, C>, Frame) {
(self.transition(), empty_frame(b'S'))
}
}
impl<S, C> Conn<S, SimpleQuery, C> {
pub fn offer(self, item: SessionItem) -> Result<SimpleTransition<S, C>, (Self, SessionItem)> {
match (
project_session_item(frontend::RuntimeState::Simple, &item),
item,
) {
(
Some(frontend::Event::CopyIn),
SessionItem::Message(BackendMessage::CopyInResponse(response)),
) => Ok(SimpleTransition::CopyIn(self.transition(), response)),
(
Some(frontend::Event::CopyOut),
SessionItem::Message(BackendMessage::CopyOutResponse(response)),
) => Ok(SimpleTransition::CopyOut(self.transition(), response)),
(
Some(frontend::Event::CopyBoth),
SessionItem::Message(BackendMessage::CopyBothResponse(response)),
) => Ok(SimpleTransition::CopyBoth(self.transition(), response)),
(
Some(frontend::Event::Ready),
SessionItem::ReadyForQuery {
status,
parameters_changed,
},
) => Ok(SimpleTransition::Ready(ready_state(
self,
status,
parameters_changed,
))),
(
Some(frontend::Event::Error),
SessionItem::Message(BackendMessage::ErrorResponse(error)),
) => Ok(SimpleTransition::Error(self.transition(), error)),
(
Some(frontend::Event::Continue),
item @ (SessionItem::CommandComplete { .. }
| SessionItem::Message(
BackendMessage::RowDescription(_)
| BackendMessage::DataRow(_)
| BackendMessage::EmptyQueryResponse,
)),
) => Ok(SimpleTransition::Continue(self, item)),
(_, item) => Err((self, item)),
}
}
}
impl<S, C> Conn<S, CopyIn, C> {
pub fn push_copy_data(self, data: Bytes) -> (Self, Frame) {
(
self,
Frame {
tag: b'd',
body: data,
},
)
}
pub fn push_copy_done(self) -> (Conn<S, AwaitingReady, C>, Frame) {
(self.transition(), empty_frame(b'c'))
}
pub fn push_copy_fail(self, message: &[u8]) -> io::Result<(Conn<S, AwaitingReady, C>, Frame)> {
Ok((self.transition(), cstr_frame(b'f', message)?))
}
pub fn offer(self, item: SessionItem) -> Result<CopyInTransition<S, C>, (Self, SessionItem)> {
match (
project_session_item(frontend::RuntimeState::CopyIn, &item),
item,
) {
(
Some(frontend::Event::Error),
SessionItem::Message(BackendMessage::ErrorResponse(error)),
) => Ok(CopyInTransition::Error(self.transition(), error)),
(_, item) => Err((self, item)),
}
}
}
impl<S, C> Conn<S, CopyBothClientDone, C> {
pub fn offer(
self,
item: SessionItem,
) -> Result<CopyBothClientDoneReceive<S, C>, (Self, SessionItem)> {
match (
project_session_item(frontend::RuntimeState::CopyBothClientDone, &item),
item,
) {
(
Some(frontend::Event::ReceiveCopyData),
SessionItem::Message(BackendMessage::CopyData(data)),
) => Ok(CopyBothClientDoneReceive::Data(self, data)),
(
Some(frontend::Event::ReceiveCopyDone),
SessionItem::Message(BackendMessage::CopyDone),
) => Ok(CopyBothClientDoneReceive::Done(self.transition())),
(
Some(frontend::Event::Error),
SessionItem::Message(BackendMessage::ErrorResponse(error)),
) => Ok(CopyBothClientDoneReceive::Error(self.transition(), error)),
(_, item) => Err((self, item)),
}
}
}
impl<S, C> CopyBothClientDoneReceive<S, C> {
pub fn decode_replication(self) -> ReplicationClientDoneProjection<S, C> {
match self {
Self::Data(conn, data) => match BackendReplication::decode(data) {
Ok(message) => Ok(ReplicationClientDoneReceive::Message(conn, message)),
Err(error) => Err((conn, error)),
},
Self::Done(conn) => Ok(ReplicationClientDoneReceive::Done(conn)),
Self::Error(conn, error) => Ok(ReplicationClientDoneReceive::Error(conn, error)),
}
}
}
impl<S, C> Conn<S, CopyBothServerDone, C> {
pub fn push_copy_data(self, data: Bytes) -> (Self, Frame) {
(
self,
Frame {
tag: b'd',
body: data,
},
)
}
pub fn push_replication(self, message: &FrontendReplication) -> (Self, Frame) {
self.push_copy_data(message.encode())
}
pub fn push_copy_done(self) -> (Conn<S, AwaitingReady, C>, Frame) {
(self.transition(), empty_frame(b'c'))
}
}
impl<S, C> Conn<S, CopyOut, C> {
pub fn offer(self, item: SessionItem) -> Result<CopyOutTransition<S, C>, (Self, SessionItem)> {
match (
project_session_item(frontend::RuntimeState::CopyOut, &item),
item,
) {
(
Some(frontend::Event::CopyData),
SessionItem::Message(BackendMessage::CopyData(data)),
) => Ok(CopyOutTransition::Data(self, data)),
(Some(frontend::Event::CopyDone), SessionItem::Message(BackendMessage::CopyDone)) => {
Ok(CopyOutTransition::Done(self.transition()))
}
(
Some(frontend::Event::Error),
SessionItem::Message(BackendMessage::ErrorResponse(error)),
) => Ok(CopyOutTransition::Error(self.transition(), error)),
(_, item) => Err((self, item)),
}
}
}
impl<S, C> Conn<S, CopyBoth, C> {
pub fn push_copy_data(self, data: Bytes) -> (Self, Frame) {
(
self,
Frame {
tag: b'd',
body: data,
},
)
}
pub fn push_replication(self, message: &FrontendReplication) -> (Self, Frame) {
self.push_copy_data(message.encode())
}
pub fn push_copy_done(self) -> (Conn<S, CopyBothClientDone, C>, Frame) {
(self.transition(), empty_frame(b'c'))
}
pub fn offer(self, item: SessionItem) -> Result<CopyBothReceive<S, C>, (Self, SessionItem)> {
match (
project_session_item(frontend::RuntimeState::CopyBoth, &item),
item,
) {
(
Some(frontend::Event::ReceiveCopyData),
SessionItem::Message(BackendMessage::CopyData(data)),
) => Ok(CopyBothReceive::Data(self, data)),
(
Some(frontend::Event::ReceiveCopyDone),
SessionItem::Message(BackendMessage::CopyDone),
) => Ok(CopyBothReceive::Done(self.transition())),
(
Some(frontend::Event::Error),
SessionItem::Message(BackendMessage::ErrorResponse(error)),
) => Ok(CopyBothReceive::Error(self.transition(), error)),
(_, item) => Err((self, item)),
}
}
}
impl<S, C> CopyBothReceive<S, C> {
pub fn decode_replication(self) -> ReplicationProjection<S, C> {
match self {
Self::Data(conn, data) => match BackendReplication::decode(data) {
Ok(message) => Ok(ReplicationReceive::Message(conn, message)),
Err(error) => Err((conn, error)),
},
Self::Done(conn) => Ok(ReplicationReceive::Done(conn)),
Self::Error(conn, error) => Ok(ReplicationReceive::Error(conn, error)),
}
}
}
impl<S, C> Conn<S, Draining, C> {
pub fn offer(self, item: SessionItem) -> DrainingTransition<S, C> {
match (
project_session_item(frontend::RuntimeState::Draining, &item),
item,
) {
(
Some(frontend::Event::Ready),
SessionItem::ReadyForQuery {
status,
parameters_changed,
},
) => DrainingTransition::Ready(ready_state(self, status, parameters_changed)),
(_, item) => DrainingTransition::Continue(self, item),
}
}
}
impl<S, C> Conn<S, AwaitingReady, C> {
pub fn offer(self, item: SessionItem) -> AwaitingReadyTransition<S, C> {
match (
project_session_item(frontend::RuntimeState::AwaitingReady, &item),
item,
) {
(
Some(frontend::Event::Ready),
SessionItem::ReadyForQuery {
status,
parameters_changed,
},
) => AwaitingReadyTransition::Ready(ready_state(self, status, parameters_changed)),
(
Some(frontend::Event::Error),
SessionItem::Message(BackendMessage::ErrorResponse(error)),
) => AwaitingReadyTransition::Error(self.transition(), error),
(_, item) => AwaitingReadyTransition::Continue(self, item),
}
}
}
impl<S> Conn<S, Resetting, Dirty> {
#[must_use]
pub fn offer(self, item: SessionItem) -> ResettingTransition<S> {
match (
project_session_item(frontend::RuntimeState::Resetting, &item),
item,
) {
(Some(frontend::Event::DiscardComplete), SessionItem::CommandComplete { tag, .. })
if tag == b"DISCARD ALL".as_slice() =>
{
ResettingTransition::Complete(self.transition())
}
(
Some(frontend::Event::Error),
SessionItem::Message(BackendMessage::ErrorResponse(error)),
) => ResettingTransition::Error(self.transition(), error),
(_, item) => ResettingTransition::Continue(self, item),
}
}
}
impl<S> Conn<S, ResetComplete, Dirty> {
#[must_use]
pub fn offer(self, item: SessionItem) -> ResetCompleteTransition<S> {
match (
project_session_item(frontend::RuntimeState::ResetComplete, &item),
item,
) {
(
Some(frontend::Event::ReadyClean),
SessionItem::ReadyForQuery {
status: TransactionStatus::Idle,
parameters_changed: false,
},
) => ResetCompleteTransition::Ready(self.transition()),
(
Some(frontend::Event::ReadyClean | frontend::Event::ReadyDirty),
SessionItem::ReadyForQuery {
status,
parameters_changed,
},
) => ResetCompleteTransition::Dirty {
conn: self.transition(),
status,
parameters_changed,
},
(
Some(frontend::Event::Error),
SessionItem::Message(BackendMessage::ErrorResponse(error)),
) => ResetCompleteTransition::Error(self.transition(), error),
(_, item) => ResetCompleteTransition::Continue(self, item),
}
}
}
impl<S, P> Conn<S, P, Pristine> {
pub fn mark_dirty(self) -> Conn<S, P, Dirty> {
self.transition()
}
}
fn project_session_item(
state: frontend::RuntimeState,
item: &SessionItem,
) -> Option<frontend::Event> {
match item {
SessionItem::Message(message) => frontend::project_external(state, message),
SessionItem::CommandComplete { tag, .. } => {
frontend::project_external(state, &BackendMessage::CommandComplete(tag.clone()))
}
SessionItem::ReadyForQuery { status, .. } => {
frontend::project_external(state, &BackendMessage::ReadyForQuery(*status))
}
}
}
fn cstr_frame(tag: u8, value: &[u8]) -> io::Result<Frame> {
if value.contains(&0) {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"message string contains a NUL byte",
));
}
let mut body = BytesMut::with_capacity(value.len() + 1);
body.extend_from_slice(value);
body.put_u8(0);
Ok(Frame {
tag,
body: body.freeze(),
})
}
fn empty_frame(tag: u8) -> Frame {
Frame {
tag,
body: Bytes::new(),
}
}
fn ready_state<S, P, C>(
conn: Conn<S, P, C>,
status: TransactionStatus,
parameters_changed: bool,
) -> ReadyState<S, C> {
if status == TransactionStatus::Idle && !parameters_changed {
ReadyState::Clean(conn.transition())
} else {
ReadyState::Dirty {
conn: conn.transition(),
status,
parameters_changed,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn extended_building_self_loops_then_syncs() {
let ready: Conn<(), Ready> = Conn::new(()).transition();
let building = ready.begin_extended();
let (building, _) = building
.push_parse(&Parse {
statement: Bytes::from_static(b"statement"),
query: Bytes::from_static(b"select $1"),
parameter_types: vec![23],
})
.expect("valid Parse");
let (bound, _) = building
.push_bind(&Bind {
portal: Bytes::from_static(b"portal"),
statement: Bytes::from_static(b"statement"),
parameter_formats: vec![0],
parameters: vec![Some(Bytes::from_static(b"42"))],
result_formats: vec![0],
})
.expect("valid Bind");
let (bound, _) = bound
.push_execute(&Execute {
portal: Bytes::from_static(b"portal"),
max_rows: 0,
})
.expect("valid Execute");
let (awaiting_ready, sync) = bound.push_sync();
assert_eq!(sync.tag, b'S');
awaiting_ready.into_transport();
}
#[test]
fn function_call_requires_result_then_ready() {
let ready: Conn<(), Ready> = Conn::new(()).transition();
let call = FunctionCall {
function_oid: 42,
argument_formats: vec![1],
arguments: vec![Some(Bytes::from_static(b"argument"))],
result_format: 1,
};
let (calling, frame) = ready.push_function_call(&call).unwrap();
assert_eq!(frame.tag, b'F');
let FunctionCallTransition::Response(awaiting_ready, result) = calling
.offer(BackendMessage::FunctionCallResponse(Bytes::from_static(
b"result",
)))
.unwrap()
else {
panic!("function result projected to the wrong branch")
};
assert_eq!(result, Bytes::from_static(b"result"));
let AwaitingReadyTransition::Ready(ReadyState::Clean(ready)) =
awaiting_ready.offer(SessionItem::ReadyForQuery {
status: TransactionStatus::Idle,
parameters_changed: false,
})
else {
panic!("function call did not return to ready")
};
ready.into_transport();
}
#[test]
fn ready_session_can_terminate_gracefully() {
let ready: Conn<(), Ready> = Conn::new(()).transition();
let (terminated, frame) = ready.push_terminate();
assert_eq!(frame.tag, b'X');
assert!(frame.body.is_empty());
terminated.into_transport();
}
#[test]
fn copy_both_waits_for_both_half_closes() {
use crate::grammar::frontend::{Event, RuntimeFsm, RuntimeState};
let mut client_first = RuntimeFsm::new();
client_first.step(Event::Query).unwrap();
client_first.step(Event::CopyBoth).unwrap();
let open: Conn<(), CopyBoth> = Conn::new(()).transition();
let (client_done, frame) = open.push_copy_done();
client_first.step(Event::SendCopyDone).unwrap();
assert_eq!(frame.tag, b'c');
let CopyBothClientDoneReceive::Data(client_done, data) = client_done
.offer(SessionItem::Message(BackendMessage::CopyData(
Bytes::from_static(b"after client close"),
)))
.unwrap()
else {
panic!("backend data projected to the wrong branch")
};
client_first.step(Event::ReceiveCopyData).unwrap();
assert_eq!(data, Bytes::from_static(b"after client close"));
let CopyBothClientDoneReceive::Done(awaiting) = client_done
.offer(SessionItem::Message(BackendMessage::CopyDone))
.unwrap()
else {
panic!("backend close projected to the wrong branch")
};
client_first.step(Event::ReceiveCopyDone).unwrap();
assert_eq!(client_first.state(), RuntimeState::AwaitingReady);
awaiting.into_transport();
let mut server_first = RuntimeFsm::new();
server_first.step(Event::Query).unwrap();
server_first.step(Event::CopyBoth).unwrap();
let open: Conn<(), CopyBoth> = Conn::new(()).transition();
let CopyBothReceive::Done(server_done) = open
.offer(SessionItem::Message(BackendMessage::CopyDone))
.unwrap()
else {
panic!("backend close projected to the wrong branch")
};
server_first.step(Event::ReceiveCopyDone).unwrap();
let (server_done, data) =
server_done.push_copy_data(Bytes::from_static(b"after server close"));
server_first.step(Event::SendCopyData).unwrap();
assert_eq!(data.tag, b'd');
let (awaiting, done) = server_done.push_copy_done();
server_first.step(Event::SendCopyDone).unwrap();
assert_eq!(done.tag, b'c');
assert_eq!(server_first.state(), RuntimeState::AwaitingReady);
awaiting.into_transport();
}
#[test]
fn copy_in_can_receive_an_early_backend_error() {
let copy: Conn<(), CopyIn> = Conn::new(()).transition();
let error = DiagnosticResponse {
fields: vec![crate::codec::DiagnosticField {
code: b'M',
value: Bytes::from_static(b"copy cancelled"),
}],
};
let CopyInTransition::Error(draining, received) = copy
.offer(SessionItem::Message(BackendMessage::ErrorResponse(
error.clone(),
)))
.unwrap();
assert_eq!(received, error);
let DrainingTransition::Ready(ReadyState::Clean(ready)) =
draining.offer(SessionItem::ReadyForQuery {
status: TransactionStatus::Idle,
parameters_changed: false,
})
else {
panic!("COPY failure did not drain to readiness")
};
ready.release();
}
#[test]
fn copy_both_projects_typed_replication_without_losing_connection() {
let open: Conn<(), CopyBoth> = Conn::new(()).transition();
let status = FrontendReplication::StandbyStatus {
written: 10,
flushed: 9,
applied: 8,
client_time: 7,
reply_requested: true,
};
let (open, frame) = open.push_replication(&status);
assert_eq!(frame.body, status.encode());
let keepalive = BackendReplication::PrimaryKeepalive {
wal_end: 11,
server_time: 12,
reply_requested: true,
};
let receive = open
.offer(SessionItem::Message(BackendMessage::CopyData(
keepalive.encode(),
)))
.unwrap();
let ReplicationReceive::Message(open, decoded) = receive.decode_replication().unwrap()
else {
panic!("keepalive projected to the wrong branch")
};
assert_eq!(decoded, keepalive);
open.into_transport();
let open: Conn<(), CopyBoth> = Conn::new(()).transition();
let receive = open
.offer(SessionItem::Message(BackendMessage::CopyData(
Bytes::from_static(b"kshort"),
)))
.unwrap();
let (open, _) = receive.decode_replication().unwrap_err();
open.into_transport();
}
#[test]
fn transaction_status_taints_ready_connection() {
let query: Conn<(), SimpleQuery> = Conn::new(()).transition();
let transition = query
.offer(SessionItem::ReadyForQuery {
status: TransactionStatus::InTransaction,
parameters_changed: false,
})
.expect("ReadyForQuery is valid evidence");
let SimpleTransition::Ready(ReadyState::Dirty {
conn,
status: TransactionStatus::InTransaction,
parameters_changed: false,
}) = transition
else {
panic!("transaction should taint readiness")
};
conn.into_transport();
}
#[test]
fn changed_parameters_taint_idle_connection() {
let query: Conn<(), SimpleQuery> = Conn::new(()).transition();
let transition = query
.offer(SessionItem::ReadyForQuery {
status: TransactionStatus::Idle,
parameters_changed: true,
})
.expect("ReadyForQuery is valid evidence");
let SimpleTransition::Ready(ReadyState::Dirty {
conn,
status: TransactionStatus::Idle,
parameters_changed: true,
}) = transition
else {
panic!("parameter change should taint readiness")
};
conn.into_transport();
}
#[test]
fn simple_queries_are_dirty_unless_inspection_proves_them_stateless() {
fn require_dirty<S>(conn: Conn<S, Ready, Dirty>) {
conn.into_transport();
}
let ready: Conn<(), Ready> = Conn::new(()).transition();
let (query, _) = ready.push_query(b"LISTEN events").unwrap();
let SimpleTransition::Ready(ReadyState::Clean(dirty)) = query
.offer(SessionItem::ReadyForQuery {
status: TransactionStatus::Idle,
parameters_changed: false,
})
.unwrap()
else {
panic!("idle readiness should preserve the query's dirty evidence")
};
require_dirty(dirty);
let ready: Conn<(), Ready> = Conn::new(()).transition();
let (query, _) = ready.push_stateless_query(b"SELECT 1").unwrap();
let SimpleTransition::Ready(ReadyState::Clean(pristine)) = query
.offer(SessionItem::ReadyForQuery {
status: TransactionStatus::Idle,
parameters_changed: false,
})
.unwrap()
else {
panic!("stateless query should retain pristine evidence")
};
pristine.release();
}
#[test]
fn discard_all_evidence_recovers_pool_cleanliness() {
let ready: Conn<(), Ready> = Conn::new(()).transition();
let (resetting, frame) = ready.mark_dirty().begin_reset().unwrap();
assert_eq!(frame.body, Bytes::from_static(b"ROLLBACK; DISCARD ALL\0"));
let ResettingTransition::Continue(resetting, _) =
resetting.offer(SessionItem::CommandComplete {
tag: Bytes::from_static(b"ROLLBACK"),
command: crate::demux::CommandIndex(0),
notices: vec![],
})
else {
panic!("ROLLBACK incorrectly completed reset")
};
let ResettingTransition::Complete(reset_complete) =
resetting.offer(SessionItem::CommandComplete {
tag: Bytes::from_static(b"DISCARD ALL"),
command: crate::demux::CommandIndex(1),
notices: vec![],
})
else {
panic!("DISCARD ALL did not advance reset")
};
let ResetCompleteTransition::Ready(ready) =
reset_complete.offer(SessionItem::ReadyForQuery {
status: TransactionStatus::Idle,
parameters_changed: false,
})
else {
panic!("clean ready evidence did not restore pristine state")
};
ready.release();
}
}