消息与读写缓冲

// Note: we can't use the same enum for both frontend and backend message formats
// because there are duplicated format codes between them.
//
// For example, `Close` (frontend) and `CommandComplete` (backend) both use format code `C`.
// <https://www.postgresql.org/docs/current/protocol-message-formats.html>
#[derive(Debug, PartialOrd, PartialEq)]
#[repr(u8)]
pub enum FrontendMessageFormat {
    Bind = b'B',
    Close = b'C',
    // ...
    Describe = b'D',
    Execute = b'E',
    #[expect(dead_code)]
    Flush = b'H',
    Parse = b'P',
    // ...
    PasswordPolymorphic = b'p',
    Query = b'Q',
    Sync = b'S',
    Terminate = b'X',
}

pub(crate) trait FrontendMessage: Sized {
    /// The format prefix of this message.
    const FORMAT: FrontendMessageFormat;

    /// Return the amount of space, in bytes, to reserve in the buffer passed to [`Self::encode_body()`].
    fn body_size_hint(&self) -> Saturating<usize>;

    /// Encode this type as a Frontend message in the Postgres protocol.
    ///
    /// The implementation should *not* include `Self::FORMAT` or the length prefix.
    fn encode_body(&self, buf: &mut Vec<u8>) -> Result<(), Error>;
    // ...
}

pub(crate) trait BackendMessage: Sized {
    /// The expected message format.
    ///
    /// <https://www.postgresql.org/docs/current/protocol-message-formats.html>
    const FORMAT: BackendMessageFormat;

    /// Decode this type from a Backend message in the Postgres protocol.
    ///
    /// The format code and length prefix have already been read and are not at the start of `bytes`.
    fn decode_body(buf: Bytes) -> Result<Self, Error>;
}

    fn put_length_prefixed<F>(&mut self, write_contents: F) -> Result<(), crate::Error>
    where
        F: FnOnce(&mut Vec<u8>) -> Result<(), crate::Error>,
    {
        // reserve space to write the prefixed length
        let offset = self.len();
        self.extend(&[0; 4]);

        // write the main body of the message
        let write_result = write_contents(self);

        let size_result = write_result.and_then(|_| {
            let size = self.len() - offset;
            i32::try_from(size)
                .map_err(|_| err_protocol!("message size out of range for protocol: {size}"))
        });

        match size_result {
            Ok(size) => {
                // now calculate the size of what we wrote and set the length value
                self[offset..(offset + 4)].copy_from_slice(&size.to_be_bytes());
                Ok(())
            }
            Err(e) => {
                // Put the buffer back to where it was.
                self.truncate(offset);
                Err(e)
            }
        }
    }

    #[inline(always)]
    pub(crate) fn write_msg(&mut self, message: impl FrontendMessage) -> Result<(), Error> {
        self.write(EncodeMessage(message))
    }

    pub(crate) async fn send<T>(&mut self, message: T) -> Result<(), Error>
    where
        T: FrontendMessage,
    {
        self.write_msg(message)?;
        self.flush().await?;
        Ok(())
    }

    pub(crate) async fn recv_unchecked(&mut self) -> Result<ReceivedMessage, Error> {
        // NOTE: to not break everything, this should be cancel-safe;
        // DO NOT modify `buf` unless a full message has been read
        self.inner
            .try_read(|buf| {
                // all packets in postgres start with a 5-byte header
                // this header contains the message type and the total length of the message
                let Some(mut header) = buf.get(..5) else {
                    return Ok(ControlFlow::Continue(5));
                };

                let format = BackendMessageFormat::try_from_u8(header.get_u8())?;

                let message_len = header.get_u32() as usize;

                let expected_len = message_len
                    .checked_add(1)
                    // this shouldn't really happen but is mostly a sanity check
                    .ok_or_else(|| {
                        err_protocol!("message_len + 1 overflows usize: {message_len}")
                    })?;

                if buf.len() < expected_len {
                    return Ok(ControlFlow::Continue(expected_len));
                }

                // `buf` SHOULD NOT be modified ABOVE this line

                // pop off the format code since it's not counted in `message_len`
                buf.advance(1);

                // consume the message, including the length prefix
                let mut contents = buf.split_to(message_len).freeze();

                // cut off the length prefix
                contents.advance(4);

                Ok(ControlFlow::Break(ReceivedMessage { format, contents }))
            })
            .await
    }

    pub async fn try_read<F, R>(&mut self, mut try_read: F) -> Result<R, Error>
    where
        F: FnMut(&mut BytesMut) -> Result<ControlFlow<R, usize>, Error>,
    {
        loop {
            let read_len = match try_read(&mut self.read_buf.read)? {
                ControlFlow::Continue(read_len) => read_len,
                ControlFlow::Break(ret) => return Ok(ret),
            };

            self.read_buf.read(read_len, &mut self.socket).await?;
        }
    }

    pub(crate) async fn recv(&mut self) -> Result<ReceivedMessage, Error> {
        loop {
            let message = self.recv_unchecked().await?;

            match message.format {
                BackendMessageFormat::ErrorResponse => {
                    // An error returned from the database server.
                    return Err(message.decode::<PgDatabaseError>()?.into());
                }

                BackendMessageFormat::NotificationResponse => {
                    if let Some(buffer) = &mut self.notifications {
                        let notification: Notification = message.decode()?;
                        let _ = buffer.send(notification).await;

                        continue;
                    }
                }

                BackendMessageFormat::ParameterStatus => {
                    // informs the frontend about the current (initial)
                    // setting of backend parameters

                    let ParameterStatus { name, value } = message.decode()?;
                    // TODO: handle `client_encoding`, `DateStyle` change

                    match name.as_str() {
                        "server_version" => {
                            self.server_version_num = parse_server_version(&value);
                        }
                        _ => {
                            self.parameter_statuses.insert(name, value);
                        }
                    }

                    continue;
                }

                BackendMessageFormat::NoticeResponse => {
                    // do we need this to be more configurable?
                    // if you are reading this comment and think so, open an issue

                    let notice: Notice = message.decode()?;

                    let (log_level, tracing_level) = match notice.severity() {
                        PgSeverity::Fatal | PgSeverity::Panic | PgSeverity::Error => {
                            (Level::Error, tracing::Level::ERROR)
                        }
                        PgSeverity::Warning => (Level::Warn, tracing::Level::WARN),
                        PgSeverity::Notice => (Level::Info, tracing::Level::INFO),
                        PgSeverity::Debug => (Level::Debug, tracing::Level::DEBUG),
                        PgSeverity::Info | PgSeverity::Log => (Level::Trace, tracing::Level::TRACE),
                    };
                    // ...
                    if log_is_enabled {
                        sqlx_core::private_tracing_dynamic_event!(
                            target: "sqlx::postgres::notice",
                            tracing_level,
                            message = notice.message()
                        );
                    }

                    continue;
                }

                _ => {}
            }

            return Ok(message);
        }
    }

  println!("--- notice ---");
  let r = sqlx::raw_sql("do $$ begin raise notice 'hello from pg1'; raise warning 'careful'; end $$")
    .execute(&mut conn)
    .await?;
  settle().await;
  println!("rows_affected = {}", r.rows_affected());

--- notice ---
  [conn 1] C->S Query("do $$ begin raise notice 'hello from pg1…")
  [conn 1] S->C NoticeResponse(S=NOTICE C=00000 M=hello from pg1)
  [conn 1] S->C NoticeResponse(S=WARNING C=01000 M=careful)
  [conn 1] S->C CommandComplete("DO"), ReadyForQuery(I)
 INFO sqlx::postgres::notice: hello from pg1
 WARN sqlx::postgres::notice: careful
rows_affected = 0

  println!("--- set DateStyle ---");
  conn.execute("set DateStyle = 'German'").await?;
  settle().await;
  quiet(&p, true);

  // simple query: text format, parsed by sqlx
  let row = conn.fetch_one("select date '2026-10-24' as d").await?;
  let text = row.try_get::<chrono::NaiveDate, _>("d");
  println!("simple query (text):    {text:?}");

  // prepared statement: binary format, DateStyle does not matter
  let row = sqlx::query("select date '2026-10-24' as d").fetch_one(&mut conn).await?;
  let bin = row.try_get::<chrono::NaiveDate, _>("d");
  println!("prepared query (binary): {bin:?}");

--- set DateStyle ---
  [conn 1] C->S Query("set DateStyle = 'German'")
  [conn 1] S->C CommandComplete("SET"), ParameterStatus(DateStyle=German, DMY), ReadyForQuery(I)
simple query (text):    Err(ColumnDecode { index: "\"d\"", source: ParseError(Invalid) })
prepared query (binary): Ok(2026-10-24)

  // LISTEN on a plain PgConnection (not a PgListener), then NOTIFY ourselves
  conn.execute("listen pg1_chan").await?;
  settle().await;
  let r = conn.execute("notify pg1_chan, 'hi'").await;
  settle().await;
  println!("notify: {r:?}");
  let r = sqlx::query_scalar::<_, i32>("select 1").fetch_one(&mut conn).await;
  settle().await;
  println!("next query: {r:?}");

  [conn 1] C->S Query("listen pg1_chan")
  [conn 1] S->C CommandComplete("LISTEN"), ReadyForQuery(I)
  [conn 1] C->S Query("notify pg1_chan, 'hi'")
  [conn 1] S->C CommandComplete("NOTIFY"), NotificationResponse, ReadyForQuery(I)
notify: Err(Protocol("execute: unexpected message: NotificationResponse (sqlx_postgres::connection::executor:366)"))
  [conn 1] C->S Parse(sqlx_s_1, "select 1"), Describe(S, sqlx_s_1), Sync
  ...
next query: Ok(1)

  println!("--- error during Execute, after two rows ---");
  let r = sqlx::query_scalar::<_, i32>("select 10 / (3 - g) from generate_series(1, $1) g")
    .bind(3_i32)
    .fetch_all(&mut conn)
    .await;
  settle().await;
  println!("result: {}", r.unwrap_err());

  println!("--- next query on the same connection ---");
  let v: Vec<i32> = sqlx::query_scalar("select 10 / (3 - g) from generate_series(1, $1) g")
    .bind(2_i32)
    .fetch_all(&mut conn)
    .await?;
  settle().await;
  println!("result: {v:?}");

  println!("--- error during Parse ---");
  let r = sqlx::query("selec 1").execute(&mut conn).await;
  settle().await;
  println!("result: {}", r.unwrap_err());
  let v: i32 = sqlx::query_scalar("select 2").fetch_one(&mut conn).await?;
  settle().await;
  println!("result: {v}");

  // ... 一段简单查询，见动脑筋

  println!("--- error inside a transaction ---");
  let mut tx = conn.begin().await?;
  let _ = sqlx::query("select 1 / $1").bind(0_i32).execute(&mut *tx).await;
  let r = sqlx::query("select 1").execute(&mut *tx).await;
  settle().await;
  println!("result: {}", r.unwrap_err());
  tx.rollback().await?;
  settle().await;

--- error during Execute, after two rows ---
  [conn 1] C->S Parse(sqlx_s_1, "select 10 / (3 - g) from generate_series…"), Describe(S, sqlx_s_1), Sync
  [conn 1] S->C ParseComplete, ParameterDescription, RowDescription, ReadyForQuery(I)
  [conn 1] C->S Bind(portal unnamed, stmt sqlx_s_1), Execute(limit 0), Close(P, unnamed), Sync
  [conn 1] S->C BindComplete, DataRow(列长 4), DataRow(列长 4), ErrorResponse(S=ERROR C=22012 M=division by zero)
  [conn 1] S->C ReadyForQuery(I)
result: error returned from database: division by zero at line 870
--- next query on the same connection ---
  [conn 1] C->S Bind(portal unnamed, stmt sqlx_s_1), Execute(limit 0), Close(P, unnamed), Sync
  [conn 1] S->C BindComplete, DataRow(列长 4), DataRow(列长 4), CommandComplete("SELECT 2"), CloseComplete, ReadyForQuery(I)
result: [5, 10]
--- error during Parse ---
  [conn 1] C->S Parse(sqlx_s_2, "selec 1"), Describe(S, sqlx_s_2), Sync
  [conn 1] S->C ErrorResponse(S=ERROR C=42601 M=syntax error at or near "selec")
  [conn 1] S->C ReadyForQuery(I)
result: error returned from database: syntax error at or near "selec" at line 1244
  [conn 1] C->S Parse(sqlx_s_3, "select 2"), Describe(S, sqlx_s_3), Sync
  [conn 1] S->C ParseComplete, ParameterDescription, RowDescription, ReadyForQuery(I)
  [conn 1] C->S Bind(portal unnamed, stmt sqlx_s_3), Execute(limit 0), Close(P, unnamed), Sync
  [conn 1] S->C BindComplete, DataRow(列长 4), CommandComplete("SELECT 1"), CloseComplete, ReadyForQuery(I)
result: 2
...
--- error inside a transaction ---
  [conn 1] C->S Query("BEGIN")
  [conn 1] S->C CommandComplete("BEGIN"), ReadyForQuery(T)
  [conn 1] C->S Parse(sqlx_s_4, "select 1 / $1"), Describe(S, sqlx_s_4), Sync
  [conn 1] S->C ParseComplete, ParameterDescription, RowDescription, ReadyForQuery(T)
  [conn 1] C->S Bind(portal unnamed, stmt sqlx_s_4), Execute(limit 0), Close(P, unnamed), Sync
  [conn 1] S->C ErrorResponse(S=ERROR C=22012 M=division by zero)
  [conn 1] S->C ReadyForQuery(E)
  [conn 1] C->S Parse(sqlx_s_5, "select 1"), Describe(S, sqlx_s_5), Sync
  [conn 1] S->C ErrorResponse(S=ERROR C=25P02 M=current transaction is aborted, commands ignored until end of transaction block)
  [conn 1] S->C ReadyForQuery(E)
result: error returned from database: current transaction is aborted, commands ignored until end of transaction block at line 1498
  [conn 1] C->S Query("ROLLBACK")
  [conn 1] S->C CommandComplete("ROLLBACK"), ReadyForQuery(I)

--- simple query, error in the second statement ---
  [conn 1] C->S Query("select 1; select 1 / 0; select 3")
  [conn 1] S->C RowDescription, DataRow(列长 1), CommandComplete("SELECT 1"), ErrorResponse(S=ERROR C=22012 M=division by zero), ReadyForQuery(I)
result: error returned from database: division by zero at line 870