sqlx_postgres/io/
buf_mut.rs

1use crate::io::{PortalId, StatementId};
2
3pub trait PgBufMutExt {
4    fn put_length_prefixed<F>(&mut self, f: F) -> Result<(), crate::Error>
5    where
6        F: FnOnce(&mut Vec<u8>) -> Result<(), crate::Error>;
7
8    fn put_statement_name(&mut self, id: StatementId);
9
10    fn put_portal_name(&mut self, id: PortalId);
11}
12
13impl PgBufMutExt for Vec<u8> {
14    // writes a length-prefixed message, this is used when encoding nearly all messages as postgres
15    // wants us to send the length of the often-variable-sized messages up front
16    fn put_length_prefixed<F>(&mut self, write_contents: F) -> Result<(), crate::Error>
17    where
18        F: FnOnce(&mut Vec<u8>) -> Result<(), crate::Error>,
19    {
20        // reserve space to write the prefixed length
21        let offset = self.len();
22        self.extend(&[0; 4]);
23
24        // write the main body of the message
25        let write_result = write_contents(self);
26
27        let size_result = write_result.and_then(|_| {
28            let size = self.len() - offset;
29            i32::try_from(size)
30                .map_err(|_| err_protocol!("message size out of range for protocol: {size}"))
31        });
32
33        match size_result {
34            Ok(size) => {
35                // now calculate the size of what we wrote and set the length value
36                self[offset..(offset + 4)].copy_from_slice(&size.to_be_bytes());
37                Ok(())
38            }
39            Err(e) => {
40                // Put the buffer back to where it was.
41                self.truncate(offset);
42                Err(e)
43            }
44        }
45    }
46
47    // writes a statement name by ID
48    #[inline]
49    fn put_statement_name(&mut self, id: StatementId) {
50        id.put_name_with_nul(self);
51    }
52
53    // writes a portal name by ID
54    #[inline]
55    fn put_portal_name(&mut self, id: PortalId) {
56        id.put_name_with_nul(self);
57    }
58}