An example for using pgwire as proxy

Cargo.toml

@@ -49,6 +49,8 @@ gluesql = "0.13"
 ## for datafusion example
 datafusion = "33"
 sqlparser = "0.40"
+## for proxy example
+tokio-postgres = "0.7"
 
 [features]
 default = ["tokio", "time-format"]

examples/proxy.rs

@@ -0,0 +1,99 @@
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use tokio::net::TcpListener;
+use tokio_postgres::{Client, NoTls, SimpleQueryMessage};
+
+use pgwire::api::auth::noop::NoopStartupHandler;
+use pgwire::api::query::{PlaceholderExtendedQueryHandler, SimpleQueryHandler};
+use pgwire::api::results::{DataRowEncoder, FieldFormat, FieldInfo, QueryResponse, Response, Tag};
+use pgwire::api::{ClientInfo, MakeHandler, StatelessMakeHandler, Type};
+use pgwire::error::{PgWireError, PgWireResult};
+use pgwire::tokio::process_socket;
+
+pub struct ProxyProcessor {
+    upstream_client: Client,
+}
+
+#[async_trait]
+impl SimpleQueryHandler for ProxyProcessor {
+    async fn do_query<'a, C>(&self, _client: &C, query: &'a str) -> PgWireResult<Vec<Response<'a>>>
+    where
+        C: ClientInfo + Unpin + Send + Sync,
+    {
+        self.upstream_client
+            .simple_query(query)
+            .await
+            .map_err(|e| PgWireError::ApiError(Box::new(e)))
+            .map(|resp_msgs| {
+                let mut downstream_response = Vec::new();
+                let mut row_buf = Vec::new();
+                for resp in resp_msgs {
+                    match resp {
+                        SimpleQueryMessage::CommandComplete(count) => {
+                            if row_buf.is_empty() {
+                                downstream_response.push(Response::Execution(
+                                    Tag::new_for_execution("", Some(count as usize)),
+                                ));
+                            } else {
+                                // FIXME: convert SimpleQueryRows to
+                                // QueryResponse: including schema and data
+                                let query_response = row_buf.into();
+                                downstream_response.push(Response::Query(query_response));
+                            }
+                        }
+                        SimpleQueryMessage::Row(row) => {
+                            // TODO: convert simple query row to pgwire data row
+                            // and process the response when needed
+                            row_buf.push(row);
+                        }
+                        _ => {}
+                    }
+                }
+
+                downstream_response
+            })
+    }
+}
+
+#[tokio::main]
+pub async fn main() {
+    let (client, connection) = tokio_postgres::connect("host=127.0.0.1 user=postgres", NoTls)
+        .await
+        .expect("Cannot client upstream connection");
+    tokio::spawn(async move {
+        if let Err(e) = connection.await {
+            eprintln!("Upstream connection error: {}", e);
+        }
+    });
+
+    let processor = Arc::new(StatelessMakeHandler::new(Arc::new(ProxyProcessor {
+        upstream_client: client,
+    })));
+
+    // We have not implemented extended query in this server, use placeholder instead
+    let placeholder = Arc::new(StatelessMakeHandler::new(Arc::new(
+        PlaceholderExtendedQueryHandler,
+    )));
+    let authenticator = Arc::new(StatelessMakeHandler::new(Arc::new(NoopStartupHandler)));
+
+    let server_addr = "127.0.0.1:5431";
+    let listener = TcpListener::bind(server_addr).await.unwrap();
+    println!("Listening to {}", server_addr);
+    loop {
+        let incoming_socket = listener.accept().await.unwrap();
+        let authenticator_ref = authenticator.make();
+        let processor_ref = processor.make();
+        let placeholder_ref = placeholder.make();
+        tokio::spawn(async move {
+            process_socket(
+                incoming_socket.0,
+                None,
+                authenticator_ref,
+                processor_ref,
+                placeholder_ref,
+            )
+            .await
+        });
+    }
+}