sqlx_postgres/message/
copy.rs

1use crate::error::Result;
2use crate::io::BufMutExt;
3use crate::message::{
4    BackendMessage, BackendMessageFormat, FrontendMessage, FrontendMessageFormat,
5};
6use sqlx_core::bytes::{Buf, Bytes};
7use sqlx_core::Error;
8use std::num::Saturating;
9use std::ops::Deref;
10
11/// The same structure is sent for both `CopyInResponse` and `CopyOutResponse`
12pub struct CopyResponseData {
13    pub format: i8,
14    pub num_columns: i16,
15    pub format_codes: Vec<i16>,
16}
17
18pub struct CopyInResponse(pub CopyResponseData);
19
20#[allow(dead_code)]
21pub struct CopyOutResponse(pub CopyResponseData);
22
23pub struct CopyData<B>(pub B);
24
25pub struct CopyFail {
26    pub message: String,
27}
28
29pub struct CopyDone;
30
31impl CopyResponseData {
32    #[inline]
33    fn decode(mut buf: Bytes) -> Result<Self> {
34        let format = buf.get_i8();
35        let num_columns = buf.get_i16();
36
37        let format_codes = (0..num_columns).map(|_| buf.get_i16()).collect();
38
39        Ok(CopyResponseData {
40            format,
41            num_columns,
42            format_codes,
43        })
44    }
45}
46
47impl BackendMessage for CopyInResponse {
48    const FORMAT: BackendMessageFormat = BackendMessageFormat::CopyInResponse;
49
50    #[inline(always)]
51    fn decode_body(buf: Bytes) -> std::result::Result<Self, Error> {
52        Ok(Self(CopyResponseData::decode(buf)?))
53    }
54}
55
56impl BackendMessage for CopyOutResponse {
57    const FORMAT: BackendMessageFormat = BackendMessageFormat::CopyOutResponse;
58
59    #[inline(always)]
60    fn decode_body(buf: Bytes) -> std::result::Result<Self, Error> {
61        Ok(Self(CopyResponseData::decode(buf)?))
62    }
63}
64
65impl BackendMessage for CopyData<Bytes> {
66    const FORMAT: BackendMessageFormat = BackendMessageFormat::CopyData;
67
68    #[inline(always)]
69    fn decode_body(buf: Bytes) -> std::result::Result<Self, Error> {
70        Ok(Self(buf))
71    }
72}
73
74impl<B: Deref<Target = [u8]>> FrontendMessage for CopyData<B> {
75    const FORMAT: FrontendMessageFormat = FrontendMessageFormat::CopyData;
76
77    #[inline(always)]
78    fn body_size_hint(&self) -> Saturating<usize> {
79        Saturating(self.0.len())
80    }
81
82    #[inline(always)]
83    fn encode_body(&self, buf: &mut Vec<u8>) -> Result<(), Error> {
84        buf.extend_from_slice(&self.0);
85        Ok(())
86    }
87}
88
89impl FrontendMessage for CopyFail {
90    const FORMAT: FrontendMessageFormat = FrontendMessageFormat::CopyFail;
91
92    #[inline(always)]
93    fn body_size_hint(&self) -> Saturating<usize> {
94        Saturating(self.message.len())
95    }
96
97    #[inline(always)]
98    fn encode_body(&self, buf: &mut Vec<u8>) -> std::result::Result<(), Error> {
99        buf.put_str_nul(&self.message);
100        Ok(())
101    }
102}
103
104impl CopyFail {
105    #[inline(always)]
106    pub fn new(msg: impl Into<String>) -> CopyFail {
107        CopyFail {
108            message: msg.into(),
109        }
110    }
111}
112
113impl FrontendMessage for CopyDone {
114    const FORMAT: FrontendMessageFormat = FrontendMessageFormat::CopyDone;
115    #[inline(always)]
116    fn body_size_hint(&self) -> Saturating<usize> {
117        Saturating(0)
118    }
119
120    #[inline(always)]
121    fn encode_body(&self, _buf: &mut Vec<u8>) -> std::result::Result<(), Error> {
122        Ok(())
123    }
124}
125
126impl BackendMessage for CopyDone {
127    const FORMAT: BackendMessageFormat = BackendMessageFormat::CopyDone;
128
129    #[inline(always)]
130    fn decode_body(bytes: Bytes) -> std::result::Result<Self, Error> {
131        if !bytes.is_empty() {
132            // Not fatal but may indicate a protocol change
133            tracing::debug!(
134                "Postgres backend returned non-empty message for CopyDone: \"{}\"",
135                bytes.escape_ascii()
136            )
137        }
138
139        Ok(CopyDone)
140    }
141}