zero_postgres/protocol/frontend/
copy.rs

1//! COPY protocol frontend messages.
2
3use crate::protocol::codec::MessageBuilder;
4
5/// Write a CopyData message.
6pub fn write_copy_data(buf: &mut Vec<u8>, data: &[u8]) {
7    let mut msg = MessageBuilder::new(buf, super::msg_type::COPY_DATA);
8    msg.write_bytes(data);
9    msg.finish();
10}
11
12/// Write a CopyDone message.
13pub fn write_copy_done(buf: &mut Vec<u8>) {
14    let msg = MessageBuilder::new(buf, super::msg_type::COPY_DONE);
15    msg.finish();
16}
17
18/// Write a CopyFail message.
19pub fn write_copy_fail(buf: &mut Vec<u8>, error_message: &str) {
20    let mut msg = MessageBuilder::new(buf, super::msg_type::COPY_FAIL);
21    msg.write_cstr(error_message);
22    msg.finish();
23}
24
25#[cfg(test)]
26mod tests {
27    use super::*;
28
29    #[test]
30    fn test_copy_data() {
31        let mut buf = Vec::new();
32        write_copy_data(&mut buf, b"hello\tworld\n");
33
34        assert_eq!(buf[0], b'd');
35        let len = i32::from_be_bytes([buf[1], buf[2], buf[3], buf[4]]);
36        assert_eq!(len as usize, buf.len() - 1);
37    }
38
39    #[test]
40    fn test_copy_done() {
41        let mut buf = Vec::new();
42        write_copy_done(&mut buf);
43
44        assert_eq!(buf.len(), 5);
45        assert_eq!(buf[0], b'c');
46        assert_eq!(&buf[1..5], &4_i32.to_be_bytes());
47    }
48
49    #[test]
50    fn test_copy_fail() {
51        let mut buf = Vec::new();
52        write_copy_fail(&mut buf, "error occurred");
53
54        assert_eq!(buf[0], b'f');
55    }
56}