pg_proto/
server_session.rs

1//! Server-role query sessions used when a proxy terminates the client protocol.
2
3use std::io;
4use std::marker::PhantomData;
5
6use bytes::Bytes;
7
8use crate::{
9    Conn, Dirty,
10    auth::Ready,
11    codec::{
12        BackendMessage, Bind, Close, CopyResponse, Describe, DiagnosticResponse, Execute, Frame,
13        FrontendMessage, FunctionCall, Parse, RowDescription, TransactionStatus,
14    },
15    grammar::backend,
16    pre_startup::Terminated,
17    replication::{BackendReplication, FrontendReplication},
18};
19
20#[derive(Debug)]
21/// A client simple query is being served.
22pub enum ServerSimpleQuery {}
23
24#[derive(Debug)]
25/// A simple-query error was sent and readiness must follow.
26pub enum ServerSimpleError {}
27
28#[derive(Debug)]
29/// A legacy function call is being served.
30pub enum ServerFunctionCall {}
31
32#[derive(Debug)]
33/// A function result was sent and readiness must follow.
34pub enum ServerFunctionCallDone {}
35
36#[derive(Debug)]
37/// A function-call error was sent and readiness must follow.
38pub enum ServerFunctionCallError {}
39
40#[derive(Debug)]
41/// The server is accepting an extended-query pipeline.
42pub enum ServerBuilding {}
43
44#[derive(Debug)]
45/// An inspected `Parse` awaits its response.
46pub enum ServerParse {}
47
48#[derive(Debug)]
49/// An inspected `Bind` awaits its response.
50pub enum ServerBind {}
51
52#[derive(Debug)]
53/// An inspected `Describe` awaits its response.
54pub enum ServerDescribe {}
55
56#[derive(Debug)]
57/// An inspected `Execute` is being served.
58pub enum ServerExecute {}
59
60#[derive(Debug)]
61/// An inspected `Close` awaits its response.
62pub enum ServerClose {}
63
64#[derive(Debug)]
65/// A client `Sync` awaits `ReadyForQuery`.
66pub enum ServerSync {}
67
68#[derive(Debug)]
69/// A failed extended pipeline is discarded until `Sync`.
70pub enum ServerExtendedError {}
71
72#[derive(Debug)]
73/// COPY resumes in a simple-query session.
74pub enum CopySimple {}
75
76#[derive(Debug)]
77/// COPY resumes in an extended-query session.
78pub enum CopyExtended {}
79
80/// Maps a COPY resumption marker to its generated nested-session states.
81pub trait CopyResume {
82    /// Generated state while both COPY directions remain open.
83    const BOTH_OPEN_STATE: backend::RuntimeState;
84    /// Generated state after the server closes its COPY direction.
85    const BOTH_SERVER_DONE_STATE: backend::RuntimeState;
86}
87
88impl CopyResume for CopySimple {
89    const BOTH_OPEN_STATE: backend::RuntimeState = backend::RuntimeState::SimpleCopyBoth;
90    const BOTH_SERVER_DONE_STATE: backend::RuntimeState =
91        backend::RuntimeState::SimpleCopyBothServerDone;
92}
93
94impl CopyResume for CopyExtended {
95    const BOTH_OPEN_STATE: backend::RuntimeState = backend::RuntimeState::ExtendedCopyBoth;
96    const BOTH_SERVER_DONE_STATE: backend::RuntimeState =
97        backend::RuntimeState::ExtendedCopyBothServerDone;
98}
99
100#[derive(Debug)]
101/// Server-role COPY IN stream, resumed according to `Resume`.
102pub struct ServerCopyIn<Resume>(PhantomData<Resume>);
103
104#[derive(Debug)]
105/// Client completed a server-role COPY IN stream.
106pub struct ServerCopyInDone<Resume>(PhantomData<Resume>);
107
108#[derive(Debug)]
109/// Client failed a server-role COPY IN stream.
110pub struct ServerCopyInFailed<Resume>(PhantomData<Resume>);
111
112#[derive(Debug)]
113/// Server-role COPY OUT stream, resumed according to `Resume`.
114pub struct ServerCopyOut<Resume>(PhantomData<Resume>);
115
116#[derive(Debug)]
117/// Server completed a server-role COPY OUT stream.
118pub struct ServerCopyOutDone<Resume>(PhantomData<Resume>);
119
120#[derive(Debug)]
121/// Both halves of a COPY BOTH stream remain open.
122pub enum BothOpen {}
123
124#[derive(Debug)]
125/// The client half of a COPY BOTH stream is closed.
126pub enum BothClientDone {}
127
128#[derive(Debug)]
129/// The server half of a COPY BOTH stream is closed.
130pub enum BothServerDone {}
131
132#[derive(Debug)]
133/// Both halves of a COPY BOTH stream are closed.
134pub enum BothDone {}
135
136#[derive(Debug)]
137/// Server-role COPY BOTH stream parameterised by resumption and half-close state.
138pub struct ServerCopyBoth<Resume, Ends>(PhantomData<(Resume, Ends)>);
139
140#[derive(Debug)]
141/// Client failed a server-role COPY BOTH stream.
142pub struct ServerCopyBothFailed<Resume>(PhantomData<Resume>);
143
144/// Client choice while both COPY BOTH directions remain open.
145#[derive(Debug)]
146pub enum ServerCopyBothOpenOffer<S, C, Resume> {
147    /// The client sent one opaque data chunk.
148    Data {
149        /// Connection remaining in COPY BOTH.
150        conn: Conn<S, ServerCopyBoth<Resume, BothOpen>, C>,
151        /// Copy payload.
152        data: Bytes,
153    },
154    /// The client closed its sending half.
155    Done(Conn<S, ServerCopyBoth<Resume, BothClientDone>, C>),
156    /// The client aborted COPY.
157    Fail {
158        /// Failed COPY connection.
159        conn: Conn<S, ServerCopyBothFailed<Resume>, C>,
160        /// Client error message without its terminating NUL.
161        message: Bytes,
162    },
163}
164
165/// Client choice after the server closes its COPY BOTH direction.
166#[derive(Debug)]
167pub enum ServerCopyBothServerDoneOffer<S, C, Resume> {
168    /// The client sent one final opaque data chunk.
169    Data {
170        /// Connection with only the client direction open.
171        conn: Conn<S, ServerCopyBoth<Resume, BothServerDone>, C>,
172        /// Copy payload.
173        data: Bytes,
174    },
175    /// The client closed the remaining direction.
176    Done(Conn<S, ServerCopyBoth<Resume, BothDone>, C>),
177    /// The client aborted COPY.
178    Fail {
179        /// Failed COPY connection.
180        conn: Conn<S, ServerCopyBothFailed<Resume>, C>,
181        /// Client error message without its terminating NUL.
182        message: Bytes,
183    },
184}
185
186/// Typed standby choice while both replication directions remain open.
187#[derive(Debug)]
188pub enum ServerReplicationOpenOffer<S, C, Resume> {
189    /// The standby sent one decoded replication message.
190    Message {
191        /// Connection remaining in COPY BOTH.
192        conn: Conn<S, ServerCopyBoth<Resume, BothOpen>, C>,
193        /// Decoded standby message.
194        message: FrontendReplication,
195    },
196    /// The standby closed its sending half.
197    Done(Conn<S, ServerCopyBoth<Resume, BothClientDone>, C>),
198    /// The standby aborted replication.
199    Fail {
200        /// Failed replication connection.
201        conn: Conn<S, ServerCopyBothFailed<Resume>, C>,
202        /// Standby error message without its terminating NUL.
203        message: Bytes,
204    },
205}
206
207/// Typed standby choice after the walsender closes its direction.
208#[derive(Debug)]
209pub enum ServerReplicationServerDoneOffer<S, C, Resume> {
210    /// The standby sent one final decoded replication message.
211    Message {
212        /// Connection with only the standby direction open.
213        conn: Conn<S, ServerCopyBoth<Resume, BothServerDone>, C>,
214        /// Decoded standby message.
215        message: FrontendReplication,
216    },
217    /// The standby closed the remaining direction.
218    Done(Conn<S, ServerCopyBoth<Resume, BothDone>, C>),
219    /// The standby aborted replication.
220    Fail {
221        /// Failed replication connection.
222        conn: Conn<S, ServerCopyBothFailed<Resume>, C>,
223        /// Standby error message without its terminating NUL.
224        message: Bytes,
225    },
226}
227
228/// Replication projection preserving the open connection when decoding fails.
229pub type ServerReplicationOpenProjection<S, C, Resume> = Result<
230    ServerReplicationOpenOffer<S, C, Resume>,
231    (Conn<S, ServerCopyBoth<Resume, BothOpen>, C>, io::Error),
232>;
233/// Replication projection after server half-close, preserving decode failures.
234pub type ServerReplicationServerDoneProjection<S, C, Resume> = Result<
235    ServerReplicationServerDoneOffer<S, C, Resume>,
236    (
237        Conn<S, ServerCopyBoth<Resume, BothServerDone>, C>,
238        io::Error,
239    ),
240>;
241
242/// Client choice inside a server-role COPY IN sub-session.
243#[derive(Debug)]
244pub enum ServerCopyInOffer<S, C, Resume> {
245    /// The client sent one data chunk.
246    Data {
247        /// Connection remaining in COPY IN.
248        conn: Conn<S, ServerCopyIn<Resume>, C>,
249        /// Copy payload.
250        data: Bytes,
251    },
252    /// The client completed COPY IN.
253    Done(Conn<S, ServerCopyInDone<Resume>, C>),
254    /// The client aborted COPY IN.
255    Fail {
256        /// Failed COPY connection.
257        conn: Conn<S, ServerCopyInFailed<Resume>, C>,
258        /// Client error message without its terminating NUL.
259        message: Bytes,
260    },
261}
262
263/// COPY IN projection preserving the connection and message on mismatch.
264pub type CopyInProjection<S, C, Resume> = Result<
265    ServerCopyInOffer<S, C, Resume>,
266    Box<(Conn<S, ServerCopyIn<Resume>, C>, FrontendMessage)>,
267>;
268/// Result of starting a server-role COPY IN stream.
269pub type CopyInStart<S, C, Resume> = io::Result<(Conn<S, ServerCopyIn<Resume>, C>, Frame)>;
270/// Result of starting a server-role COPY OUT stream.
271pub type CopyOutStart<S, C, Resume> = io::Result<(Conn<S, ServerCopyOut<Resume>, C>, Frame)>;
272/// Result of closing a server-role COPY OUT stream.
273pub type CopyOutCompletion<S, C, Resume> =
274    io::Result<(Conn<S, ServerCopyOutDone<Resume>, C>, Frame)>;
275/// Result of starting a server-role COPY BOTH stream.
276pub type CopyBothStart<S, C, Resume> =
277    io::Result<(Conn<S, ServerCopyBoth<Resume, BothOpen>, C>, Frame)>;
278/// COPY BOTH projection while both directions remain open.
279pub type CopyBothOpenProjection<S, C, Resume> = Result<
280    ServerCopyBothOpenOffer<S, C, Resume>,
281    Box<(
282        Conn<S, ServerCopyBoth<Resume, BothOpen>, C>,
283        FrontendMessage,
284    )>,
285>;
286/// COPY BOTH projection after the server direction closes.
287pub type CopyBothServerDoneProjection<S, C, Resume> = Result<
288    ServerCopyBothServerDoneOffer<S, C, Resume>,
289    Box<(
290        Conn<S, ServerCopyBoth<Resume, BothServerDone>, C>,
291        FrontendMessage,
292    )>,
293>;
294/// Result of closing the server half of COPY BOTH.
295pub type CopyBothServerHalfClose<S, C, Resume> =
296    io::Result<(Conn<S, ServerCopyBoth<Resume, BothServerDone>, C>, Frame)>;
297/// Result of completing COPY BOTH after both halves close.
298pub type CopyBothCompletion<S, C, Resume> =
299    io::Result<(Conn<S, ServerCopyBoth<Resume, BothDone>, C>, Frame)>;
300
301/// External choice offered by a client while the server role is ready.
302#[derive(Debug)]
303pub enum ServerReadyOffer<S, C> {
304    /// A simple query, conservatively marking the session dirty.
305    Query {
306        /// Connection serving the query.
307        conn: Conn<S, ServerSimpleQuery, Dirty>,
308        /// Inspectable and replaceable SQL bytes.
309        query: Bytes,
310    },
311    /// A legacy function-call request.
312    FunctionCall {
313        /// Connection serving the function call.
314        conn: Conn<S, ServerFunctionCall, Dirty>,
315        /// Fully decoded function-call message.
316        message: FunctionCall,
317    },
318    /// One message in an extended-query pipeline.
319    Extended(ServerExtendedOffer<S, C>),
320    /// The client terminated the session.
321    Terminate(Conn<S, Terminated, C>),
322}
323
324/// A response-specific branch of the extended-query building loop.
325#[derive(Debug)]
326pub enum ServerExtendedOffer<S, C> {
327    /// An inspected and reconstructable `Parse` request.
328    Parse {
329        /// Connection awaiting a parse response.
330        conn: Conn<S, ServerParse, Dirty>,
331        /// Decoded request available to application policy.
332        message: Parse,
333    },
334    /// An inspected and reconstructable `Bind` request.
335    Bind {
336        /// Connection awaiting a bind response.
337        conn: Conn<S, ServerBind, Dirty>,
338        /// Decoded request available to application policy.
339        message: Bind,
340    },
341    /// An inspected and reconstructable `Describe` request.
342    Describe {
343        /// Connection awaiting a description response.
344        conn: Conn<S, ServerDescribe, C>,
345        /// Decoded request available to application policy.
346        message: Describe,
347    },
348    /// An inspected and reconstructable `Execute` request.
349    Execute {
350        /// Connection serving the portal execution.
351        conn: Conn<S, ServerExecute, C>,
352        /// Decoded request available to application policy.
353        message: Execute,
354    },
355    /// An inspected and reconstructable `Close` request.
356    Close {
357        /// Connection awaiting a close response.
358        conn: Conn<S, ServerClose, C>,
359        /// Decoded request available to application policy.
360        message: Close,
361    },
362    /// The client requested immediate delivery of buffered responses.
363    Flush(Conn<S, ServerBuilding, C>),
364    /// The client ended the pipeline.
365    Sync(Conn<S, ServerSync, C>),
366}
367
368/// Projection while discarding a failed pipeline up to its synchronisation point.
369#[derive(Debug)]
370pub enum ServerDiscard<S, C> {
371    /// A pipeline message was discarded; continue until synchronisation.
372    Continue(Conn<S, ServerExtendedError, C>),
373    /// `Sync` ended the failed pipeline.
374    Sync(Conn<S, ServerSync, C>),
375}
376
377/// Ready state produced from the status byte sent to the client.
378#[derive(Debug)]
379pub enum ServerReadyState<S, C> {
380    /// Idle readiness retained the existing cleanliness index.
381    Ready(Conn<S, Ready, C>),
382    /// Non-idle readiness made the connection dirty.
383    Dirty {
384        /// Ready connection carrying the dirty marker.
385        conn: Conn<S, Ready, Dirty>,
386        /// Transaction status sent to the client.
387        status: TransactionStatus,
388    },
389}
390
391/// Projection of a client ready-state choice, preserving invalid input.
392pub type ReadyProjection<S, C> =
393    Result<ServerReadyOffer<S, C>, Box<(Conn<S, Ready, C>, FrontendMessage)>>;
394/// Projection of an extended-query choice, preserving invalid input.
395pub type ExtendedProjection<S, Phase, C> =
396    Result<ServerExtendedOffer<S, C>, Box<(Conn<S, Phase, C>, FrontendMessage)>>;
397
398impl<S, C> Conn<S, Ready, C> {
399    /// Projects an inspected client message into the server-role ready state.
400    ///
401    /// # Errors
402    ///
403    /// Returns the unchanged connection and message for choices not yet legal in
404    /// this simple-query projection.
405    pub fn offer_frontend(self, message: FrontendMessage) -> ReadyProjection<S, C> {
406        match (
407            backend::project_external(backend::RuntimeState::Ready, &message),
408            message,
409        ) {
410            (Some(backend::Event::Query), FrontendMessage::Query(query)) => {
411                Ok(ServerReadyOffer::Query {
412                    conn: self.transition(),
413                    query,
414                })
415            }
416            (Some(backend::Event::FunctionCall), FrontendMessage::FunctionCall(message)) => {
417                Ok(ServerReadyOffer::FunctionCall {
418                    conn: self.transition(),
419                    message,
420                })
421            }
422            (Some(backend::Event::Terminate), FrontendMessage::Terminate) => {
423                Ok(ServerReadyOffer::Terminate(self.transition()))
424            }
425            (Some(_), other) => project_extended(self, backend::RuntimeState::Ready, other)
426                .map(ServerReadyOffer::Extended),
427            (None, other) => Err(Box::new((self, other))),
428        }
429    }
430
431    /// Accepts inspected query text which cannot retain client session state.
432    pub fn accept_stateless_query(self, query: Bytes) -> (Conn<S, ServerSimpleQuery, C>, Bytes) {
433        (self.transition(), query)
434    }
435
436    /// Accepts an allow-listed function call known not to retain session state.
437    pub fn accept_stateless_function_call(
438        self,
439        message: FunctionCall,
440    ) -> (Conn<S, ServerFunctionCall, C>, FunctionCall) {
441        (self.transition(), message)
442    }
443}
444
445impl<S, C> Conn<S, ServerFunctionCall, C> {
446    /// Sends the typed function result before the mandatory ready message.
447    ///
448    /// # Errors
449    ///
450    /// Returns an error if the result is too large for a wire frame.
451    pub fn respond(self, value: Bytes) -> io::Result<(Conn<S, ServerFunctionCallDone, C>, Frame)> {
452        Ok((
453            self.transition(),
454            BackendMessage::FunctionCallResponse(value).to_frame()?,
455        ))
456    }
457
458    /// Rejects the call before the mandatory ready message.
459    ///
460    /// # Errors
461    ///
462    /// Returns an error if a diagnostic field is invalid.
463    pub fn error(
464        self,
465        response: DiagnosticResponse,
466    ) -> io::Result<(Conn<S, ServerFunctionCallError, C>, Frame)> {
467        Ok((
468            self.transition(),
469            BackendMessage::ErrorResponse(response).to_frame()?,
470        ))
471    }
472}
473
474impl<S, C> Conn<S, ServerFunctionCallDone, C> {
475    /// Sends readiness after a successful function call.
476    ///
477    /// # Errors
478    ///
479    /// Returns an error only if the fixed ready message cannot be encoded.
480    pub fn ready(self, status: TransactionStatus) -> io::Result<(ServerReadyState<S, C>, Frame)> {
481        ready(self, status)
482    }
483}
484
485impl<S, C> Conn<S, ServerFunctionCallError, C> {
486    /// Sends readiness after a failed function call.
487    ///
488    /// # Errors
489    ///
490    /// Returns an error only if the fixed ready message cannot be encoded.
491    pub fn ready(self, status: TransactionStatus) -> io::Result<(ServerReadyState<S, C>, Frame)> {
492        ready(self, status)
493    }
494}
495
496impl<S, C> Conn<S, ServerBuilding, C> {
497    /// Projects the next inspected message in an extended-query pipeline.
498    ///
499    /// # Errors
500    ///
501    /// Returns the unchanged state and message if it is not legal before `Sync`.
502    pub fn offer_frontend(
503        self,
504        message: FrontendMessage,
505    ) -> ExtendedProjection<S, ServerBuilding, C> {
506        project_extended(self, backend::RuntimeState::Building, message)
507    }
508}
509
510impl<S, C> Conn<S, ServerSimpleQuery, C> {
511    /// Sends a non-terminal typed result message after proxy inspection or rewriting.
512    ///
513    /// # Errors
514    ///
515    /// Returns an error if the message cannot be reconstructed, or if it would
516    /// prematurely change the simple-query state.
517    pub fn send(self, message: &BackendMessage) -> io::Result<(Self, Frame)> {
518        if matches!(
519            message,
520            BackendMessage::ErrorResponse(_)
521                | BackendMessage::ReadyForQuery(_)
522                | BackendMessage::CopyInResponse(_)
523                | BackendMessage::CopyOutResponse(_)
524                | BackendMessage::CopyBothResponse(_)
525        ) {
526            return Err(io::Error::new(
527                io::ErrorKind::InvalidInput,
528                "state-changing response requires its typed transition",
529            ));
530        }
531        Ok((self, message.to_frame()?))
532    }
533
534    /// Sends an error response before the mandatory `ReadyForQuery`.
535    ///
536    /// # Errors
537    ///
538    /// Returns an error if a diagnostic field is invalid.
539    pub fn error(
540        self,
541        response: DiagnosticResponse,
542    ) -> io::Result<(Conn<S, ServerSimpleError, C>, Frame)> {
543        Ok((
544            self.transition(),
545            BackendMessage::ErrorResponse(response).to_frame()?,
546        ))
547    }
548
549    /// Starts a simple-query COPY IN sub-session.
550    ///
551    /// # Errors
552    ///
553    /// Returns an error if the format count overflows the protocol field.
554    pub fn copy_in(self, response: CopyResponse) -> CopyInStart<S, C, CopySimple> {
555        Ok((
556            self.transition(),
557            BackendMessage::CopyInResponse(response).to_frame()?,
558        ))
559    }
560
561    /// Starts a simple-query COPY OUT sub-session.
562    ///
563    /// # Errors
564    ///
565    /// Returns an error if the format count overflows the protocol field.
566    pub fn copy_out(self, response: CopyResponse) -> CopyOutStart<S, C, CopySimple> {
567        Ok((
568            self.transition(),
569            BackendMessage::CopyOutResponse(response).to_frame()?,
570        ))
571    }
572
573    /// Starts a simple-query COPY BOTH sub-session.
574    ///
575    /// # Errors
576    ///
577    /// Returns an error if the format count overflows the protocol field.
578    pub fn copy_both(self, response: CopyResponse) -> CopyBothStart<S, C, CopySimple> {
579        Ok((
580            self.transition(),
581            BackendMessage::CopyBothResponse(response).to_frame()?,
582        ))
583    }
584
585    /// Ends a successful simple-query exchange and surfaces transaction status.
586    ///
587    /// # Errors
588    ///
589    /// Returns an error only if the fixed ready message cannot be encoded.
590    pub fn ready(self, status: TransactionStatus) -> io::Result<(ServerReadyState<S, C>, Frame)> {
591        ready(self, status)
592    }
593}
594
595impl<S, C> Conn<S, ServerSimpleError, C> {
596    /// Ends an errored simple-query exchange and surfaces transaction status.
597    ///
598    /// # Errors
599    ///
600    /// Returns an error only if the fixed ready message cannot be encoded.
601    pub fn ready(self, status: TransactionStatus) -> io::Result<(ServerReadyState<S, C>, Frame)> {
602        ready(self, status)
603    }
604}
605
606impl<S, C> Conn<S, ServerParse, C> {
607    /// Confirms a successful `Parse` and returns to the building loop.
608    ///
609    /// # Errors
610    ///
611    /// Returns an error only if the fixed response cannot be encoded.
612    pub fn complete(self) -> io::Result<(Conn<S, ServerBuilding, C>, Frame)> {
613        Ok((self.transition(), BackendMessage::ParseComplete.to_frame()?))
614    }
615
616    /// Rejects `Parse` and begins discarding the pipeline until `Sync`.
617    ///
618    /// # Errors
619    ///
620    /// Returns an error if a diagnostic field is invalid.
621    pub fn error(
622        self,
623        response: DiagnosticResponse,
624    ) -> io::Result<(Conn<S, ServerExtendedError, C>, Frame)> {
625        extended_error(self, response)
626    }
627}
628
629impl<S, C> Conn<S, ServerBind, C> {
630    /// Confirms a successful `Bind` and returns to the building loop.
631    ///
632    /// # Errors
633    ///
634    /// Returns an error only if the fixed response cannot be encoded.
635    pub fn complete(self) -> io::Result<(Conn<S, ServerBuilding, C>, Frame)> {
636        Ok((self.transition(), BackendMessage::BindComplete.to_frame()?))
637    }
638
639    /// Rejects `Bind` and begins discarding the pipeline until `Sync`.
640    ///
641    /// # Errors
642    ///
643    /// Returns an error if a diagnostic field is invalid.
644    pub fn error(
645        self,
646        response: DiagnosticResponse,
647    ) -> io::Result<(Conn<S, ServerExtendedError, C>, Frame)> {
648        extended_error(self, response)
649    }
650}
651
652impl<S, C> Conn<S, ServerClose, C> {
653    /// Confirms `Close` and returns to the building loop.
654    ///
655    /// # Errors
656    ///
657    /// Returns an error only if the fixed response cannot be encoded.
658    pub fn complete(self) -> io::Result<(Conn<S, ServerBuilding, C>, Frame)> {
659        Ok((self.transition(), BackendMessage::CloseComplete.to_frame()?))
660    }
661
662    /// Rejects `Close` and begins discarding the pipeline until `Sync`.
663    ///
664    /// # Errors
665    ///
666    /// Returns an error if a diagnostic field is invalid.
667    pub fn error(
668        self,
669        response: DiagnosticResponse,
670    ) -> io::Result<(Conn<S, ServerExtendedError, C>, Frame)> {
671        extended_error(self, response)
672    }
673}
674
675impl<S, C> Conn<S, ServerDescribe, C> {
676    /// Sends statement parameter OIDs before its row metadata.
677    ///
678    /// # Errors
679    ///
680    /// Returns an error if the OID count overflows the protocol field.
681    pub fn parameter_description(self, oids: Vec<u32>) -> io::Result<(Self, Frame)> {
682        Ok((self, BackendMessage::ParameterDescription(oids).to_frame()?))
683    }
684
685    /// Sends reconstructable row metadata and returns to the building loop.
686    ///
687    /// # Errors
688    ///
689    /// Returns an error if field metadata is invalid.
690    pub fn row_description(
691        self,
692        description: RowDescription,
693    ) -> io::Result<(Conn<S, ServerBuilding, C>, Frame)> {
694        Ok((
695            self.transition(),
696            BackendMessage::RowDescription(description).to_frame()?,
697        ))
698    }
699
700    /// Sends `NoData` and returns to the building loop.
701    ///
702    /// # Errors
703    ///
704    /// Returns an error only if the fixed response cannot be encoded.
705    pub fn no_data(self) -> io::Result<(Conn<S, ServerBuilding, C>, Frame)> {
706        Ok((self.transition(), BackendMessage::NoData.to_frame()?))
707    }
708
709    /// Rejects `Describe` and begins discarding the pipeline until `Sync`.
710    ///
711    /// # Errors
712    ///
713    /// Returns an error if a diagnostic field is invalid.
714    pub fn error(
715        self,
716        response: DiagnosticResponse,
717    ) -> io::Result<(Conn<S, ServerExtendedError, C>, Frame)> {
718        extended_error(self, response)
719    }
720}
721
722impl<S, C> Conn<S, ServerExecute, C> {
723    /// Sends a non-terminal result message for `Execute`.
724    ///
725    /// # Errors
726    ///
727    /// Returns an error if the message cannot be reconstructed or requires a
728    /// dedicated state transition.
729    pub fn send(self, message: &BackendMessage) -> io::Result<(Self, Frame)> {
730        if matches!(
731            message,
732            BackendMessage::CommandComplete(_)
733                | BackendMessage::PortalSuspended
734                | BackendMessage::ErrorResponse(_)
735                | BackendMessage::ReadyForQuery(_)
736                | BackendMessage::CopyInResponse(_)
737                | BackendMessage::CopyOutResponse(_)
738                | BackendMessage::CopyBothResponse(_)
739        ) {
740            return Err(io::Error::new(
741                io::ErrorKind::InvalidInput,
742                "state-changing response requires its typed transition",
743            ));
744        }
745        Ok((self, message.to_frame()?))
746    }
747
748    /// Completes execution and returns to the building loop.
749    ///
750    /// # Errors
751    ///
752    /// Returns an error if the command tag contains a NUL byte.
753    pub fn command_complete(self, tag: Bytes) -> io::Result<(Conn<S, ServerBuilding, C>, Frame)> {
754        Ok((
755            self.transition(),
756            BackendMessage::CommandComplete(tag).to_frame()?,
757        ))
758    }
759
760    /// Suspends a portal and returns to the building loop.
761    ///
762    /// # Errors
763    ///
764    /// Returns an error only if the fixed response cannot be encoded.
765    pub fn portal_suspended(self) -> io::Result<(Conn<S, ServerBuilding, C>, Frame)> {
766        Ok((
767            self.transition(),
768            BackendMessage::PortalSuspended.to_frame()?,
769        ))
770    }
771
772    /// Rejects `Execute` and begins discarding the pipeline until `Sync`.
773    ///
774    /// # Errors
775    ///
776    /// Returns an error if a diagnostic field is invalid.
777    pub fn error(
778        self,
779        response: DiagnosticResponse,
780    ) -> io::Result<(Conn<S, ServerExtendedError, C>, Frame)> {
781        extended_error(self, response)
782    }
783
784    /// Starts an extended-query COPY IN sub-session.
785    ///
786    /// # Errors
787    ///
788    /// Returns an error if the format count overflows the protocol field.
789    pub fn copy_in(self, response: CopyResponse) -> CopyInStart<S, C, CopyExtended> {
790        Ok((
791            self.transition(),
792            BackendMessage::CopyInResponse(response).to_frame()?,
793        ))
794    }
795
796    /// Starts an extended-query COPY OUT sub-session.
797    ///
798    /// # Errors
799    ///
800    /// Returns an error if the format count overflows the protocol field.
801    pub fn copy_out(self, response: CopyResponse) -> CopyOutStart<S, C, CopyExtended> {
802        Ok((
803            self.transition(),
804            BackendMessage::CopyOutResponse(response).to_frame()?,
805        ))
806    }
807
808    /// Starts an extended-query COPY BOTH sub-session.
809    ///
810    /// # Errors
811    ///
812    /// Returns an error if the format count overflows the protocol field.
813    pub fn copy_both(self, response: CopyResponse) -> CopyBothStart<S, C, CopyExtended> {
814        Ok((
815            self.transition(),
816            BackendMessage::CopyBothResponse(response).to_frame()?,
817        ))
818    }
819}
820
821impl<S, C> Conn<S, ServerCopyIn<CopySimple>, C> {
822    /// Projects one inspected frontend message inside COPY IN.
823    ///
824    /// # Errors
825    ///
826    /// Returns the unchanged state and message for anything other than COPY data,
827    /// completion, or failure.
828    pub fn offer_frontend(self, message: FrontendMessage) -> CopyInProjection<S, C, CopySimple> {
829        project_copy_in(self, backend::RuntimeState::SimpleCopyIn, message)
830    }
831}
832
833impl<S, C> Conn<S, ServerCopyIn<CopyExtended>, C> {
834    /// Projects one inspected frontend message inside extended-query COPY IN.
835    ///
836    /// # Errors
837    ///
838    /// Returns the unchanged state and message for anything other than COPY data,
839    /// completion, or failure.
840    pub fn offer_frontend(self, message: FrontendMessage) -> CopyInProjection<S, C, CopyExtended> {
841        project_copy_in(self, backend::RuntimeState::ExtendedCopyIn, message)
842    }
843}
844
845impl<S, C> Conn<S, ServerCopyInDone<CopySimple>, C> {
846    /// Completes simple-query COPY IN before `ReadyForQuery`.
847    ///
848    /// # Errors
849    ///
850    /// Returns an error if the command tag contains a NUL byte.
851    pub fn command_complete(
852        self,
853        tag: Bytes,
854    ) -> io::Result<(Conn<S, ServerSimpleQuery, C>, Frame)> {
855        Ok((
856            self.transition(),
857            BackendMessage::CommandComplete(tag).to_frame()?,
858        ))
859    }
860}
861
862impl<S, C> Conn<S, ServerCopyInDone<CopyExtended>, C> {
863    /// Completes extended-query COPY IN and returns to the building loop.
864    ///
865    /// # Errors
866    ///
867    /// Returns an error if the command tag contains a NUL byte.
868    pub fn command_complete(self, tag: Bytes) -> io::Result<(Conn<S, ServerBuilding, C>, Frame)> {
869        Ok((
870            self.transition(),
871            BackendMessage::CommandComplete(tag).to_frame()?,
872        ))
873    }
874}
875
876impl<S, C> Conn<S, ServerCopyInFailed<CopySimple>, C> {
877    /// Reports a client COPY failure before simple-query readiness.
878    ///
879    /// # Errors
880    ///
881    /// Returns an error if a diagnostic field is invalid.
882    pub fn error(
883        self,
884        response: DiagnosticResponse,
885    ) -> io::Result<(Conn<S, ServerSimpleError, C>, Frame)> {
886        Ok((
887            self.transition(),
888            BackendMessage::ErrorResponse(response).to_frame()?,
889        ))
890    }
891}
892
893impl<S, C> Conn<S, ServerCopyInFailed<CopyExtended>, C> {
894    /// Reports a client COPY failure and discards the pipeline until `Sync`.
895    ///
896    /// # Errors
897    ///
898    /// Returns an error if a diagnostic field is invalid.
899    pub fn error(
900        self,
901        response: DiagnosticResponse,
902    ) -> io::Result<(Conn<S, ServerExtendedError, C>, Frame)> {
903        extended_error(self, response)
904    }
905}
906
907impl<S, C, Resume> Conn<S, ServerCopyOut<Resume>, C> {
908    /// Sends one COPY OUT data chunk and remains in the nested session.
909    ///
910    /// # Errors
911    ///
912    /// Returns an error only if the data frame cannot be encoded.
913    pub fn data(self, data: Bytes) -> io::Result<(Self, Frame)> {
914        Ok((self, BackendMessage::CopyData(data).to_frame()?))
915    }
916
917    /// Sends a structured WAL or keepalive payload.
918    ///
919    /// # Errors
920    ///
921    /// Returns an error only if the data frame cannot be encoded.
922    pub fn replication(self, message: &BackendReplication) -> io::Result<(Self, Frame)> {
923        self.data(message.encode())
924    }
925
926    /// Ends the COPY data stream before its command completion.
927    ///
928    /// # Errors
929    ///
930    /// Returns an error only if the fixed response cannot be encoded.
931    pub fn done(self) -> CopyOutCompletion<S, C, Resume> {
932        Ok((self.transition(), BackendMessage::CopyDone.to_frame()?))
933    }
934}
935
936impl<S, C> Conn<S, ServerCopyOutDone<CopySimple>, C> {
937    /// Completes simple-query COPY OUT before `ReadyForQuery`.
938    ///
939    /// # Errors
940    ///
941    /// Returns an error if the command tag contains a NUL byte.
942    pub fn command_complete(
943        self,
944        tag: Bytes,
945    ) -> io::Result<(Conn<S, ServerSimpleQuery, C>, Frame)> {
946        Ok((
947            self.transition(),
948            BackendMessage::CommandComplete(tag).to_frame()?,
949        ))
950    }
951}
952
953impl<S, C> Conn<S, ServerCopyOutDone<CopyExtended>, C> {
954    /// Completes extended-query COPY OUT and returns to the building loop.
955    ///
956    /// # Errors
957    ///
958    /// Returns an error if the command tag contains a NUL byte.
959    pub fn command_complete(self, tag: Bytes) -> io::Result<(Conn<S, ServerBuilding, C>, Frame)> {
960        Ok((
961            self.transition(),
962            BackendMessage::CommandComplete(tag).to_frame()?,
963        ))
964    }
965}
966
967impl<S, C, Resume: CopyResume> Conn<S, ServerCopyBoth<Resume, BothOpen>, C> {
968    /// Projects client data, half-close, or failure while both directions are open.
969    ///
970    /// # Errors
971    ///
972    /// Returns the unchanged state and message if it is not COPY traffic.
973    pub fn offer_frontend(self, message: FrontendMessage) -> CopyBothOpenProjection<S, C, Resume> {
974        match (
975            backend::project_external(Resume::BOTH_OPEN_STATE, &message),
976            message,
977        ) {
978            (Some(backend::Event::ReceiveData), FrontendMessage::CopyData(data)) => {
979                Ok(ServerCopyBothOpenOffer::Data { conn: self, data })
980            }
981            (Some(backend::Event::ReceiveDone), FrontendMessage::CopyDone) => {
982                Ok(ServerCopyBothOpenOffer::Done(self.transition()))
983            }
984            (Some(backend::Event::Fail), FrontendMessage::CopyFail(message)) => {
985                Ok(ServerCopyBothOpenOffer::Fail {
986                    conn: self.transition(),
987                    message,
988                })
989            }
990            (_, other) => Err(Box::new((self, other))),
991        }
992    }
993
994    /// Sends backend COPY data while its direction remains open.
995    ///
996    /// # Errors
997    ///
998    /// Returns an error only if the data frame cannot be encoded.
999    pub fn data(self, data: Bytes) -> io::Result<(Self, Frame)> {
1000        Ok((self, BackendMessage::CopyData(data).to_frame()?))
1001    }
1002
1003    /// Sends a structured WAL or keepalive payload while both halves are open.
1004    ///
1005    /// # Errors
1006    ///
1007    /// Returns an error only if the data frame cannot be encoded.
1008    pub fn replication(self, message: &BackendReplication) -> io::Result<(Self, Frame)> {
1009        self.data(message.encode())
1010    }
1011
1012    /// Half-closes the backend direction while the client direction remains open.
1013    ///
1014    /// # Errors
1015    ///
1016    /// Returns an error only if the fixed completion frame cannot be encoded.
1017    pub fn done(self) -> CopyBothServerHalfClose<S, C, Resume> {
1018        Ok((self.transition(), BackendMessage::CopyDone.to_frame()?))
1019    }
1020}
1021
1022impl<S, C, Resume> Conn<S, ServerCopyBoth<Resume, BothClientDone>, C> {
1023    /// Sends remaining backend data after the client has half-closed.
1024    ///
1025    /// # Errors
1026    ///
1027    /// Returns an error only if the data frame cannot be encoded.
1028    pub fn data(self, data: Bytes) -> io::Result<(Self, Frame)> {
1029        Ok((self, BackendMessage::CopyData(data).to_frame()?))
1030    }
1031
1032    /// Sends a structured WAL or keepalive payload after the client half-close.
1033    ///
1034    /// # Errors
1035    ///
1036    /// Returns an error only if the data frame cannot be encoded.
1037    pub fn replication(self, message: &BackendReplication) -> io::Result<(Self, Frame)> {
1038        self.data(message.encode())
1039    }
1040
1041    /// Half-closes the backend direction, completing both COPY streams.
1042    ///
1043    /// # Errors
1044    ///
1045    /// Returns an error only if the fixed completion frame cannot be encoded.
1046    pub fn done(self) -> CopyBothCompletion<S, C, Resume> {
1047        Ok((self.transition(), BackendMessage::CopyDone.to_frame()?))
1048    }
1049}
1050
1051impl<S, C, Resume: CopyResume> Conn<S, ServerCopyBoth<Resume, BothServerDone>, C> {
1052    /// Projects remaining client traffic after the backend has half-closed.
1053    ///
1054    /// # Errors
1055    ///
1056    /// Returns the unchanged state and message if it is not COPY traffic.
1057    pub fn offer_frontend(
1058        self,
1059        message: FrontendMessage,
1060    ) -> CopyBothServerDoneProjection<S, C, Resume> {
1061        match (
1062            backend::project_external(Resume::BOTH_SERVER_DONE_STATE, &message),
1063            message,
1064        ) {
1065            (Some(backend::Event::ReceiveData), FrontendMessage::CopyData(data)) => {
1066                Ok(ServerCopyBothServerDoneOffer::Data { conn: self, data })
1067            }
1068            (Some(backend::Event::ReceiveDone), FrontendMessage::CopyDone) => {
1069                Ok(ServerCopyBothServerDoneOffer::Done(self.transition()))
1070            }
1071            (Some(backend::Event::Fail), FrontendMessage::CopyFail(message)) => {
1072                Ok(ServerCopyBothServerDoneOffer::Fail {
1073                    conn: self.transition(),
1074                    message,
1075                })
1076            }
1077            (_, other) => Err(Box::new((self, other))),
1078        }
1079    }
1080}
1081
1082impl<S, C, Resume> ServerCopyBothOpenOffer<S, C, Resume> {
1083    /// Decodes client COPY data as a structured standby message.
1084    ///
1085    /// # Errors
1086    ///
1087    /// Returns the live connection with a decoding error for malformed known payloads.
1088    pub fn decode_replication(self) -> ServerReplicationOpenProjection<S, C, Resume> {
1089        match self {
1090            Self::Data { conn, data } => match FrontendReplication::decode(data) {
1091                Ok(message) => Ok(ServerReplicationOpenOffer::Message { conn, message }),
1092                Err(error) => Err((conn, error)),
1093            },
1094            Self::Done(conn) => Ok(ServerReplicationOpenOffer::Done(conn)),
1095            Self::Fail { conn, message } => Ok(ServerReplicationOpenOffer::Fail { conn, message }),
1096        }
1097    }
1098}
1099
1100impl<S, C, Resume> ServerCopyBothServerDoneOffer<S, C, Resume> {
1101    /// Decodes remaining client COPY data as a structured standby message.
1102    ///
1103    /// # Errors
1104    ///
1105    /// Returns the live connection with a decoding error for malformed known payloads.
1106    pub fn decode_replication(self) -> ServerReplicationServerDoneProjection<S, C, Resume> {
1107        match self {
1108            Self::Data { conn, data } => match FrontendReplication::decode(data) {
1109                Ok(message) => Ok(ServerReplicationServerDoneOffer::Message { conn, message }),
1110                Err(error) => Err((conn, error)),
1111            },
1112            Self::Done(conn) => Ok(ServerReplicationServerDoneOffer::Done(conn)),
1113            Self::Fail { conn, message } => {
1114                Ok(ServerReplicationServerDoneOffer::Fail { conn, message })
1115            }
1116        }
1117    }
1118}
1119
1120impl<S, C> Conn<S, ServerCopyBoth<CopySimple, BothDone>, C> {
1121    /// Completes simple-query COPY BOTH before `ReadyForQuery`.
1122    ///
1123    /// # Errors
1124    ///
1125    /// Returns an error if the command tag contains a NUL byte.
1126    pub fn command_complete(
1127        self,
1128        tag: Bytes,
1129    ) -> io::Result<(Conn<S, ServerSimpleQuery, C>, Frame)> {
1130        Ok((
1131            self.transition(),
1132            BackendMessage::CommandComplete(tag).to_frame()?,
1133        ))
1134    }
1135}
1136
1137impl<S, C> Conn<S, ServerCopyBoth<CopyExtended, BothDone>, C> {
1138    /// Completes extended-query COPY BOTH and returns to the building loop.
1139    ///
1140    /// # Errors
1141    ///
1142    /// Returns an error if the command tag contains a NUL byte.
1143    pub fn command_complete(self, tag: Bytes) -> io::Result<(Conn<S, ServerBuilding, C>, Frame)> {
1144        Ok((
1145            self.transition(),
1146            BackendMessage::CommandComplete(tag).to_frame()?,
1147        ))
1148    }
1149}
1150
1151impl<S, C> Conn<S, ServerCopyBothFailed<CopySimple>, C> {
1152    /// Reports a client COPY failure before simple-query readiness.
1153    ///
1154    /// # Errors
1155    ///
1156    /// Returns an error if a diagnostic field is invalid.
1157    pub fn error(
1158        self,
1159        response: DiagnosticResponse,
1160    ) -> io::Result<(Conn<S, ServerSimpleError, C>, Frame)> {
1161        Ok((
1162            self.transition(),
1163            BackendMessage::ErrorResponse(response).to_frame()?,
1164        ))
1165    }
1166}
1167
1168impl<S, C> Conn<S, ServerCopyBothFailed<CopyExtended>, C> {
1169    /// Reports a client COPY failure and discards the pipeline until `Sync`.
1170    ///
1171    /// # Errors
1172    ///
1173    /// Returns an error if a diagnostic field is invalid.
1174    pub fn error(
1175        self,
1176        response: DiagnosticResponse,
1177    ) -> io::Result<(Conn<S, ServerExtendedError, C>, Frame)> {
1178        extended_error(self, response)
1179    }
1180}
1181
1182impl<S, C> Conn<S, ServerExtendedError, C> {
1183    /// Discards one pipelined message; only `Sync` exits error recovery.
1184    #[must_use]
1185    pub fn discard(self, message: &FrontendMessage) -> ServerDiscard<S, C> {
1186        match backend::project_external(backend::RuntimeState::ExtendedError, message) {
1187            Some(backend::Event::Sync) => ServerDiscard::Sync(self.transition()),
1188            Some(backend::Event::Discard) | None => ServerDiscard::Continue(self),
1189            Some(_) => unreachable!("extended-error grammar has only discard and sync events"),
1190        }
1191    }
1192}
1193
1194fn project_copy_in<S, C, Resume>(
1195    conn: Conn<S, ServerCopyIn<Resume>, C>,
1196    state: backend::RuntimeState,
1197    message: FrontendMessage,
1198) -> CopyInProjection<S, C, Resume> {
1199    match (backend::project_external(state, &message), message) {
1200        (Some(backend::Event::Data), FrontendMessage::CopyData(data)) => {
1201            Ok(ServerCopyInOffer::Data { conn, data })
1202        }
1203        (Some(backend::Event::Done), FrontendMessage::CopyDone) => {
1204            Ok(ServerCopyInOffer::Done(conn.transition()))
1205        }
1206        (Some(backend::Event::Fail), FrontendMessage::CopyFail(message)) => {
1207            Ok(ServerCopyInOffer::Fail {
1208                conn: conn.transition(),
1209                message,
1210            })
1211        }
1212        (_, other) => Err(Box::new((conn, other))),
1213    }
1214}
1215
1216impl<S, C> Conn<S, ServerSync, C> {
1217    /// Answers `Sync` with `ReadyForQuery` and surfaces transaction status.
1218    ///
1219    /// # Errors
1220    ///
1221    /// Returns an error only if the fixed ready message cannot be encoded.
1222    pub fn ready(self, status: TransactionStatus) -> io::Result<(ServerReadyState<S, C>, Frame)> {
1223        ready(self, status)
1224    }
1225}
1226
1227fn project_extended<S, Phase, C>(
1228    conn: Conn<S, Phase, C>,
1229    state: backend::RuntimeState,
1230    message: FrontendMessage,
1231) -> ExtendedProjection<S, Phase, C> {
1232    Ok(
1233        match (backend::project_external(state, &message), message) {
1234            (Some(backend::Event::Parse), FrontendMessage::Parse(message)) => {
1235                ServerExtendedOffer::Parse {
1236                    conn: conn.transition(),
1237                    message,
1238                }
1239            }
1240            (Some(backend::Event::Bind), FrontendMessage::Bind(message)) => {
1241                ServerExtendedOffer::Bind {
1242                    conn: conn.transition(),
1243                    message,
1244                }
1245            }
1246            (Some(backend::Event::Describe), FrontendMessage::Describe(message)) => {
1247                ServerExtendedOffer::Describe {
1248                    conn: conn.transition(),
1249                    message,
1250                }
1251            }
1252            (Some(backend::Event::Execute), FrontendMessage::Execute(message)) => {
1253                ServerExtendedOffer::Execute {
1254                    conn: conn.transition(),
1255                    message,
1256                }
1257            }
1258            (Some(backend::Event::Close), FrontendMessage::Close(message)) => {
1259                ServerExtendedOffer::Close {
1260                    conn: conn.transition(),
1261                    message,
1262                }
1263            }
1264            (Some(backend::Event::Flush), FrontendMessage::Flush) => {
1265                ServerExtendedOffer::Flush(conn.transition())
1266            }
1267            (Some(backend::Event::Sync), FrontendMessage::Sync) => {
1268                ServerExtendedOffer::Sync(conn.transition())
1269            }
1270            (_, other) => return Err(Box::new((conn, other))),
1271        },
1272    )
1273}
1274
1275fn extended_error<S, Phase, C>(
1276    conn: Conn<S, Phase, C>,
1277    response: DiagnosticResponse,
1278) -> io::Result<(Conn<S, ServerExtendedError, C>, Frame)> {
1279    Ok((
1280        conn.transition(),
1281        BackendMessage::ErrorResponse(response).to_frame()?,
1282    ))
1283}
1284
1285fn ready<S, Phase, C>(
1286    conn: Conn<S, Phase, C>,
1287    status: TransactionStatus,
1288) -> io::Result<(ServerReadyState<S, C>, Frame)> {
1289    let frame = BackendMessage::ReadyForQuery(status).to_frame()?;
1290    let state = if status == TransactionStatus::Idle {
1291        ServerReadyState::Ready(conn.transition())
1292    } else {
1293        ServerReadyState::Dirty {
1294            conn: conn.transition(),
1295            status,
1296        }
1297    };
1298    Ok((state, frame))
1299}
1300
1301#[cfg(test)]
1302mod tests {
1303    use super::*;
1304    use crate::{
1305        Pristine,
1306        codec::{DataRow, DiagnosticField},
1307    };
1308
1309    #[test]
1310    fn simple_query_allows_rewriting_before_ready() {
1311        fn require_dirty<S>(conn: Conn<S, Ready, Dirty>) {
1312            conn.into_transport();
1313        }
1314
1315        let ready: Conn<(), Ready> = Conn::new(()).transition();
1316        let ServerReadyOffer::Query { conn, query } = ready
1317            .offer_frontend(FrontendMessage::Query(Bytes::from_static(b"select 1")))
1318            .unwrap()
1319        else {
1320            panic!("query projected to the wrong branch")
1321        };
1322        assert_eq!(query, Bytes::from_static(b"select 1"));
1323
1324        let rewritten = BackendMessage::DataRow(DataRow {
1325            columns: vec![Some(Bytes::from_static(b"2"))],
1326        });
1327        let (conn, frame) = conn.send(&rewritten).unwrap();
1328        assert_eq!(frame.tag, b'D');
1329        let (state, frame) = conn.ready(TransactionStatus::Idle).unwrap();
1330        assert_eq!(frame.body, Bytes::from_static(b"I"));
1331        let ServerReadyState::Ready(ready) = state else {
1332            panic!("idle response unexpectedly changed the transaction state")
1333        };
1334        require_dirty(ready);
1335
1336        let ready: Conn<(), Ready> = Conn::new(()).transition();
1337        let (query, inspected) = ready.accept_stateless_query(Bytes::from_static(b"select 1"));
1338        assert_eq!(inspected, Bytes::from_static(b"select 1"));
1339        let (state, _) = query.ready(TransactionStatus::Idle).unwrap();
1340        let ServerReadyState::Ready(pristine) = state else {
1341            panic!("stateless query did not return to ready")
1342        };
1343        pristine.release();
1344    }
1345
1346    #[test]
1347    fn function_call_is_inspectable_and_replaceable() {
1348        let ready: Conn<(), Ready> = Conn::new(()).transition();
1349        let call = FunctionCall {
1350            function_oid: 42,
1351            argument_formats: vec![1],
1352            arguments: vec![Some(Bytes::from_static(b"original"))],
1353            result_format: 1,
1354        };
1355        let ServerReadyOffer::FunctionCall { conn, message } = ready
1356            .offer_frontend(FrontendMessage::FunctionCall(call.clone()))
1357            .unwrap()
1358        else {
1359            panic!("function call projected to the wrong branch")
1360        };
1361        assert_eq!(message, call);
1362
1363        let (done, frame) = conn.respond(Bytes::from_static(b"replacement")).unwrap();
1364        assert_eq!(frame.tag, b'V');
1365        let (state, _) = done.ready(TransactionStatus::Idle).unwrap();
1366        let ServerReadyState::Ready(ready) = state else {
1367            panic!("idle function call was marked dirty")
1368        };
1369        ready.into_transport();
1370    }
1371
1372    #[test]
1373    fn transaction_status_taints_the_server_connection() {
1374        let query: Conn<(), ServerSimpleQuery, Pristine> = Conn::new(()).transition();
1375        let (state, _) = query.ready(TransactionStatus::InTransaction).unwrap();
1376        let ServerReadyState::Dirty { conn, status } = state else {
1377            panic!("transactional response was marked clean")
1378        };
1379        assert_eq!(status, TransactionStatus::InTransaction);
1380        conn.into_transport();
1381    }
1382
1383    #[test]
1384    fn extended_pipeline_rewrites_parse_and_exits_only_through_sync() {
1385        fn require_dirty<S>(conn: Conn<S, Ready, Dirty>) {
1386            conn.into_transport();
1387        }
1388
1389        let ready: Conn<(), Ready> = Conn::new(()).transition();
1390        let parse = Parse {
1391            statement: Bytes::from_static(b"statement"),
1392            query: Bytes::from_static(b"select $1"),
1393            parameter_types: vec![23],
1394        };
1395        let ServerReadyOffer::Extended(ServerExtendedOffer::Parse { conn, message }) = ready
1396            .offer_frontend(FrontendMessage::Parse(parse.clone()))
1397            .unwrap()
1398        else {
1399            panic!("parse projected to the wrong branch")
1400        };
1401        assert_eq!(message, parse);
1402        let (building, complete) = conn.complete().unwrap();
1403        assert_eq!(complete.tag, b'1');
1404
1405        let ServerExtendedOffer::Bind { conn, message } = building
1406            .offer_frontend(FrontendMessage::Bind(Bind {
1407                portal: Bytes::new(),
1408                statement: Bytes::from_static(b"statement"),
1409                parameter_formats: vec![],
1410                parameters: vec![Some(Bytes::from_static(b"42"))],
1411                result_formats: vec![],
1412            }))
1413            .unwrap()
1414        else {
1415            panic!("bind projected to the wrong branch")
1416        };
1417        assert_eq!(message.parameters[0], Some(Bytes::from_static(b"42")));
1418        let (building, _) = conn.complete().unwrap();
1419        let ServerExtendedOffer::Sync(sync) =
1420            building.offer_frontend(FrontendMessage::Sync).unwrap()
1421        else {
1422            panic!("sync projected to the wrong branch")
1423        };
1424        let (state, _) = sync.ready(TransactionStatus::Idle).unwrap();
1425        let ServerReadyState::Ready(ready) = state else {
1426            panic!("idle sync unexpectedly changed transaction state")
1427        };
1428        require_dirty(ready);
1429    }
1430
1431    #[test]
1432    fn extended_error_discards_everything_before_sync() {
1433        let parse: Conn<(), ServerParse> = Conn::new(()).transition();
1434        let (error, frame) = parse
1435            .error(DiagnosticResponse {
1436                fields: vec![DiagnosticField {
1437                    code: b'C',
1438                    value: Bytes::from_static(b"42601"),
1439                }],
1440            })
1441            .unwrap();
1442        assert_eq!(frame.tag, b'E');
1443        let ServerDiscard::Continue(error) = error.discard(&FrontendMessage::Flush) else {
1444            panic!("flush escaped error recovery")
1445        };
1446        let ServerDiscard::Sync(sync) = error.discard(&FrontendMessage::Sync) else {
1447            panic!("sync did not exit error recovery")
1448        };
1449        let (state, _) = sync.ready(TransactionStatus::Idle).unwrap();
1450        let ServerReadyState::Ready(ready) = state else {
1451            panic!("idle sync was marked dirty")
1452        };
1453        ready.into_transport();
1454    }
1455
1456    #[test]
1457    fn extended_copy_in_is_a_nested_client_choice() {
1458        let execute: Conn<(), ServerExecute> = Conn::new(()).transition();
1459        let (copy, response) = execute
1460            .copy_in(CopyResponse {
1461                overall_format: 0,
1462                column_formats: vec![],
1463            })
1464            .unwrap();
1465        assert_eq!(response.tag, b'G');
1466        let ServerCopyInOffer::Data { conn: copy, data } = copy
1467            .offer_frontend(FrontendMessage::CopyData(Bytes::from_static(b"one\n")))
1468            .unwrap()
1469        else {
1470            panic!("COPY data projected to the wrong branch")
1471        };
1472        assert_eq!(data, Bytes::from_static(b"one\n"));
1473        let ServerCopyInOffer::Done(done) = copy.offer_frontend(FrontendMessage::CopyDone).unwrap()
1474        else {
1475            panic!("COPY completion projected to the wrong branch")
1476        };
1477        let (building, complete) = done
1478            .command_complete(Bytes::from_static(b"COPY 1"))
1479            .unwrap();
1480        assert_eq!(complete.tag, b'C');
1481        let ServerExtendedOffer::Sync(sync) =
1482            building.offer_frontend(FrontendMessage::Sync).unwrap()
1483        else {
1484            panic!("sync projected to the wrong branch")
1485        };
1486        let (state, _) = sync.ready(TransactionStatus::Idle).unwrap();
1487        let ServerReadyState::Ready(ready) = state else {
1488            panic!("idle sync was marked dirty")
1489        };
1490        ready.into_transport();
1491    }
1492
1493    #[test]
1494    fn simple_copy_out_requires_done_before_command_completion() {
1495        let query: Conn<(), ServerSimpleQuery> = Conn::new(()).transition();
1496        let (copy, response) = query
1497            .copy_out(CopyResponse {
1498                overall_format: 0,
1499                column_formats: vec![],
1500            })
1501            .unwrap();
1502        assert_eq!(response.tag, b'H');
1503        let (copy, data) = copy.data(Bytes::from_static(b"one\n")).unwrap();
1504        assert_eq!(data.tag, b'd');
1505        let (done, done_frame) = copy.done().unwrap();
1506        assert_eq!(done_frame.tag, b'c');
1507        let (query, complete) = done
1508            .command_complete(Bytes::from_static(b"COPY 1"))
1509            .unwrap();
1510        assert_eq!(complete.tag, b'C');
1511        let (state, _) = query.ready(TransactionStatus::Idle).unwrap();
1512        let ServerReadyState::Ready(ready) = state else {
1513            panic!("idle COPY was marked dirty")
1514        };
1515        ready.into_transport();
1516    }
1517
1518    #[test]
1519    fn copy_both_tracks_half_closes_independently() {
1520        use crate::grammar::backend::{Event, RuntimeFsm, RuntimeState};
1521
1522        let mut generated = RuntimeFsm::new();
1523        generated.step(Event::Execute).unwrap();
1524        let execute: Conn<(), ServerExecute> = Conn::new(()).transition();
1525        let (both, response) = execute
1526            .copy_both(CopyResponse {
1527                overall_format: 0,
1528                column_formats: vec![],
1529            })
1530            .unwrap();
1531        generated.step(Event::CopyBoth).unwrap();
1532        assert_eq!(response.tag, b'W');
1533        let ServerCopyBothOpenOffer::Done(client_done) =
1534            both.offer_frontend(FrontendMessage::CopyDone).unwrap()
1535        else {
1536            panic!("client half-close projected to the wrong branch")
1537        };
1538        generated.step(Event::ReceiveDone).unwrap();
1539        let (client_done, data) = client_done
1540            .data(Bytes::from_static(b"remaining backend data"))
1541            .unwrap();
1542        generated.step(Event::SendData).unwrap();
1543        assert_eq!(data.tag, b'd');
1544        let (done, backend_done) = client_done.done().unwrap();
1545        generated.step(Event::SendDone).unwrap();
1546        assert_eq!(backend_done.tag, b'c');
1547        let (building, _) = done
1548            .command_complete(Bytes::from_static(b"COPY 0"))
1549            .unwrap();
1550        generated.step(Event::CommandComplete).unwrap();
1551        let ServerExtendedOffer::Sync(sync) =
1552            building.offer_frontend(FrontendMessage::Sync).unwrap()
1553        else {
1554            panic!("sync projected to the wrong branch")
1555        };
1556        generated.step(Event::Sync).unwrap();
1557        let (state, _) = sync.ready(TransactionStatus::Idle).unwrap();
1558        let ServerReadyState::Ready(ready) = state else {
1559            panic!("idle sync was marked dirty")
1560        };
1561        generated.step(Event::Ready).unwrap();
1562        assert_eq!(generated.state(), RuntimeState::Ready);
1563        ready.into_transport();
1564    }
1565
1566    #[test]
1567    fn copy_both_inspects_and_replaces_replication_messages() {
1568        let both: Conn<(), ServerCopyBoth<CopySimple, BothOpen>> = Conn::new(()).transition();
1569        let status = FrontendReplication::StandbyStatus {
1570            written: 10,
1571            flushed: 9,
1572            applied: 8,
1573            client_time: 7,
1574            reply_requested: true,
1575        };
1576        let offer = both
1577            .offer_frontend(FrontendMessage::CopyData(status.encode()))
1578            .unwrap();
1579        let ServerReplicationOpenOffer::Message {
1580            conn: both,
1581            message,
1582        } = offer.decode_replication().unwrap()
1583        else {
1584            panic!("standby status projected to the wrong branch")
1585        };
1586        assert_eq!(message, status);
1587
1588        let replacement = BackendReplication::PrimaryKeepalive {
1589            wal_end: 11,
1590            server_time: 12,
1591            reply_requested: false,
1592        };
1593        let (both, frame) = both.replication(&replacement).unwrap();
1594        assert_eq!(frame.body, replacement.encode());
1595        both.into_transport();
1596
1597        let both: Conn<(), ServerCopyBoth<CopySimple, BothOpen>> = Conn::new(()).transition();
1598        let offer = both
1599            .offer_frontend(FrontendMessage::CopyData(Bytes::from_static(b"rshort")))
1600            .unwrap();
1601        let (both, _) = offer.decode_replication().unwrap_err();
1602        both.into_transport();
1603    }
1604}
pg_proto/server_session.rs

pg_proto/
server_session.rs
