sqlx_postgres/message/
mod.rs

1use sqlx_core::bytes::Bytes;
2use std::num::Saturating;
3
4use crate::error::Error;
5use crate::io::PgBufMutExt;
6
7mod authentication;
8mod backend_key_data;
9mod bind;
10mod close;
11mod command_complete;
12mod copy;
13mod data_row;
14mod describe;
15mod execute;
16mod flush;
17mod notification;
18mod parameter_description;
19mod parameter_status;
20mod parse;
21mod parse_complete;
22mod password;
23mod query;
24mod ready_for_query;
25mod response;
26mod row_description;
27mod sasl;
28mod ssl_request;
29mod startup;
30mod sync;
31mod terminate;
32
33pub use authentication::{Authentication, AuthenticationSasl};
34pub use backend_key_data::BackendKeyData;
35pub use bind::Bind;
36pub use close::Close;
37pub use command_complete::CommandComplete;
38pub use copy::{CopyData, CopyDone, CopyFail, CopyInResponse, CopyOutResponse, CopyResponseData};
39pub use data_row::DataRow;
40pub use describe::Describe;
41pub use execute::Execute;
42#[allow(unused_imports)]
43pub use flush::Flush;
44pub use notification::Notification;
45pub use parameter_description::ParameterDescription;
46pub use parameter_status::ParameterStatus;
47pub use parse::Parse;
48pub use parse_complete::ParseComplete;
49pub use password::Password;
50pub use query::Query;
51pub use ready_for_query::{ReadyForQuery, TransactionStatus};
52pub use response::{Notice, PgSeverity};
53pub use row_description::RowDescription;
54pub use sasl::{SaslInitialResponse, SaslResponse};
55use sqlx_core::io::ProtocolEncode;
56pub use ssl_request::SslRequest;
57pub use startup::Startup;
58pub use sync::Sync;
59pub use terminate::Terminate;
60
61// Note: we can't use the same enum for both frontend and backend message formats
62// because there are duplicated format codes between them.
63//
64// For example, `Close` (frontend) and `CommandComplete` (backend) both use format code `C`.
65// <https://www.postgresql.org/docs/current/protocol-message-formats.html>
66#[derive(Debug, PartialOrd, PartialEq)]
67#[repr(u8)]
68pub enum FrontendMessageFormat {
69    Bind = b'B',
70    Close = b'C',
71    CopyData = b'd',
72    CopyDone = b'c',
73    CopyFail = b'f',
74    Describe = b'D',
75    Execute = b'E',
76    Flush = b'H',
77    Parse = b'P',
78    /// This message format is polymorphic. It's used for:
79    ///
80    /// * Plain password responses
81    /// * MD5 password responses
82    /// * SASL responses
83    /// * GSSAPI/SSPI responses
84    PasswordPolymorphic = b'p',
85    Query = b'Q',
86    Sync = b'S',
87    Terminate = b'X',
88}
89
90#[derive(Debug, PartialOrd, PartialEq)]
91#[repr(u8)]
92pub enum BackendMessageFormat {
93    Authentication,
94    BackendKeyData,
95    BindComplete,
96    CloseComplete,
97    CommandComplete,
98    CopyData,
99    CopyDone,
100    CopyInResponse,
101    CopyOutResponse,
102    DataRow,
103    EmptyQueryResponse,
104    ErrorResponse,
105    NoData,
106    NoticeResponse,
107    NotificationResponse,
108    ParameterDescription,
109    ParameterStatus,
110    ParseComplete,
111    PortalSuspended,
112    ReadyForQuery,
113    RowDescription,
114}
115
116#[derive(Debug)]
117pub struct ReceivedMessage {
118    pub format: BackendMessageFormat,
119    pub contents: Bytes,
120}
121
122impl ReceivedMessage {
123    #[inline]
124    pub fn decode<T>(self) -> Result<T, Error>
125    where
126        T: BackendMessage,
127    {
128        if T::FORMAT != self.format {
129            return Err(err_protocol!(
130                "Postgres protocol error: expected {:?}, got {:?}",
131                T::FORMAT,
132                self.format
133            ));
134        }
135
136        T::decode_body(self.contents).map_err(|e| match e {
137            Error::Protocol(s) => {
138                err_protocol!("Postgres protocol error (reading {:?}): {s}", self.format)
139            }
140            other => other,
141        })
142    }
143}
144
145impl BackendMessageFormat {
146    pub fn try_from_u8(v: u8) -> Result<Self, Error> {
147        // https://www.postgresql.org/docs/current/protocol-message-formats.html
148
149        Ok(match v {
150            b'1' => BackendMessageFormat::ParseComplete,
151            b'2' => BackendMessageFormat::BindComplete,
152            b'3' => BackendMessageFormat::CloseComplete,
153            b'C' => BackendMessageFormat::CommandComplete,
154            b'd' => BackendMessageFormat::CopyData,
155            b'c' => BackendMessageFormat::CopyDone,
156            b'G' => BackendMessageFormat::CopyInResponse,
157            b'H' => BackendMessageFormat::CopyOutResponse,
158            b'D' => BackendMessageFormat::DataRow,
159            b'E' => BackendMessageFormat::ErrorResponse,
160            b'I' => BackendMessageFormat::EmptyQueryResponse,
161            b'A' => BackendMessageFormat::NotificationResponse,
162            b'K' => BackendMessageFormat::BackendKeyData,
163            b'N' => BackendMessageFormat::NoticeResponse,
164            b'R' => BackendMessageFormat::Authentication,
165            b'S' => BackendMessageFormat::ParameterStatus,
166            b'T' => BackendMessageFormat::RowDescription,
167            b'Z' => BackendMessageFormat::ReadyForQuery,
168            b'n' => BackendMessageFormat::NoData,
169            b's' => BackendMessageFormat::PortalSuspended,
170            b't' => BackendMessageFormat::ParameterDescription,
171
172            _ => return Err(err_protocol!("unknown message type: {:?}", v as char)),
173        })
174    }
175}
176
177pub(crate) trait FrontendMessage: Sized {
178    /// The format prefix of this message.
179    const FORMAT: FrontendMessageFormat;
180
181    /// Return the amount of space, in bytes, to reserve in the buffer passed to [`Self::encode_body()`].
182    fn body_size_hint(&self) -> Saturating<usize>;
183
184    /// Encode this type as a Frontend message in the Postgres protocol.
185    ///
186    /// The implementation should *not* include `Self::FORMAT` or the length prefix.
187    fn encode_body(&self, buf: &mut Vec<u8>) -> Result<(), Error>;
188
189    #[inline(always)]
190    #[cfg_attr(not(test), allow(dead_code))]
191    fn encode_msg(self, buf: &mut Vec<u8>) -> Result<(), Error> {
192        EncodeMessage(self).encode(buf)
193    }
194}
195
196pub(crate) trait BackendMessage: Sized {
197    /// The expected message format.
198    ///
199    /// <https://www.postgresql.org/docs/current/protocol-message-formats.html>
200    const FORMAT: BackendMessageFormat;
201
202    /// Decode this type from a Backend message in the Postgres protocol.
203    ///
204    /// The format code and length prefix have already been read and are not at the start of `bytes`.
205    fn decode_body(buf: Bytes) -> Result<Self, Error>;
206}
207
208pub struct EncodeMessage<F>(pub F);
209
210impl<F: FrontendMessage> ProtocolEncode<'_, ()> for EncodeMessage<F> {
211    fn encode_with(&self, buf: &mut Vec<u8>, _context: ()) -> Result<(), Error> {
212        let mut size_hint = self.0.body_size_hint();
213        // plus format code and length prefix
214        size_hint += 5;
215
216        // don't panic if `size_hint` is ridiculous
217        buf.try_reserve(size_hint.0).map_err(|e| {
218            err_protocol!(
219                "Postgres protocol: error allocating {} bytes for encoding message {:?}: {e}",
220                size_hint.0,
221                F::FORMAT,
222            )
223        })?;
224
225        buf.push(F::FORMAT as u8);
226
227        buf.put_length_prefixed(|buf| self.0.encode_body(buf))
228    }
229}
sqlx_postgres/message/mod.rs

sqlx_postgres/message/
mod.rs
