feat: add initial copy api design

examples/bench.rs

@@ -6,6 +6,7 @@ use futures::StreamExt;
 use tokio::net::TcpListener;
 
 use pgwire::api::auth::noop::NoopStartupHandler;
+use pgwire::api::copy::NoopCopyHandler;
 use pgwire::api::query::{PlaceholderExtendedQueryHandler, SimpleQueryHandler};
 use pgwire::api::results::{DataRowEncoder, FieldFormat, FieldInfo, QueryResponse, Response};
 use pgwire::api::{ClientInfo, MakeHandler, StatelessMakeHandler, Type};
@@ -73,6 +74,7 @@ pub async fn main() {
         PlaceholderExtendedQueryHandler,
     )));
     let authenticator = Arc::new(StatelessMakeHandler::new(Arc::new(NoopStartupHandler)));
+    let noop_copy_handler = Arc::new(NoopCopyHandler);
 
     let server_addr = "127.0.0.1:5433";
     let listener = TcpListener::bind(server_addr).await.unwrap();
@@ -82,13 +84,16 @@ pub async fn main() {
         let authenticator_ref = authenticator.make();
         let processor_ref = processor.make();
         let placeholder_ref = placeholder.make();
+        let copy_handler_ref = noop_copy_handler.clone();
+
         tokio::spawn(async move {
             process_socket(
                 incoming_socket.0,
                 None,
                 authenticator_ref,
                 processor_ref,
                 placeholder_ref,
+                copy_handler_ref,
             )
             .await
         });

examples/duckdb.rs

@@ -8,6 +8,7 @@ use futures::stream;
 use futures::Stream;
 use pgwire::api::auth::md5pass::{hash_md5_password, MakeMd5PasswordAuthStartupHandler};
 use pgwire::api::auth::{AuthSource, DefaultServerParameterProvider, LoginInfo, Password};
+use pgwire::api::copy::NoopCopyHandler;
 use pgwire::api::portal::{Format, Portal};
 use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler};
 use pgwire::api::results::{
@@ -350,6 +351,7 @@ pub async fn main() {
         Arc::new(parameters),
     ));
     let processor = Arc::new(MakeDuckDBBackend::new());
+    let noop_copy_handler = Arc::new(NoopCopyHandler);
 
     let server_addr = "127.0.0.1:5432";
     let listener = TcpListener::bind(server_addr).await.unwrap();
@@ -358,13 +360,16 @@ pub async fn main() {
         let incoming_socket = listener.accept().await.unwrap();
         let authenticator_ref = authenticator.make();
         let processor_ref = processor.make();
+        let copy_handler_ref = noop_copy_handler.clone();
+
         tokio::spawn(async move {
             process_socket(
                 incoming_socket.0,
                 None,
                 authenticator_ref,
                 processor_ref.clone(),
                 processor_ref,
+                copy_handler_ref,
             )
             .await
         });

examples/gluesql.rs

@@ -6,6 +6,7 @@ use tokio::net::TcpListener;
 
 use gluesql::prelude::*;
 use pgwire::api::auth::noop::NoopStartupHandler;
+use pgwire::api::copy::NoopCopyHandler;
 use pgwire::api::query::{PlaceholderExtendedQueryHandler, SimpleQueryHandler};
 use pgwire::api::results::{DataRowEncoder, FieldFormat, FieldInfo, QueryResponse, Response, Tag};
 use pgwire::api::{ClientInfo, MakeHandler, StatelessMakeHandler, Type};
@@ -170,6 +171,7 @@ pub async fn main() {
         PlaceholderExtendedQueryHandler,
     )));
     let authenticator = Arc::new(StatelessMakeHandler::new(Arc::new(NoopStartupHandler)));
+    let noop_copy_handler = Arc::new(NoopCopyHandler);
 
     let server_addr = "127.0.0.1:5432";
     let listener = TcpListener::bind(server_addr).await.unwrap();
@@ -179,13 +181,16 @@ pub async fn main() {
         let authenticator_ref = authenticator.make();
         let processor_ref = processor.make();
         let placeholder_ref = placeholder.make();
+        let copy_handler_ref = noop_copy_handler.clone();
+
         tokio::spawn(async move {
             process_socket(
                 incoming_socket.0,
                 None,
                 authenticator_ref,
                 processor_ref,
                 placeholder_ref,
+                copy_handler_ref,
             )
             .await
         });

examples/scram.rs

@@ -12,6 +12,7 @@ use tokio_rustls::TlsAcceptor;
 
 use pgwire::api::auth::scram::{gen_salted_password, MakeSASLScramAuthStartupHandler};
 use pgwire::api::auth::{AuthSource, DefaultServerParameterProvider, LoginInfo, Password};
+use pgwire::api::copy::NoopCopyHandler;
 use pgwire::api::query::{PlaceholderExtendedQueryHandler, SimpleQueryHandler};
 use pgwire::api::results::{Response, Tag};
 
@@ -83,6 +84,7 @@ pub async fn main() {
     let placeholder = Arc::new(StatelessMakeHandler::new(Arc::new(
         PlaceholderExtendedQueryHandler,
     )));
+    let noop_copy_handler = Arc::new(NoopCopyHandler);
     let mut authenticator = MakeSASLScramAuthStartupHandler::new(
         Arc::new(DummyAuthDB),
         Arc::new(DefaultServerParameterProvider::default()),
@@ -103,13 +105,16 @@ pub async fn main() {
         let authenticator_ref = authenticator.make();
         let processor_ref = processor.make();
         let placeholder_ref = placeholder.make();
+        let copy_handler_ref = noop_copy_handler.clone();
+
         tokio::spawn(async move {
             process_socket(
                 incoming_socket.0,
                 Some(tls_acceptor_ref),
                 authenticator_ref,
                 processor_ref,
                 placeholder_ref,
+                copy_handler_ref,
             )
             .await
         });

examples/secure_server.rs

@@ -11,6 +11,7 @@ use tokio_rustls::rustls::ServerConfig;
 use tokio_rustls::TlsAcceptor;
 
 use pgwire::api::auth::noop::NoopStartupHandler;
+use pgwire::api::copy::NoopCopyHandler;
 use pgwire::api::query::{PlaceholderExtendedQueryHandler, SimpleQueryHandler};
 use pgwire::api::results::{DataRowEncoder, FieldFormat, FieldInfo, QueryResponse, Response, Tag};
 use pgwire::api::{ClientInfo, MakeHandler, StatelessMakeHandler, Type};
@@ -84,6 +85,7 @@ pub async fn main() {
         PlaceholderExtendedQueryHandler,
     )));
     let authenticator = Arc::new(StatelessMakeHandler::new(Arc::new(NoopStartupHandler)));
+    let noop_copy_handler = Arc::new(NoopCopyHandler);
 
     let server_addr = "127.0.0.1:5433";
     let tls_acceptor = Arc::new(setup_tls().unwrap());
@@ -96,13 +98,15 @@ pub async fn main() {
         let authenticator_ref = authenticator.make();
         let processor_ref = processor.make();
         let placeholder_ref = placeholder.make();
+        let copy_handler_ref = noop_copy_handler.clone();
         tokio::spawn(async move {
             process_socket(
                 incoming_socket.0,
                 Some(tls_acceptor_ref),
                 authenticator_ref,
                 processor_ref,
                 placeholder_ref,
+                copy_handler_ref,
             )
             .await
         });

examples/server.rs

@@ -6,6 +6,7 @@ use futures::{stream, Sink, SinkExt, StreamExt};
 use tokio::net::TcpListener;
 
 use pgwire::api::auth::noop::NoopStartupHandler;
+use pgwire::api::copy::NoopCopyHandler;
 use pgwire::api::query::{PlaceholderExtendedQueryHandler, SimpleQueryHandler};
 use pgwire::api::results::{DataRowEncoder, FieldFormat, FieldInfo, QueryResponse, Response, Tag};
 use pgwire::api::{ClientInfo, MakeHandler, StatelessMakeHandler, Type};
@@ -76,6 +77,7 @@ pub async fn main() {
         PlaceholderExtendedQueryHandler,
     )));
     let authenticator = Arc::new(StatelessMakeHandler::new(Arc::new(NoopStartupHandler)));
+    let noop_copy_handler = Arc::new(NoopCopyHandler);
 
     let server_addr = "127.0.0.1:5432";
     let listener = TcpListener::bind(server_addr).await.unwrap();
@@ -85,13 +87,15 @@ pub async fn main() {
         let authenticator_ref = authenticator.make();
         let processor_ref = processor.make();
         let placeholder_ref = placeholder.make();
+        let copy_handler_ref = noop_copy_handler.clone();
         tokio::spawn(async move {
             process_socket(
                 incoming_socket.0,
                 None,
                 authenticator_ref,
                 processor_ref,
                 placeholder_ref,
+                copy_handler_ref,
             )
             .await
         });

examples/sqlite.rs

@@ -3,8 +3,10 @@ use std::sync::{Arc, Mutex};
 use async_trait::async_trait;
 use futures::stream;
 use futures::Stream;
+
 use pgwire::api::auth::md5pass::{hash_md5_password, MakeMd5PasswordAuthStartupHandler};
 use pgwire::api::auth::{AuthSource, DefaultServerParameterProvider, LoginInfo, Password};
+use pgwire::api::copy::NoopCopyHandler;
 use pgwire::api::portal::{Format, Portal};
 use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler};
 use pgwire::api::results::{
@@ -306,6 +308,7 @@ pub async fn main() {
         Arc::new(parameters),
     ));
     let processor = Arc::new(MakeSqliteBackend::new());
+    let noop_copy_handler = Arc::new(NoopCopyHandler);
 
     let server_addr = "127.0.0.1:5432";
     let listener = TcpListener::bind(server_addr).await.unwrap();
@@ -314,13 +317,16 @@ pub async fn main() {
         let incoming_socket = listener.accept().await.unwrap();
         let authenticator_ref = authenticator.make();
         let processor_ref = processor.make();
+        let copy_handler_ref = noop_copy_handler.clone();
+
         tokio::spawn(async move {
             process_socket(
                 incoming_socket.0,
                 None,
                 authenticator_ref,
                 processor_ref.clone(),
                 processor_ref,
+                copy_handler_ref,
             )
             .await
         });

src/api/copy.rs

@@ -0,0 +1,101 @@
+use async_trait::async_trait;
+use futures::sink::{Sink, SinkExt};
+use std::fmt::Debug;
+
+use crate::error::{PgWireError, PgWireResult};
+use crate::messages::copy::{
+    CopyBothResponse, CopyData, CopyDone, CopyFail, CopyInResponse, CopyOutResponse,
+};
+use crate::messages::PgWireBackendMessage;
+
+use super::ClientInfo;
+
+/// handler for copy messages
+#[async_trait]
+pub trait CopyHandler: Send + Sync {
+    async fn on_copy_data<C>(&self, _client: &mut C, _copy_data: CopyData) -> PgWireResult<()>
+    where
+        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
+        C::Error: Debug,
+        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
+    {
+        Ok(())
+    }
+
+    async fn on_copy_done<C>(&self, _client: &mut C, _done: CopyDone) -> PgWireResult<()>
+    where
+        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
+        C::Error: Debug,
+        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
+    {
+        Ok(())
+    }
+
+    async fn on_copy_fail<C>(&self, _client: &mut C, _fail: CopyFail) -> PgWireResult<()>
+    where
+        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
+        C::Error: Debug,
+        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
+    {
+        Ok(())
+    }
+}
+
+pub async fn send_copy_in_response<C>(
+    client: &mut C,
+    overall_format: i8,
+    columns: usize,
+    column_formats: Vec<i16>,
+) -> PgWireResult<()>
+where
+    C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
+    C::Error: Debug,
+    PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
+{
+    let resp = CopyInResponse::new(overall_format, columns as i16, column_formats);
+    client
+        .send(PgWireBackendMessage::CopyInResponse(resp))
+        .await?;
+    Ok(())
+}
+
+pub async fn send_copy_out_response<C>(
+    client: &mut C,
+    overall_format: i8,
+    columns: usize,
+    column_formats: Vec<i16>,
+) -> PgWireResult<()>
+where
+    C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
+    C::Error: Debug,
+    PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
+{
+    let resp = CopyOutResponse::new(overall_format, columns as i16, column_formats);
+    client
+        .send(PgWireBackendMessage::CopyOutResponse(resp))
+        .await?;
+    Ok(())
+}
+
+pub async fn send_copy_both_response<C>(
+    client: &mut C,
+    overall_format: i8,
+    columns: usize,
+    column_formats: Vec<i16>,
+) -> PgWireResult<()>
+where
+    C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
+    C::Error: Debug,
+    PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
+{
+    let resp = CopyBothResponse::new(overall_format, columns as i16, column_formats);
+    client
+        .send(PgWireBackendMessage::CopyBothResponse(resp))
+        .await?;
+    Ok(())
+}
+
+#[derive(Clone, Copy, Debug, Default)]
+pub struct NoopCopyHandler;
+
+impl CopyHandler for NoopCopyHandler {}

src/api/mod.rs

@@ -7,6 +7,7 @@ use std::sync::Arc;
 pub use postgres_types::Type;
 
 pub mod auth;
+pub mod copy;
 pub mod portal;
 pub mod query;
 pub mod results;
@@ -22,6 +23,7 @@ pub enum PgWireConnectionState {
     AuthenticationInProgress,
     ReadyForQuery,
     QueryInProgress,
+    CopyInProgress,
     AwaitingSync,
 }