pg_proto/
session.rs

1//! Core query, COPY, and error-draining typestates.
2
3use std::io;
4
5use bytes::{BufMut, Bytes, BytesMut};
6
7use crate::{
8    Conn, Dirty, Pristine,
9    auth::Ready,
10    codec::{
11        BackendMessage, Bind, Close, CopyResponse, Describe, DiagnosticResponse, Execute, Frame,
12        FunctionCall, Parse, TransactionStatus,
13    },
14    demux::SessionItem,
15    grammar::frontend,
16    pre_startup::Terminated,
17    replication::{BackendReplication, FrontendReplication},
18};
19
20#[derive(Debug)]
21/// A simple-query command is awaiting backend results and readiness.
22pub enum SimpleQuery {}
23
24#[derive(Debug)]
25/// A legacy function call is awaiting its result.
26pub enum FunctionCalling {}
27
28#[derive(Debug)]
29/// An extended-query pipeline is being built but has no bound portal yet.
30pub enum Building {}
31
32#[derive(Debug)]
33/// An extended-query pipeline contains a bound, executable portal.
34pub enum BoundBuilding {}
35
36#[derive(Debug)]
37/// Command results are complete and `ReadyForQuery` is required.
38pub enum AwaitingReady {}
39
40#[derive(Debug)]
41/// The client may stream COPY data to the backend.
42pub enum CopyIn {}
43
44#[derive(Debug)]
45/// The backend may stream COPY data to the client.
46pub enum CopyOut {}
47
48#[derive(Debug)]
49/// Both peers may stream COPY data.
50pub enum CopyBoth {}
51
52#[derive(Debug)]
53/// The client half of a COPY BOTH session is closed.
54pub enum CopyBothClientDone {}
55
56#[derive(Debug)]
57/// The backend half of a COPY BOTH session is closed.
58pub enum CopyBothServerDone {}
59
60#[derive(Debug)]
61/// An errored command is being drained until `ReadyForQuery`.
62pub enum Draining {}
63
64#[derive(Debug)]
65/// A pool-reset query is executing on a dirty connection.
66pub enum Resetting {}
67
68#[derive(Debug)]
69/// Reset command completion was observed and readiness is required.
70pub enum ResetComplete {}
71
72/// Structured backend error used by session transitions.
73pub type ErrorResponse = DiagnosticResponse;
74
75/// A successful transition or a backend error that enters [`Draining`].
76pub type Fallible<T, S, C = Pristine> = Result<T, (Conn<S, Draining, C>, ErrorResponse)>;
77
78/// Backend branch available while a simple query is active.
79#[derive(Debug)]
80pub enum SimpleTransition<S, C> {
81    /// A result item was consumed and more results may follow.
82    Continue(Conn<S, SimpleQuery, C>, SessionItem),
83    /// The command entered COPY IN.
84    CopyIn(Conn<S, CopyIn, C>, CopyResponse),
85    /// The command entered COPY OUT.
86    CopyOut(Conn<S, CopyOut, C>, CopyResponse),
87    /// The command entered COPY BOTH.
88    CopyBoth(Conn<S, CopyBoth, C>, CopyResponse),
89    /// Readiness completed the simple-query cycle.
90    Ready(ReadyState<S, C>),
91    /// An error entered drain-until-ready recovery.
92    Error(Conn<S, Draining, C>, ErrorResponse),
93}
94
95/// Backend branch available after a result completes.
96#[derive(Debug)]
97pub enum AwaitingReadyTransition<S, C> {
98    /// A non-terminal item was consumed; continue waiting.
99    Continue(Conn<S, AwaitingReady, C>, SessionItem),
100    /// Readiness completed the command cycle.
101    Ready(ReadyState<S, C>),
102    /// An error entered drain-until-ready recovery.
103    Error(Conn<S, Draining, C>, ErrorResponse),
104}
105
106/// Backend branch available for the legacy function-call protocol.
107#[derive(Debug)]
108pub enum FunctionCallTransition<S, C> {
109    /// The backend returned the function value.
110    Response(Conn<S, AwaitingReady, C>, Bytes),
111    /// The backend rejected the call.
112    Error(Conn<S, Draining, C>, ErrorResponse),
113}
114
115/// Backend branch while discarding messages after an error.
116#[derive(Debug)]
117pub enum DrainingTransition<S, C> {
118    /// A non-terminal item was discarded; continue draining.
119    Continue(Conn<S, Draining, C>, SessionItem),
120    /// Readiness completed recovery.
121    Ready(ReadyState<S, C>),
122}
123
124/// Backend branch during COPY OUT.
125#[derive(Debug)]
126pub enum CopyOutTransition<S, C> {
127    /// One chunk of copy data was received.
128    Data(Conn<S, CopyOut, C>, Bytes),
129    /// The backend closed the copy stream.
130    Done(Conn<S, AwaitingReady, C>),
131    /// COPY failed and entered drain-until-ready recovery.
132    Error(Conn<S, Draining, C>, ErrorResponse),
133}
134
135/// Asynchronous backend branch available while sending COPY IN data.
136#[derive(Debug)]
137pub enum CopyInTransition<S, C> {
138    /// COPY failed and entered drain-until-ready recovery.
139    Error(Conn<S, Draining, C>, ErrorResponse),
140}
141
142/// Backend branch while both COPY directions remain open.
143#[derive(Debug)]
144pub enum CopyBothReceive<S, C> {
145    /// One opaque backend data chunk was received.
146    Data(Conn<S, CopyBoth, C>, Bytes),
147    /// The backend closed its half of the stream.
148    Done(Conn<S, CopyBothServerDone, C>),
149    /// COPY failed and entered drain-until-ready recovery.
150    Error(Conn<S, Draining, C>, ErrorResponse),
151}
152
153/// Backend branch after the client closes its COPY BOTH half.
154#[derive(Debug)]
155pub enum CopyBothClientDoneReceive<S, C> {
156    /// One final opaque backend data chunk was received.
157    Data(Conn<S, CopyBothClientDone, C>, Bytes),
158    /// The backend closed its half of the stream.
159    Done(Conn<S, AwaitingReady, C>),
160    /// COPY failed and entered drain-until-ready recovery.
161    Error(Conn<S, Draining, C>, ErrorResponse),
162}
163
164/// Typed walsender branch while both replication directions remain open.
165#[derive(Debug)]
166pub enum ReplicationReceive<S, C> {
167    /// One decoded backend replication message was received.
168    Message(Conn<S, CopyBoth, C>, BackendReplication),
169    /// The backend closed its half of the stream.
170    Done(Conn<S, CopyBothServerDone, C>),
171    /// Replication failed and entered drain-until-ready recovery.
172    Error(Conn<S, Draining, C>, ErrorResponse),
173}
174
175/// Typed walsender branch after the standby closes its sending half.
176#[derive(Debug)]
177pub enum ReplicationClientDoneReceive<S, C> {
178    /// One decoded backend replication message was received.
179    Message(Conn<S, CopyBothClientDone, C>, BackendReplication),
180    /// The backend closed its half of the stream.
181    Done(Conn<S, AwaitingReady, C>),
182    /// Replication failed and entered drain-until-ready recovery.
183    Error(Conn<S, Draining, C>, ErrorResponse),
184}
185
186/// Replication projection preserving the connection when decoding fails.
187pub type ReplicationProjection<S, C> =
188    Result<ReplicationReceive<S, C>, (Conn<S, CopyBoth, C>, io::Error)>;
189/// Replication projection after client half-close, preserving decode failures.
190pub type ReplicationClientDoneProjection<S, C> =
191    Result<ReplicationClientDoneReceive<S, C>, (Conn<S, CopyBothClientDone, C>, io::Error)>;
192
193/// Readiness classified by pooling cleanliness evidence.
194#[derive(Debug)]
195pub enum ReadyState<S, C> {
196    /// The connection retained its existing cleanliness index.
197    Clean(Conn<S, Ready, C>),
198    /// Transaction or parameter evidence made the connection dirty.
199    Dirty {
200        /// Ready connection carrying the dirty cleanliness marker.
201        conn: Conn<S, Ready, Dirty>,
202        /// Transaction status reported by the backend.
203        status: TransactionStatus,
204        /// Whether reported parameters differ from startup values.
205        parameters_changed: bool,
206    },
207}
208
209/// Backend branch while a reset command is executing.
210#[derive(Debug)]
211pub enum ResettingTransition<S> {
212    /// A non-terminal item was consumed; continue waiting.
213    Continue(Conn<S, Resetting, Dirty>, SessionItem),
214    /// Reset command completion was observed.
215    Complete(Conn<S, ResetComplete, Dirty>),
216    /// Reset failed and entered drain-until-ready recovery.
217    Error(Conn<S, Draining, Dirty>, ErrorResponse),
218}
219
220/// Backend branch after reset command completion.
221#[derive(Debug)]
222pub enum ResetCompleteTransition<S> {
223    /// A non-terminal item was consumed; continue waiting.
224    Continue(Conn<S, ResetComplete, Dirty>, SessionItem),
225    /// Idle readiness with restored parameters proved the connection pristine.
226    Ready(Conn<S, Ready, Pristine>),
227    /// Readiness retained dirty transaction or parameter state.
228    Dirty {
229        /// Ready connection retaining its dirty marker.
230        conn: Conn<S, Ready, Dirty>,
231        /// Transaction status reported by the backend.
232        status: TransactionStatus,
233        /// Whether reported parameters differ from startup values.
234        parameters_changed: bool,
235    },
236    /// Reset failed and entered drain-until-ready recovery.
237    Error(Conn<S, Draining, Dirty>, ErrorResponse),
238}
239
240impl<S, C> Conn<S, Ready, C> {
241    /// Gracefully terminates a ready session without waiting for a backend reply.
242    pub fn push_terminate(self) -> (Conn<S, Terminated, C>, Frame) {
243        (self.transition(), empty_frame(b'X'))
244    }
245
246    /// Buffers a simple query and conservatively taints the pooled session.
247    ///
248    /// Simple-query text can create resources which are not reflected in
249    /// `ParameterStatus`, including listeners, prepared statements, and
250    /// advisory locks. Use [`Self::push_stateless_query`] only after custom SQL
251    /// inspection has established that the command cannot retain session state.
252    ///
253    /// # Errors
254    ///
255    /// Returns an error if the query contains a NUL byte.
256    pub fn push_query(self, query: &[u8]) -> io::Result<(Conn<S, SimpleQuery, Dirty>, Frame)> {
257        Ok((self.transition(), cstr_frame(b'Q', query)?))
258    }
259
260    /// Buffers query text which the caller has proved leaves no session state.
261    ///
262    /// # Errors
263    ///
264    /// Returns an error if the query contains a NUL byte.
265    pub fn push_stateless_query(
266        self,
267        query: &[u8],
268    ) -> io::Result<(Conn<S, SimpleQuery, C>, Frame)> {
269        Ok((self.transition(), cstr_frame(b'Q', query)?))
270    }
271
272    /// Buffers the deprecated function-call protocol message as a typed exchange.
273    ///
274    /// # Errors
275    ///
276    /// Returns an error if a count or argument length exceeds its wire field.
277    pub fn push_function_call(
278        self,
279        message: &FunctionCall,
280    ) -> io::Result<(Conn<S, FunctionCalling, Dirty>, Frame)> {
281        Ok((self.transition(), message.to_frame()?))
282    }
283
284    /// Buffers an allow-listed function call known not to retain session state.
285    ///
286    /// # Errors
287    ///
288    /// Returns an error if a count or argument length exceeds its wire field.
289    pub fn push_stateless_function_call(
290        self,
291        message: &FunctionCall,
292    ) -> io::Result<(Conn<S, FunctionCalling, C>, Frame)> {
293        Ok((self.transition(), message.to_frame()?))
294    }
295
296    /// Begins extended-query construction and consumes the ready connection.
297    ///
298    /// ```compile_fail
299    /// use pg_proto::{Conn, auth::Ready};
300    /// fn use_after_transition<S, C>(conn: Conn<S, Ready, C>) {
301    ///     let _building = conn.begin_extended();
302    ///     let _again = conn.begin_extended();
303    /// }
304    /// ```
305    pub fn begin_extended(self) -> Conn<S, Building, C> {
306        self.transition()
307    }
308}
309
310impl<S, C> Conn<S, FunctionCalling, C> {
311    /// Accepts exactly the function result or an error, after which readiness
312    /// must still be consumed.
313    ///
314    /// # Errors
315    ///
316    /// Returns the unchanged connection and message for an illegal response.
317    pub fn offer(
318        self,
319        message: BackendMessage,
320    ) -> Result<FunctionCallTransition<S, C>, (Self, BackendMessage)> {
321        match (
322            frontend::project_external(frontend::RuntimeState::FunctionCalling, &message),
323            message,
324        ) {
325            (
326                Some(frontend::Event::FunctionResponse),
327                BackendMessage::FunctionCallResponse(value),
328            ) => Ok(FunctionCallTransition::Response(self.transition(), value)),
329            (Some(frontend::Event::Error), BackendMessage::ErrorResponse(error)) => {
330                Ok(FunctionCallTransition::Error(self.transition(), error))
331            }
332            (_, other) => Err((self, other)),
333        }
334    }
335}
336
337impl<S> Conn<S, Ready, Pristine> {
338    /// Releases only a statically pristine, ready connection to a pool.
339    ///
340    /// ```compile_fail
341    /// use pg_proto::{Conn, Dirty, auth::Ready};
342    /// fn cannot_release<S>(conn: Conn<S, Ready, Dirty>) {
343    ///     let _transport = conn.release();
344    /// }
345    /// ```
346    pub fn release(self) -> S {
347        self.into_transport()
348    }
349}
350
351impl<S> Conn<S, Ready, Dirty> {
352    /// Begins the only typed path which can recover pool-safe cleanliness.
353    ///
354    /// `ROLLBACK` makes the sequence legal after either transaction status;
355    /// `DISCARD ALL` then clears session-local resources and settings.
356    ///
357    /// # Errors
358    ///
359    /// Returns an error only if the fixed reset query cannot be encoded.
360    pub fn begin_reset(self) -> io::Result<(Conn<S, Resetting, Dirty>, Frame)> {
361        Ok((
362            self.transition(),
363            cstr_frame(b'Q', b"ROLLBACK; DISCARD ALL")?,
364        ))
365    }
366}
367
368impl<S, C> Conn<S, Building, C> {
369    /// Parse is a self-loop while constructing an extended-query pipeline.
370    ///
371    /// # Errors
372    ///
373    /// Returns an error if the structured message cannot be reconstructed.
374    pub fn push_parse(self, message: &Parse) -> io::Result<(Conn<S, Building, Dirty>, Frame)> {
375        Ok((self.transition(), message.to_frame()?))
376    }
377
378    /// Describe is a self-loop while constructing an extended-query pipeline.
379    ///
380    /// # Errors
381    ///
382    /// Returns an error if the structured message cannot be reconstructed.
383    pub fn push_describe(self, message: &Describe) -> io::Result<(Self, Frame)> {
384        Ok((self, message.to_frame()?))
385    }
386
387    /// Bind introduces an executable portal.
388    ///
389    /// Execute is not available before this transition:
390    ///
391    /// ```compile_fail
392    /// use bytes::Bytes;
393    /// use pg_proto::{Conn, session::Building};
394    /// fn execute_without_bind<S, C>(conn: Conn<S, Building, C>) {
395    ///     let _ = conn.push_execute(Bytes::new());
396    /// }
397    /// ```
398    /// # Errors
399    ///
400    /// Returns an error if the structured message cannot be reconstructed.
401    pub fn push_bind(self, message: &Bind) -> io::Result<(Conn<S, BoundBuilding, Dirty>, Frame)> {
402        Ok((self.transition(), message.to_frame()?))
403    }
404
405    /// Close is legal while building, but does not make a portal executable.
406    /// # Errors
407    ///
408    /// Returns an error if the structured message cannot be reconstructed.
409    pub fn push_close(self, message: &Close) -> io::Result<(Self, Frame)> {
410        Ok((self, message.to_frame()?))
411    }
412
413    /// Requests delivery of buffered extended-query responses without ending the cycle.
414    pub fn push_flush(self) -> (Self, Frame) {
415        (self, empty_frame(b'H'))
416    }
417
418    /// Ends the extended-query pipeline and waits for readiness.
419    pub fn push_sync(self) -> (Conn<S, AwaitingReady, C>, Frame) {
420        (self.transition(), empty_frame(b'S'))
421    }
422}
423
424impl<S, C> Conn<S, BoundBuilding, C> {
425    /// # Errors
426    ///
427    /// Returns an error if the structured message cannot be reconstructed.
428    pub fn push_parse(self, message: &Parse) -> io::Result<(Conn<S, BoundBuilding, Dirty>, Frame)> {
429        Ok((self.transition(), message.to_frame()?))
430    }
431
432    /// # Errors
433    ///
434    /// Returns an error if the structured message cannot be reconstructed.
435    pub fn push_bind(self, message: &Bind) -> io::Result<(Conn<S, BoundBuilding, Dirty>, Frame)> {
436        Ok((self.transition(), message.to_frame()?))
437    }
438
439    /// # Errors
440    ///
441    /// Returns an error if the structured message cannot be reconstructed.
442    pub fn push_describe(self, message: &Describe) -> io::Result<(Self, Frame)> {
443        Ok((self, message.to_frame()?))
444    }
445
446    /// Execute is unavailable until a Bind transition has occurred.
447    ///
448    /// # Errors
449    ///
450    /// Returns an error if the structured message cannot be reconstructed.
451    pub fn push_execute(self, message: &Execute) -> io::Result<(Self, Frame)> {
452        Ok((self, message.to_frame()?))
453    }
454
455    /// # Errors
456    ///
457    /// Returns an error if the structured message cannot be reconstructed.
458    pub fn push_close(self, message: &Close) -> io::Result<(Self, Frame)> {
459        Ok((self, message.to_frame()?))
460    }
461
462    /// Requests delivery of buffered extended-query responses without ending the cycle.
463    pub fn push_flush(self) -> (Self, Frame) {
464        (self, empty_frame(b'H'))
465    }
466
467    /// Ends the extended-query pipeline and waits for readiness.
468    pub fn push_sync(self) -> (Conn<S, AwaitingReady, C>, Frame) {
469        (self.transition(), empty_frame(b'S'))
470    }
471}
472
473impl<S, C> Conn<S, SimpleQuery, C> {
474    /// Advances a simple-query session using an actual projected backend item.
475    ///
476    /// # Errors
477    ///
478    /// Returns the unchanged connection and item if it is illegal in this phase.
479    pub fn offer(self, item: SessionItem) -> Result<SimpleTransition<S, C>, (Self, SessionItem)> {
480        match (
481            project_session_item(frontend::RuntimeState::Simple, &item),
482            item,
483        ) {
484            (
485                Some(frontend::Event::CopyIn),
486                SessionItem::Message(BackendMessage::CopyInResponse(response)),
487            ) => Ok(SimpleTransition::CopyIn(self.transition(), response)),
488            (
489                Some(frontend::Event::CopyOut),
490                SessionItem::Message(BackendMessage::CopyOutResponse(response)),
491            ) => Ok(SimpleTransition::CopyOut(self.transition(), response)),
492            (
493                Some(frontend::Event::CopyBoth),
494                SessionItem::Message(BackendMessage::CopyBothResponse(response)),
495            ) => Ok(SimpleTransition::CopyBoth(self.transition(), response)),
496            (
497                Some(frontend::Event::Ready),
498                SessionItem::ReadyForQuery {
499                    status,
500                    parameters_changed,
501                },
502            ) => Ok(SimpleTransition::Ready(ready_state(
503                self,
504                status,
505                parameters_changed,
506            ))),
507            (
508                Some(frontend::Event::Error),
509                SessionItem::Message(BackendMessage::ErrorResponse(error)),
510            ) => Ok(SimpleTransition::Error(self.transition(), error)),
511            (
512                Some(frontend::Event::Continue),
513                item @ (SessionItem::CommandComplete { .. }
514                | SessionItem::Message(
515                    BackendMessage::RowDescription(_)
516                    | BackendMessage::DataRow(_)
517                    | BackendMessage::EmptyQueryResponse,
518                )),
519            ) => Ok(SimpleTransition::Continue(self, item)),
520            (_, item) => Err((self, item)),
521        }
522    }
523}
524
525impl<S, C> Conn<S, CopyIn, C> {
526    /// Query is unavailable in this nested COPY session.
527    ///
528    /// ```compile_fail
529    /// use pg_proto::{Conn, session::CopyIn};
530    /// fn query_during_copy<S, C>(conn: Conn<S, CopyIn, C>) {
531    ///     let _ = conn.push_query(b"select 1");
532    /// }
533    /// ```
534    pub fn push_copy_data(self, data: Bytes) -> (Self, Frame) {
535        (
536            self,
537            Frame {
538                tag: b'd',
539                body: data,
540            },
541        )
542    }
543
544    /// Closes the client-to-backend copy stream and waits for command completion.
545    pub fn push_copy_done(self) -> (Conn<S, AwaitingReady, C>, Frame) {
546        (self.transition(), empty_frame(b'c'))
547    }
548
549    /// Aborts COPY IN with a frontend error string.
550    ///
551    /// # Errors
552    ///
553    /// Returns an error if the message contains a NUL byte.
554    pub fn push_copy_fail(self, message: &[u8]) -> io::Result<(Conn<S, AwaitingReady, C>, Frame)> {
555        Ok((self.transition(), cstr_frame(b'f', message)?))
556    }
557
558    /// Projects an asynchronous backend failure while COPY IN data is being sent.
559    ///
560    /// This branch is reachable after cancellation or an early server-side COPY
561    /// failure. Non-error messages leave the COPY IN connection unchanged.
562    ///
563    /// # Errors
564    ///
565    /// Returns the live connection and item when it is not an error response.
566    pub fn offer(self, item: SessionItem) -> Result<CopyInTransition<S, C>, (Self, SessionItem)> {
567        match (
568            project_session_item(frontend::RuntimeState::CopyIn, &item),
569            item,
570        ) {
571            (
572                Some(frontend::Event::Error),
573                SessionItem::Message(BackendMessage::ErrorResponse(error)),
574            ) => Ok(CopyInTransition::Error(self.transition(), error)),
575            (_, item) => Err((self, item)),
576        }
577    }
578}
579
580impl<S, C> Conn<S, CopyBothClientDone, C> {
581    /// Continues receiving after the frontend half has closed.
582    ///
583    /// # Errors
584    ///
585    /// Returns the unchanged connection and item for an illegal response.
586    pub fn offer(
587        self,
588        item: SessionItem,
589    ) -> Result<CopyBothClientDoneReceive<S, C>, (Self, SessionItem)> {
590        match (
591            project_session_item(frontend::RuntimeState::CopyBothClientDone, &item),
592            item,
593        ) {
594            (
595                Some(frontend::Event::ReceiveCopyData),
596                SessionItem::Message(BackendMessage::CopyData(data)),
597            ) => Ok(CopyBothClientDoneReceive::Data(self, data)),
598            (
599                Some(frontend::Event::ReceiveCopyDone),
600                SessionItem::Message(BackendMessage::CopyDone),
601            ) => Ok(CopyBothClientDoneReceive::Done(self.transition())),
602            (
603                Some(frontend::Event::Error),
604                SessionItem::Message(BackendMessage::ErrorResponse(error)),
605            ) => Ok(CopyBothClientDoneReceive::Error(self.transition(), error)),
606            (_, item) => Err((self, item)),
607        }
608    }
609}
610
611impl<S, C> CopyBothClientDoneReceive<S, C> {
612    /// Decodes data received after the frontend half-close.
613    ///
614    /// # Errors
615    ///
616    /// Returns the connection with a decoding error for malformed known payloads.
617    pub fn decode_replication(self) -> ReplicationClientDoneProjection<S, C> {
618        match self {
619            Self::Data(conn, data) => match BackendReplication::decode(data) {
620                Ok(message) => Ok(ReplicationClientDoneReceive::Message(conn, message)),
621                Err(error) => Err((conn, error)),
622            },
623            Self::Done(conn) => Ok(ReplicationClientDoneReceive::Done(conn)),
624            Self::Error(conn, error) => Ok(ReplicationClientDoneReceive::Error(conn, error)),
625        }
626    }
627}
628
629impl<S, C> Conn<S, CopyBothServerDone, C> {
630    /// Continues sending after the backend half has closed.
631    pub fn push_copy_data(self, data: Bytes) -> (Self, Frame) {
632        (
633            self,
634            Frame {
635                tag: b'd',
636                body: data,
637            },
638        )
639    }
640
641    /// Sends a structured standby message after the backend half-close.
642    pub fn push_replication(self, message: &FrontendReplication) -> (Self, Frame) {
643        self.push_copy_data(message.encode())
644    }
645
646    /// Closes the remaining frontend half and begins readiness processing.
647    pub fn push_copy_done(self) -> (Conn<S, AwaitingReady, C>, Frame) {
648        (self.transition(), empty_frame(b'c'))
649    }
650}
651
652impl<S, C> Conn<S, CopyOut, C> {
653    /// Advances COPY OUT using backend evidence.
654    ///
655    /// # Errors
656    ///
657    /// Returns the unchanged connection and item when it is illegal in COPY OUT.
658    pub fn offer(self, item: SessionItem) -> Result<CopyOutTransition<S, C>, (Self, SessionItem)> {
659        match (
660            project_session_item(frontend::RuntimeState::CopyOut, &item),
661            item,
662        ) {
663            (
664                Some(frontend::Event::CopyData),
665                SessionItem::Message(BackendMessage::CopyData(data)),
666            ) => Ok(CopyOutTransition::Data(self, data)),
667            (Some(frontend::Event::CopyDone), SessionItem::Message(BackendMessage::CopyDone)) => {
668                Ok(CopyOutTransition::Done(self.transition()))
669            }
670            (
671                Some(frontend::Event::Error),
672                SessionItem::Message(BackendMessage::ErrorResponse(error)),
673            ) => Ok(CopyOutTransition::Error(self.transition(), error)),
674            (_, item) => Err((self, item)),
675        }
676    }
677}
678
679impl<S, C> Conn<S, CopyBoth, C> {
680    /// Sends one opaque data chunk while retaining both COPY directions.
681    pub fn push_copy_data(self, data: Bytes) -> (Self, Frame) {
682        (
683            self,
684            Frame {
685                tag: b'd',
686                body: data,
687            },
688        )
689    }
690
691    /// Sends a structured standby message in the walsender stream.
692    pub fn push_replication(self, message: &FrontendReplication) -> (Self, Frame) {
693        self.push_copy_data(message.encode())
694    }
695
696    /// Closes the client half while leaving the backend half readable.
697    pub fn push_copy_done(self) -> (Conn<S, CopyBothClientDone, C>, Frame) {
698        (self.transition(), empty_frame(b'c'))
699    }
700
701    /// Receives the backend half of a bidirectional COPY session.
702    ///
703    /// # Errors
704    ///
705    /// Returns the unchanged connection and item when it is illegal in COPY BOTH.
706    pub fn offer(self, item: SessionItem) -> Result<CopyBothReceive<S, C>, (Self, SessionItem)> {
707        match (
708            project_session_item(frontend::RuntimeState::CopyBoth, &item),
709            item,
710        ) {
711            (
712                Some(frontend::Event::ReceiveCopyData),
713                SessionItem::Message(BackendMessage::CopyData(data)),
714            ) => Ok(CopyBothReceive::Data(self, data)),
715            (
716                Some(frontend::Event::ReceiveCopyDone),
717                SessionItem::Message(BackendMessage::CopyDone),
718            ) => Ok(CopyBothReceive::Done(self.transition())),
719            (
720                Some(frontend::Event::Error),
721                SessionItem::Message(BackendMessage::ErrorResponse(error)),
722            ) => Ok(CopyBothReceive::Error(self.transition(), error)),
723            (_, item) => Err((self, item)),
724        }
725    }
726}
727
728impl<S, C> CopyBothReceive<S, C> {
729    /// Decodes a COPY BOTH data branch as a walsender message.
730    ///
731    /// # Errors
732    ///
733    /// Returns the connection with a decoding error for malformed known payloads.
734    pub fn decode_replication(self) -> ReplicationProjection<S, C> {
735        match self {
736            Self::Data(conn, data) => match BackendReplication::decode(data) {
737                Ok(message) => Ok(ReplicationReceive::Message(conn, message)),
738                Err(error) => Err((conn, error)),
739            },
740            Self::Done(conn) => Ok(ReplicationReceive::Done(conn)),
741            Self::Error(conn, error) => Ok(ReplicationReceive::Error(conn, error)),
742        }
743    }
744}
745
746impl<S, C> Conn<S, Draining, C> {
747    /// `ReadyForQuery` is the sole exit from error draining.
748    pub fn offer(self, item: SessionItem) -> DrainingTransition<S, C> {
749        match (
750            project_session_item(frontend::RuntimeState::Draining, &item),
751            item,
752        ) {
753            (
754                Some(frontend::Event::Ready),
755                SessionItem::ReadyForQuery {
756                    status,
757                    parameters_changed,
758                },
759            ) => DrainingTransition::Ready(ready_state(self, status, parameters_changed)),
760            (_, item) => DrainingTransition::Continue(self, item),
761        }
762    }
763}
764
765impl<S, C> Conn<S, AwaitingReady, C> {
766    /// Consumes responses after Sync until `ReadyForQuery` proves readiness.
767    pub fn offer(self, item: SessionItem) -> AwaitingReadyTransition<S, C> {
768        match (
769            project_session_item(frontend::RuntimeState::AwaitingReady, &item),
770            item,
771        ) {
772            (
773                Some(frontend::Event::Ready),
774                SessionItem::ReadyForQuery {
775                    status,
776                    parameters_changed,
777                },
778            ) => AwaitingReadyTransition::Ready(ready_state(self, status, parameters_changed)),
779            (
780                Some(frontend::Event::Error),
781                SessionItem::Message(BackendMessage::ErrorResponse(error)),
782            ) => AwaitingReadyTransition::Error(self.transition(), error),
783            (_, item) => AwaitingReadyTransition::Continue(self, item),
784        }
785    }
786}
787
788impl<S> Conn<S, Resetting, Dirty> {
789    /// Waits for evidence that `DISCARD ALL` itself completed.
790    #[must_use]
791    pub fn offer(self, item: SessionItem) -> ResettingTransition<S> {
792        match (
793            project_session_item(frontend::RuntimeState::Resetting, &item),
794            item,
795        ) {
796            (Some(frontend::Event::DiscardComplete), SessionItem::CommandComplete { tag, .. })
797                if tag == b"DISCARD ALL".as_slice() =>
798            {
799                ResettingTransition::Complete(self.transition())
800            }
801            (
802                Some(frontend::Event::Error),
803                SessionItem::Message(BackendMessage::ErrorResponse(error)),
804            ) => ResettingTransition::Error(self.transition(), error),
805            (_, item) => ResettingTransition::Continue(self, item),
806        }
807    }
808}
809
810impl<S> Conn<S, ResetComplete, Dirty> {
811    /// Restores `Pristine` only from idle readiness and the startup parameter baseline.
812    #[must_use]
813    pub fn offer(self, item: SessionItem) -> ResetCompleteTransition<S> {
814        match (
815            project_session_item(frontend::RuntimeState::ResetComplete, &item),
816            item,
817        ) {
818            (
819                Some(frontend::Event::ReadyClean),
820                SessionItem::ReadyForQuery {
821                    status: TransactionStatus::Idle,
822                    parameters_changed: false,
823                },
824            ) => ResetCompleteTransition::Ready(self.transition()),
825            (
826                Some(frontend::Event::ReadyClean | frontend::Event::ReadyDirty),
827                SessionItem::ReadyForQuery {
828                    status,
829                    parameters_changed,
830                },
831            ) => ResetCompleteTransition::Dirty {
832                conn: self.transition(),
833                status,
834                parameters_changed,
835            },
836            (
837                Some(frontend::Event::Error),
838                SessionItem::Message(BackendMessage::ErrorResponse(error)),
839            ) => ResetCompleteTransition::Error(self.transition(), error),
840            (_, item) => ResetCompleteTransition::Continue(self, item),
841        }
842    }
843}
844
845impl<S, P> Conn<S, P, Pristine> {
846    /// Conservatively records session-local state without changing protocol phase.
847    pub fn mark_dirty(self) -> Conn<S, P, Dirty> {
848        self.transition()
849    }
850}
851
852fn project_session_item(
853    state: frontend::RuntimeState,
854    item: &SessionItem,
855) -> Option<frontend::Event> {
856    match item {
857        SessionItem::Message(message) => frontend::project_external(state, message),
858        SessionItem::CommandComplete { tag, .. } => {
859            frontend::project_external(state, &BackendMessage::CommandComplete(tag.clone()))
860        }
861        SessionItem::ReadyForQuery { status, .. } => {
862            frontend::project_external(state, &BackendMessage::ReadyForQuery(*status))
863        }
864    }
865}
866
867fn cstr_frame(tag: u8, value: &[u8]) -> io::Result<Frame> {
868    if value.contains(&0) {
869        return Err(io::Error::new(
870            io::ErrorKind::InvalidInput,
871            "message string contains a NUL byte",
872        ));
873    }
874    let mut body = BytesMut::with_capacity(value.len() + 1);
875    body.extend_from_slice(value);
876    body.put_u8(0);
877    Ok(Frame {
878        tag,
879        body: body.freeze(),
880    })
881}
882
883fn empty_frame(tag: u8) -> Frame {
884    Frame {
885        tag,
886        body: Bytes::new(),
887    }
888}
889
890fn ready_state<S, P, C>(
891    conn: Conn<S, P, C>,
892    status: TransactionStatus,
893    parameters_changed: bool,
894) -> ReadyState<S, C> {
895    if status == TransactionStatus::Idle && !parameters_changed {
896        ReadyState::Clean(conn.transition())
897    } else {
898        ReadyState::Dirty {
899            conn: conn.transition(),
900            status,
901            parameters_changed,
902        }
903    }
904}
905
906#[cfg(test)]
907mod tests {
908    use super::*;
909
910    #[test]
911    fn extended_building_self_loops_then_syncs() {
912        let ready: Conn<(), Ready> = Conn::new(()).transition();
913        let building = ready.begin_extended();
914        let (building, _) = building
915            .push_parse(&Parse {
916                statement: Bytes::from_static(b"statement"),
917                query: Bytes::from_static(b"select $1"),
918                parameter_types: vec![23],
919            })
920            .expect("valid Parse");
921        let (bound, _) = building
922            .push_bind(&Bind {
923                portal: Bytes::from_static(b"portal"),
924                statement: Bytes::from_static(b"statement"),
925                parameter_formats: vec![0],
926                parameters: vec![Some(Bytes::from_static(b"42"))],
927                result_formats: vec![0],
928            })
929            .expect("valid Bind");
930        let (bound, _) = bound
931            .push_execute(&Execute {
932                portal: Bytes::from_static(b"portal"),
933                max_rows: 0,
934            })
935            .expect("valid Execute");
936        let (awaiting_ready, sync) = bound.push_sync();
937        assert_eq!(sync.tag, b'S');
938        awaiting_ready.into_transport();
939    }
940
941    #[test]
942    fn function_call_requires_result_then_ready() {
943        let ready: Conn<(), Ready> = Conn::new(()).transition();
944        let call = FunctionCall {
945            function_oid: 42,
946            argument_formats: vec![1],
947            arguments: vec![Some(Bytes::from_static(b"argument"))],
948            result_format: 1,
949        };
950        let (calling, frame) = ready.push_function_call(&call).unwrap();
951        assert_eq!(frame.tag, b'F');
952
953        let FunctionCallTransition::Response(awaiting_ready, result) = calling
954            .offer(BackendMessage::FunctionCallResponse(Bytes::from_static(
955                b"result",
956            )))
957            .unwrap()
958        else {
959            panic!("function result projected to the wrong branch")
960        };
961        assert_eq!(result, Bytes::from_static(b"result"));
962        let AwaitingReadyTransition::Ready(ReadyState::Clean(ready)) =
963            awaiting_ready.offer(SessionItem::ReadyForQuery {
964                status: TransactionStatus::Idle,
965                parameters_changed: false,
966            })
967        else {
968            panic!("function call did not return to ready")
969        };
970        ready.into_transport();
971    }
972
973    #[test]
974    fn ready_session_can_terminate_gracefully() {
975        let ready: Conn<(), Ready> = Conn::new(()).transition();
976        let (terminated, frame) = ready.push_terminate();
977        assert_eq!(frame.tag, b'X');
978        assert!(frame.body.is_empty());
979        terminated.into_transport();
980    }
981
982    #[test]
983    fn copy_both_waits_for_both_half_closes() {
984        use crate::grammar::frontend::{Event, RuntimeFsm, RuntimeState};
985
986        let mut client_first = RuntimeFsm::new();
987        client_first.step(Event::Query).unwrap();
988        client_first.step(Event::CopyBoth).unwrap();
989        let open: Conn<(), CopyBoth> = Conn::new(()).transition();
990        let (client_done, frame) = open.push_copy_done();
991        client_first.step(Event::SendCopyDone).unwrap();
992        assert_eq!(frame.tag, b'c');
993        let CopyBothClientDoneReceive::Data(client_done, data) = client_done
994            .offer(SessionItem::Message(BackendMessage::CopyData(
995                Bytes::from_static(b"after client close"),
996            )))
997            .unwrap()
998        else {
999            panic!("backend data projected to the wrong branch")
1000        };
1001        client_first.step(Event::ReceiveCopyData).unwrap();
1002        assert_eq!(data, Bytes::from_static(b"after client close"));
1003        let CopyBothClientDoneReceive::Done(awaiting) = client_done
1004            .offer(SessionItem::Message(BackendMessage::CopyDone))
1005            .unwrap()
1006        else {
1007            panic!("backend close projected to the wrong branch")
1008        };
1009        client_first.step(Event::ReceiveCopyDone).unwrap();
1010        assert_eq!(client_first.state(), RuntimeState::AwaitingReady);
1011        awaiting.into_transport();
1012
1013        let mut server_first = RuntimeFsm::new();
1014        server_first.step(Event::Query).unwrap();
1015        server_first.step(Event::CopyBoth).unwrap();
1016        let open: Conn<(), CopyBoth> = Conn::new(()).transition();
1017        let CopyBothReceive::Done(server_done) = open
1018            .offer(SessionItem::Message(BackendMessage::CopyDone))
1019            .unwrap()
1020        else {
1021            panic!("backend close projected to the wrong branch")
1022        };
1023        server_first.step(Event::ReceiveCopyDone).unwrap();
1024        let (server_done, data) =
1025            server_done.push_copy_data(Bytes::from_static(b"after server close"));
1026        server_first.step(Event::SendCopyData).unwrap();
1027        assert_eq!(data.tag, b'd');
1028        let (awaiting, done) = server_done.push_copy_done();
1029        server_first.step(Event::SendCopyDone).unwrap();
1030        assert_eq!(done.tag, b'c');
1031        assert_eq!(server_first.state(), RuntimeState::AwaitingReady);
1032        awaiting.into_transport();
1033    }
1034
1035    #[test]
1036    fn copy_in_can_receive_an_early_backend_error() {
1037        let copy: Conn<(), CopyIn> = Conn::new(()).transition();
1038        let error = DiagnosticResponse {
1039            fields: vec![crate::codec::DiagnosticField {
1040                code: b'M',
1041                value: Bytes::from_static(b"copy cancelled"),
1042            }],
1043        };
1044        let CopyInTransition::Error(draining, received) = copy
1045            .offer(SessionItem::Message(BackendMessage::ErrorResponse(
1046                error.clone(),
1047            )))
1048            .unwrap();
1049        assert_eq!(received, error);
1050
1051        let DrainingTransition::Ready(ReadyState::Clean(ready)) =
1052            draining.offer(SessionItem::ReadyForQuery {
1053                status: TransactionStatus::Idle,
1054                parameters_changed: false,
1055            })
1056        else {
1057            panic!("COPY failure did not drain to readiness")
1058        };
1059        ready.release();
1060    }
1061
1062    #[test]
1063    fn copy_both_projects_typed_replication_without_losing_connection() {
1064        let open: Conn<(), CopyBoth> = Conn::new(()).transition();
1065        let status = FrontendReplication::StandbyStatus {
1066            written: 10,
1067            flushed: 9,
1068            applied: 8,
1069            client_time: 7,
1070            reply_requested: true,
1071        };
1072        let (open, frame) = open.push_replication(&status);
1073        assert_eq!(frame.body, status.encode());
1074
1075        let keepalive = BackendReplication::PrimaryKeepalive {
1076            wal_end: 11,
1077            server_time: 12,
1078            reply_requested: true,
1079        };
1080        let receive = open
1081            .offer(SessionItem::Message(BackendMessage::CopyData(
1082                keepalive.encode(),
1083            )))
1084            .unwrap();
1085        let ReplicationReceive::Message(open, decoded) = receive.decode_replication().unwrap()
1086        else {
1087            panic!("keepalive projected to the wrong branch")
1088        };
1089        assert_eq!(decoded, keepalive);
1090        open.into_transport();
1091
1092        let open: Conn<(), CopyBoth> = Conn::new(()).transition();
1093        let receive = open
1094            .offer(SessionItem::Message(BackendMessage::CopyData(
1095                Bytes::from_static(b"kshort"),
1096            )))
1097            .unwrap();
1098        let (open, _) = receive.decode_replication().unwrap_err();
1099        open.into_transport();
1100    }
1101
1102    #[test]
1103    fn transaction_status_taints_ready_connection() {
1104        let query: Conn<(), SimpleQuery> = Conn::new(()).transition();
1105        let transition = query
1106            .offer(SessionItem::ReadyForQuery {
1107                status: TransactionStatus::InTransaction,
1108                parameters_changed: false,
1109            })
1110            .expect("ReadyForQuery is valid evidence");
1111        let SimpleTransition::Ready(ReadyState::Dirty {
1112            conn,
1113            status: TransactionStatus::InTransaction,
1114            parameters_changed: false,
1115        }) = transition
1116        else {
1117            panic!("transaction should taint readiness")
1118        };
1119        conn.into_transport();
1120    }
1121
1122    #[test]
1123    fn changed_parameters_taint_idle_connection() {
1124        let query: Conn<(), SimpleQuery> = Conn::new(()).transition();
1125        let transition = query
1126            .offer(SessionItem::ReadyForQuery {
1127                status: TransactionStatus::Idle,
1128                parameters_changed: true,
1129            })
1130            .expect("ReadyForQuery is valid evidence");
1131        let SimpleTransition::Ready(ReadyState::Dirty {
1132            conn,
1133            status: TransactionStatus::Idle,
1134            parameters_changed: true,
1135        }) = transition
1136        else {
1137            panic!("parameter change should taint readiness")
1138        };
1139        conn.into_transport();
1140    }
1141
1142    #[test]
1143    fn simple_queries_are_dirty_unless_inspection_proves_them_stateless() {
1144        fn require_dirty<S>(conn: Conn<S, Ready, Dirty>) {
1145            conn.into_transport();
1146        }
1147
1148        let ready: Conn<(), Ready> = Conn::new(()).transition();
1149        let (query, _) = ready.push_query(b"LISTEN events").unwrap();
1150        let SimpleTransition::Ready(ReadyState::Clean(dirty)) = query
1151            .offer(SessionItem::ReadyForQuery {
1152                status: TransactionStatus::Idle,
1153                parameters_changed: false,
1154            })
1155            .unwrap()
1156        else {
1157            panic!("idle readiness should preserve the query's dirty evidence")
1158        };
1159        require_dirty(dirty);
1160
1161        let ready: Conn<(), Ready> = Conn::new(()).transition();
1162        let (query, _) = ready.push_stateless_query(b"SELECT 1").unwrap();
1163        let SimpleTransition::Ready(ReadyState::Clean(pristine)) = query
1164            .offer(SessionItem::ReadyForQuery {
1165                status: TransactionStatus::Idle,
1166                parameters_changed: false,
1167            })
1168            .unwrap()
1169        else {
1170            panic!("stateless query should retain pristine evidence")
1171        };
1172        pristine.release();
1173    }
1174
1175    #[test]
1176    fn discard_all_evidence_recovers_pool_cleanliness() {
1177        let ready: Conn<(), Ready> = Conn::new(()).transition();
1178        let (resetting, frame) = ready.mark_dirty().begin_reset().unwrap();
1179        assert_eq!(frame.body, Bytes::from_static(b"ROLLBACK; DISCARD ALL\0"));
1180        let ResettingTransition::Continue(resetting, _) =
1181            resetting.offer(SessionItem::CommandComplete {
1182                tag: Bytes::from_static(b"ROLLBACK"),
1183                command: crate::demux::CommandIndex(0),
1184                notices: vec![],
1185            })
1186        else {
1187            panic!("ROLLBACK incorrectly completed reset")
1188        };
1189        let ResettingTransition::Complete(reset_complete) =
1190            resetting.offer(SessionItem::CommandComplete {
1191                tag: Bytes::from_static(b"DISCARD ALL"),
1192                command: crate::demux::CommandIndex(1),
1193                notices: vec![],
1194            })
1195        else {
1196            panic!("DISCARD ALL did not advance reset")
1197        };
1198        let ResetCompleteTransition::Ready(ready) =
1199            reset_complete.offer(SessionItem::ReadyForQuery {
1200                status: TransactionStatus::Idle,
1201                parameters_changed: false,
1202            })
1203        else {
1204            panic!("clean ready evidence did not restore pristine state")
1205        };
1206        ready.release();
1207    }
1208}
pg_proto/session.rs

pg_proto/
session.rs
