pgwire/api/
copy.rs

1use async_trait::async_trait;
2use futures::sink::{Sink, SinkExt};
3use futures::stream::StreamExt;
4use std::fmt::Debug;
5
6use crate::error::{ErrorInfo, PgWireError, PgWireResult};
7use crate::messages::PgWireBackendMessage;
8use crate::messages::copy::{
9    CopyBothResponse, CopyData, CopyDone, CopyFail, CopyInResponse, CopyOutResponse,
10};
11
12use super::ClientInfo;
13use super::results::{CopyResponse, Tag};
14
15/// handler for copy messages
16#[async_trait]
17pub trait CopyHandler: Send + Sync {
18    /// Called when copy data is received from the client.
19    async fn on_copy_data<C>(&self, _client: &mut C, _copy_data: CopyData) -> PgWireResult<()>
20    where
21        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
22        C::Error: Debug,
23        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>;
24
25    /// Called when the client signals the copy-in operation is complete.
26    async fn on_copy_done<C>(&self, _client: &mut C, _done: CopyDone) -> PgWireResult<()>
27    where
28        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
29        C::Error: Debug,
30        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>;
31
32    /// Called when the client aborts the copy-in operation.
33    async fn on_copy_fail<C>(&self, _client: &mut C, fail: CopyFail) -> PgWireError
34    where
35        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
36        C::Error: Debug,
37        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
38    {
39        PgWireError::UserError(Box::new(ErrorInfo::new(
40            "ERROR".to_owned(),
41            "XX000".to_owned(),
42            format!("COPY IN mode terminated by the user: {}", fail.message),
43        )))
44    }
45}
46
47/// Send a CopyInResponse message to the client.
48pub async fn send_copy_in_response<C>(client: &mut C, resp: CopyResponse) -> PgWireResult<()>
49where
50    C: Sink<PgWireBackendMessage> + Unpin,
51    C::Error: Debug,
52    PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
53{
54    let column_formats = resp.column_formats();
55    let resp = CopyInResponse::new(resp.format, resp.columns as i16, column_formats);
56    client
57        .send(PgWireBackendMessage::CopyInResponse(resp))
58        .await?;
59    Ok(())
60}
61
62/// Send a CopyOutResponse header and stream copy data to the client.
63pub async fn send_copy_out_response<C>(client: &mut C, resp: CopyResponse) -> PgWireResult<()>
64where
65    C: Sink<PgWireBackendMessage> + Unpin,
66    C::Error: Debug,
67    PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
68{
69    let column_formats = resp.column_formats();
70    let CopyResponse {
71        format,
72        columns,
73        mut data_stream,
74    } = resp;
75    let copy_resp = CopyOutResponse::new(format, columns as i16, column_formats);
76    client
77        .send(PgWireBackendMessage::CopyOutResponse(copy_resp))
78        .await?;
79
80    let mut rows = 0;
81
82    while let Some(copy_data) = data_stream.next().await {
83        match copy_data {
84            Ok(data) => {
85                if !data.data.is_empty() {
86                    // do not count trailer
87                    if data.data.as_ref() != [0xFF, 0xFF] {
88                        rows += 1;
89                    }
90                    client.feed(PgWireBackendMessage::CopyData(data)).await?;
91                }
92            }
93            Err(e) => {
94                let copy_fail = CopyFail::new(format!("{}", e));
95                client
96                    .send(PgWireBackendMessage::CopyFail(copy_fail))
97                    .await?;
98                return Err(e);
99            }
100        }
101    }
102
103    let copy_done = CopyDone::new();
104    client
105        .send(PgWireBackendMessage::CopyDone(copy_done))
106        .await?;
107
108    let tag = Tag::new("COPY").with_rows(rows);
109    client
110        .send(PgWireBackendMessage::CommandComplete(tag.into()))
111        .await?;
112
113    Ok(())
114}
115
116/// Send a CopyBothResponse header and stream copy data to the client.
117pub async fn send_copy_both_response<C>(client: &mut C, resp: CopyResponse) -> PgWireResult<()>
118where
119    C: Sink<PgWireBackendMessage> + Unpin,
120    C::Error: Debug,
121    PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
122{
123    let column_formats = resp.column_formats();
124    let CopyResponse {
125        format,
126        columns,
127        mut data_stream,
128    } = resp;
129    let copy_resp = CopyBothResponse::new(format, columns as i16, column_formats);
130    client
131        .send(PgWireBackendMessage::CopyBothResponse(copy_resp))
132        .await?;
133
134    let mut rows = 0;
135
136    while let Some(copy_data) = data_stream.next().await {
137        match copy_data {
138            Ok(data) => {
139                if !data.data.is_empty() {
140                    // do not count trailer
141                    if data.data.as_ref() != [0xFF, 0xFF] {
142                        rows += 1;
143                    }
144                    client.feed(PgWireBackendMessage::CopyData(data)).await?;
145                }
146            }
147            Err(e) => {
148                let copy_fail = CopyFail::new(format!("{}", e));
149                client
150                    .send(PgWireBackendMessage::CopyFail(copy_fail))
151                    .await?;
152                return Err(e);
153            }
154        }
155    }
156
157    let copy_done = CopyDone::new();
158    client
159        .send(PgWireBackendMessage::CopyDone(copy_done))
160        .await?;
161
162    let tag = Tag::new("COPY").with_rows(rows);
163    client
164        .send(PgWireBackendMessage::CommandComplete(tag.into()))
165        .await?;
166
167    Ok(())
168}
169
170#[async_trait]
171impl CopyHandler for super::NoopHandler {
172    async fn on_copy_data<C>(&self, _client: &mut C, _copy_data: CopyData) -> PgWireResult<()>
173    where
174        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
175        C::Error: Debug,
176        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
177    {
178        Err(PgWireError::UserError(Box::new(ErrorInfo::new(
179            "FATAL".to_owned(),
180            "08P01".to_owned(),
181            "This feature is not implemented.".to_string(),
182        ))))
183    }
184
185    async fn on_copy_done<C>(&self, _client: &mut C, _done: CopyDone) -> PgWireResult<()>
186    where
187        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
188        C::Error: Debug,
189        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
190    {
191        Err(PgWireError::UserError(Box::new(ErrorInfo::new(
192            "FATAL".to_owned(),
193            "08P01".to_owned(),
194            "This feature is not implemented.".to_string(),
195        ))))
196    }
197}
pgwire/api/copy.rs

pgwire/api/
copy.rs
