pgWire 통신 연결 구현

src/lib/errors/mod.rs

@@ -3,3 +3,4 @@ pub mod into_error;
 pub mod lexing_error;
 pub mod parsing_error;
 pub mod predule;
+pub mod server_error;

src/lib/errors/predule.rs

@@ -2,3 +2,4 @@ pub use super::execute_error::*;
 pub use super::into_error::*;
 pub use super::lexing_error::*;
 pub use super::parsing_error::*;
+pub use super::server_error::*;

src/lib/errors/server_error.rs

@@ -0,0 +1,26 @@
+use std::string::ToString;
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct ServerError {
+    pub message: String,
+}
+
+impl ServerError {
+    pub fn new<T: ToString>(message: T) -> Self {
+        Self {
+            message: message.to_string(),
+        }
+    }
+
+    pub fn boxed<T: ToString>(message: T) -> Box<Self> {
+        Box::new(Self::new(message))
+    }
+}
+
+impl std::error::Error for ServerError {}
+
+impl std::fmt::Display for ServerError {
+    fn fmt(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(formatter, "server error: {}", self.message)
+    }
+}

src/lib/pgwire/connection/engine_func.rs

@@ -0,0 +1,4 @@
+use std::{pin::Pin, sync::Arc};
+
+pub type EngineFunc<E> =
+    Arc<dyn Fn() -> Pin<Box<dyn futures::Future<Output = E> + Send>> + Send + Sync>;

src/lib/pgwire/connection/mod.rs

@@ -13,3 +13,6 @@ pub use prepared_statement::*;
 
 pub mod bound_portal;
 pub use bound_portal::*;
+
+pub mod engine_func;
+pub use engine_func::*;

src/lib/pgwire/engine/mod.rs

@@ -4,3 +4,6 @@ pub use engine_impl::*;
 
 pub mod portal;
 pub use portal::*;
+
+pub mod rrdb;
+pub use rrdb::*;

src/lib/pgwire/engine/rrdb.rs

@@ -0,0 +1,65 @@
+use async_trait::async_trait;
+use sqlparser::ast::{Expr, SelectItem, SetExpr, Statement};
+
+use crate::lib::pgwire::engine::{Engine, Portal};
+use crate::lib::pgwire::protocol::{
+    DataRowBatch, DataTypeOid, ErrorResponse, FieldDescription, SqlState,
+};
+
+pub struct RRDBPortal;
+
+#[async_trait]
+impl Portal for RRDBPortal {
+    async fn fetch(&mut self, batch: &mut DataRowBatch) -> Result<(), ErrorResponse> {
+        let mut row = batch.create_row();
+        row.write_int4(1);
+        Ok(())
+    }
+}
+
+pub struct RRDBEngine;
+
+#[async_trait]
+impl Engine for RRDBEngine {
+    type PortalType = RRDBPortal;
+
+    async fn prepare(
+        &mut self,
+        statement: &Statement,
+    ) -> Result<Vec<FieldDescription>, ErrorResponse> {
+        if let Statement::Query(query) = &statement {
+            if let SetExpr::Select(select) = *query.body.clone() {
+                if select.projection.len() == 1 {
+                    if let SelectItem::UnnamedExpr(Expr::Identifier(column_name)) =
+                        &select.projection[0]
+                    {
+                        match column_name.value.as_str() {
+                            "test_error" => {
+                                return Err(ErrorResponse::error(
+                                    SqlState::DATA_EXCEPTION,
+                                    "test error",
+                                ))
+                            }
+                            "test_fatal" => {
+                                return Err(ErrorResponse::fatal(
+                                    SqlState::DATA_EXCEPTION,
+                                    "fatal error",
+                                ))
+                            }
+                            _ => (),
+                        }
+                    }
+                }
+            }
+        }
+
+        Ok(vec![FieldDescription {
+            name: "test".to_owned(),
+            data_type: DataTypeOid::Int4,
+        }])
+    }
+
+    async fn create_portal(&mut self, _: &Statement) -> Result<Self::PortalType, ErrorResponse> {
+        Ok(RRDBPortal)
+    }
+}

src/lib/server/core.rs

@@ -1,17 +1,21 @@
+use std::error::Error;
+use std::sync::Arc;
+
 use crate::lib::ast::predule::{DDLStatement, SQLStatement};
+use crate::lib::errors::server_error::ServerError;
 use crate::lib::executor::predule::Executor;
 use crate::lib::optimizer::predule::Optimizer;
 use crate::lib::parser::predule::Parser;
+use crate::lib::pgwire::predule::{Connection, RRDBEngine};
 use crate::lib::server::predule::ServerOption;
 
-use tokio::io::AsyncReadExt;
 use tokio::net::TcpListener;
 
 pub struct Server {
     pub option: ServerOption,
 }
 
-async fn process_query(query: String) -> Result<(), Box<dyn std::error::Error>> {
+async fn _process_query(query: String) -> Result<(), Box<dyn std::error::Error>> {
     let mut parser = Parser::new(query)?;
     let executor = Executor::new();
 
@@ -41,54 +45,33 @@ impl Server {
         Self { option }
     }
 
-    pub async fn run(&self) -> Result<(), Box<dyn std::error::Error>> {
-        println!("# Server is Running at {}", self.option.port);
-
-        let listener = TcpListener::bind(format!("0.0.0.0:{}", self.option.port)).await?;
-
-        loop {
-            match listener.accept().await {
-                Ok((mut socket, address)) => {
-                    println!("# Connected: {}:{}", address.ip(), address.port());
+    /// Starts a server using a function responsible for producing engine instances and set of bind options.
+    ///
+    /// Returns once the server is listening for connections, with the accept loop
+    /// running as a background task, and returns the listener's local port.
+    ///
+    /// Useful for creating test harnesses binding to port 0 to select a random port.
+    pub async fn run(&self) -> Result<(), Box<dyn Error>> {
+        let listener =
+            TcpListener::bind((self.option.host.to_owned(), self.option.port as u16)).await?;
 
-                    tokio::spawn(async move {
-                        let mut buffer = [0; 1024];
-
-                        loop {
-                            match socket.read(&mut buffer).await {
-                                // socket closed
-                                Ok(n) => {
-                                    if n != 0 {
-                                        let query = String::from_utf8_lossy(&buffer).to_string();
-                                        println!("QUERY: {}", query);
-
-                                        match process_query(query.clone()).await {
-                                            Ok(_) => {
-                                                let response = "OK".to_string();
-                                                println!("RESPONSE: {}", response);
-                                            }
-                                            Err(error) => {
-                                                println!("ERROR: {}", error);
-                                            }
-                                        }
-                                        // TODO: 쿼리 실행 후 리턴값 반환
-                                    } else {
-                                        eprintln!("# 연결 종료");
-                                        break;
-                                    }
-                                }
-                                Err(e) => {
-                                    eprintln!("failed to read from socket; err = {:?}", e);
-                                    break;
-                                }
-                            };
-                        }
-                    });
-                }
-                Err(error) => {
-                    println!("# Error: {}", error);
-                }
+        let result = tokio::spawn(async move {
+            loop {
+                println!("??");
+                let (stream, _) = listener.accept().await.unwrap();
+                println!("22");
+                let engine_func = Arc::new(|| Box::pin(async { RRDBEngine }));
+                tokio::spawn(async move {
+                    let mut conn = Connection::new(engine_func().await);
+                    conn.run(stream).await.unwrap();
+                });
             }
+        })
+        .await;
+
+        match result {
+            Ok(_) => Ok(()),
+            Err(error) => Err(ServerError::boxed(error.to_string())),
         }
     }
 }