sqlx_core_guts/postgres/message/
copy.rs

1use crate::error::Result;
2use crate::io::{BufExt, BufMutExt, Decode, Encode};
3use bytes::{Buf, BufMut, Bytes};
4use std::ops::Deref;
5
6/// The same structure is sent for both `CopyInResponse` and `CopyOutResponse`
7pub struct CopyResponse {
8    pub format: i8,
9    pub num_columns: i16,
10    pub format_codes: Vec<i16>,
11}
12
13pub struct CopyData<B>(pub B);
14
15pub struct CopyFail {
16    pub message: String,
17}
18
19pub struct CopyDone;
20
21impl Decode<'_> for CopyResponse {
22    fn decode_with(mut buf: Bytes, _: ()) -> Result<Self> {
23        let format = buf.get_i8();
24        let num_columns = buf.get_i16();
25
26        let format_codes = (0..num_columns).map(|_| buf.get_i16()).collect();
27
28        Ok(CopyResponse {
29            format,
30            num_columns,
31            format_codes,
32        })
33    }
34}
35
36impl Decode<'_> for CopyData<Bytes> {
37    fn decode_with(buf: Bytes, _: ()) -> Result<Self> {
38        // well.. that was easy
39        Ok(CopyData(buf))
40    }
41}
42
43impl<B: Deref<Target = [u8]>> Encode<'_> for CopyData<B> {
44    fn encode_with(&self, buf: &mut Vec<u8>, _context: ()) {
45        buf.push(b'd');
46        buf.put_u32(self.0.len() as u32 + 4);
47        buf.extend_from_slice(&self.0);
48    }
49}
50
51impl Decode<'_> for CopyFail {
52    fn decode_with(mut buf: Bytes, _: ()) -> Result<Self> {
53        Ok(CopyFail {
54            message: buf.get_str_nul()?,
55        })
56    }
57}
58
59impl Encode<'_> for CopyFail {
60    fn encode_with(&self, buf: &mut Vec<u8>, _: ()) {
61        let len = 4 + self.message.len() + 1;
62
63        buf.push(b'f'); // to pay respects
64        buf.put_u32(len as u32);
65        buf.put_str_nul(&self.message);
66    }
67}
68
69impl CopyFail {
70    pub fn new(msg: impl Into<String>) -> CopyFail {
71        CopyFail {
72            message: msg.into(),
73        }
74    }
75}
76
77impl Decode<'_> for CopyDone {
78    fn decode_with(buf: Bytes, _: ()) -> Result<Self> {
79        if buf.is_empty() {
80            Ok(CopyDone)
81        } else {
82            Err(err_protocol!(
83                "expected no data for CopyDone, got: {:?}",
84                buf
85            ))
86        }
87    }
88}
89
90impl Encode<'_> for CopyDone {
91    fn encode_with(&self, buf: &mut Vec<u8>, _: ()) {
92        buf.reserve(4);
93        buf.push(b'c');
94        buf.put_u32(4);
95    }
96}