pgwire/messages/
copy.rs

1use bytes::{Buf, BufMut, Bytes, BytesMut};
2
3use super::{codec, DecodeContext, Message};
4use crate::error::PgWireResult;
5
6pub const MESSAGE_TYPE_BYTE_COPY_DATA: u8 = b'd';
7
8#[non_exhaustive]
9#[derive(PartialEq, Eq, Debug, Default, new)]
10pub struct CopyData {
11    pub data: Bytes,
12}
13
14impl Message for CopyData {
15    #[inline]
16    fn message_type() -> Option<u8> {
17        Some(MESSAGE_TYPE_BYTE_COPY_DATA)
18    }
19
20    #[inline]
21    fn max_message_length() -> usize {
22        super::LARGE_PACKET_SIZE_LIMIT
23    }
24
25    fn message_length(&self) -> usize {
26        4 + self.data.len()
27    }
28
29    fn encode_body(&self, buf: &mut BytesMut) -> PgWireResult<()> {
30        buf.put(self.data.as_ref());
31        Ok(())
32    }
33
34    fn decode_body(buf: &mut BytesMut, len: usize, _ctx: &DecodeContext) -> PgWireResult<Self> {
35        let data = buf.split_to(len - 4).freeze();
36        Ok(Self::new(data))
37    }
38}
39
40pub const MESSAGE_TYPE_BYTE_COPY_DONE: u8 = b'c';
41
42#[non_exhaustive]
43#[derive(PartialEq, Eq, Debug, Default, new)]
44pub struct CopyDone;
45
46impl Message for CopyDone {
47    #[inline]
48    fn message_type() -> Option<u8> {
49        Some(MESSAGE_TYPE_BYTE_COPY_DONE)
50    }
51
52    fn message_length(&self) -> usize {
53        4
54    }
55
56    fn encode_body(&self, _buf: &mut BytesMut) -> PgWireResult<()> {
57        Ok(())
58    }
59
60    fn decode_body(_buf: &mut BytesMut, _len: usize, _ctx: &DecodeContext) -> PgWireResult<Self> {
61        Ok(Self::new())
62    }
63}
64
65pub const MESSAGE_TYPE_BYTE_COPY_FAIL: u8 = b'f';
66
67#[non_exhaustive]
68#[derive(PartialEq, Eq, Debug, Default, new)]
69pub struct CopyFail {
70    pub message: String,
71}
72
73impl Message for CopyFail {
74    #[inline]
75    fn message_type() -> Option<u8> {
76        Some(MESSAGE_TYPE_BYTE_COPY_DONE)
77    }
78
79    fn message_length(&self) -> usize {
80        4 + self.message.len() + 1
81    }
82
83    fn encode_body(&self, buf: &mut BytesMut) -> PgWireResult<()> {
84        codec::put_cstring(buf, &self.message);
85        Ok(())
86    }
87
88    fn decode_body(buf: &mut BytesMut, _len: usize, _ctx: &DecodeContext) -> PgWireResult<Self> {
89        let msg = codec::get_cstring(buf).unwrap_or_else(|| "".to_owned());
90        Ok(Self::new(msg))
91    }
92}
93
94pub const MESSAGE_TYPE_BYTE_COPY_IN_RESPONSE: u8 = b'G';
95
96#[non_exhaustive]
97#[derive(PartialEq, Eq, Debug, Default, new)]
98pub struct CopyInResponse {
99    pub format: i8,
100    pub columns: i16,
101    pub column_formats: Vec<i16>,
102}
103
104impl Message for CopyInResponse {
105    #[inline]
106    fn message_type() -> Option<u8> {
107        Some(MESSAGE_TYPE_BYTE_COPY_IN_RESPONSE)
108    }
109
110    #[inline]
111    fn max_message_length() -> usize {
112        super::SMALL_BACKEND_PACKET_SIZE_LIMIT
113    }
114
115    fn message_length(&self) -> usize {
116        4 + 1 + 2 + self.column_formats.len() * 2
117    }
118
119    fn encode_body(&self, buf: &mut BytesMut) -> PgWireResult<()> {
120        buf.put_i8(self.format);
121        buf.put_i16(self.columns);
122        for cf in &self.column_formats {
123            buf.put_i16(*cf);
124        }
125        Ok(())
126    }
127
128    fn decode_body(buf: &mut BytesMut, _len: usize, _ctx: &DecodeContext) -> PgWireResult<Self> {
129        let format = buf.get_i8();
130        let columns = buf.get_i16();
131        let mut column_formats = Vec::with_capacity(columns as usize);
132        for _ in 0..columns {
133            column_formats.push(buf.get_i16());
134        }
135
136        Ok(Self::new(format, columns, column_formats))
137    }
138}
139
140pub const MESSAGE_TYPE_BYTE_COPY_OUT_RESPONSE: u8 = b'H';
141
142#[non_exhaustive]
143#[derive(PartialEq, Eq, Debug, Default, new)]
144pub struct CopyOutResponse {
145    pub format: i8,
146    pub columns: i16,
147    pub column_formats: Vec<i16>,
148}
149
150impl Message for CopyOutResponse {
151    #[inline]
152    fn message_type() -> Option<u8> {
153        Some(MESSAGE_TYPE_BYTE_COPY_OUT_RESPONSE)
154    }
155
156    #[inline]
157    fn max_message_length() -> usize {
158        super::SMALL_BACKEND_PACKET_SIZE_LIMIT
159    }
160
161    fn message_length(&self) -> usize {
162        4 + 1 + 2 + self.column_formats.len() * 2
163    }
164
165    fn encode_body(&self, buf: &mut BytesMut) -> PgWireResult<()> {
166        buf.put_i8(self.format);
167        buf.put_i16(self.columns);
168        for cf in &self.column_formats {
169            buf.put_i16(*cf);
170        }
171        Ok(())
172    }
173
174    fn decode_body(buf: &mut BytesMut, _len: usize, _ctx: &DecodeContext) -> PgWireResult<Self> {
175        let format = buf.get_i8();
176        let columns = buf.get_i16();
177        let mut column_formats = Vec::with_capacity(columns as usize);
178        for _ in 0..columns {
179            column_formats.push(buf.get_i16());
180        }
181
182        Ok(Self::new(format, columns, column_formats))
183    }
184}
185
186pub const MESSAGE_TYPE_BYTE_COPY_BOTH_RESPONSE: u8 = b'W';
187
188#[non_exhaustive]
189#[derive(PartialEq, Eq, Debug, Default, new)]
190pub struct CopyBothResponse {
191    pub format: i8,
192    pub columns: i16,
193    pub column_formats: Vec<i16>,
194}
195
196impl Message for CopyBothResponse {
197    #[inline]
198    fn message_type() -> Option<u8> {
199        Some(MESSAGE_TYPE_BYTE_COPY_BOTH_RESPONSE)
200    }
201
202    #[inline]
203    fn max_message_length() -> usize {
204        super::SMALL_BACKEND_PACKET_SIZE_LIMIT
205    }
206
207    fn message_length(&self) -> usize {
208        4 + 1 + 2 + self.column_formats.len() * 2
209    }
210
211    fn encode_body(&self, buf: &mut BytesMut) -> PgWireResult<()> {
212        buf.put_i8(self.format);
213        buf.put_i16(self.columns);
214        for cf in &self.column_formats {
215            buf.put_i16(*cf);
216        }
217        Ok(())
218    }
219
220    fn decode_body(buf: &mut BytesMut, _len: usize, _ctx: &DecodeContext) -> PgWireResult<Self> {
221        let format = buf.get_i8();
222        let columns = buf.get_i16();
223        let mut column_formats = Vec::with_capacity(columns as usize);
224        for _ in 0..columns {
225            column_formats.push(buf.get_i16());
226        }
227
228        Ok(Self::new(format, columns, column_formats))
229    }
230}