refactor: create trait for writing response from interpreter

server/src/proto_wrapper.rs

@@ -1,5 +1,6 @@
 use async_trait::async_trait;
 use minisql::operation::ColumnSelection;
+use minisql::response_writer::{CompleteStatus, ResponseWriter};
 use minisql::restricted_row::RestrictedRow;
 use minisql::schema::{Column, TableSchema};
 use proto::message::backend::{
@@ -9,71 +10,57 @@ use proto::message::backend::{
 use proto::message::primitive::pglist::PgList;
 use proto::message::primitive::pgoid::PgOid;
 use proto::writer::backend::BackendProtoWriter;
-use std::fmt;
+use proto::writer::protowriter::ProtoFlush;
+use std::io::Error;
+use std::time::Duration;
 
-pub enum CompleteStatus {
-    Insert { oid: i32, rows: i32 },
-    Delete(usize),
-    Select(usize),
-    CreateTable,
-    CreateIndex,
-}
+pub struct ServerProtoWrapper<W>(W, Option<Duration>);
 
-impl fmt::Display for CompleteStatus {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        match self {
-            CompleteStatus::Insert { oid, rows } => write!(f, "INSERT {} {}", oid, rows),
-            CompleteStatus::Delete(rows) => write!(f, "DELETE {}", rows),
-            CompleteStatus::Select(rows) => write!(f, "SELECT {}", rows),
-            CompleteStatus::CreateTable => write!(f, "CREATE TABLE"),
-            CompleteStatus::CreateIndex => write!(f, "CREATE INDEX"),
-        }
-    }
-}
-
-#[async_trait]
-pub trait ServerProto {
-    async fn write_error_message(&mut self, error_message: &str) -> anyhow::Result<()>;
-    async fn write_ready_for_query(&mut self) -> anyhow::Result<()>;
-    async fn write_empty_query(&mut self) -> anyhow::Result<()>;
-    async fn write_table_header(
-        &mut self,
-        table_schema: &TableSchema,
-        columns: &ColumnSelection,
-    ) -> anyhow::Result<()>;
-    async fn write_table_row(&mut self, row: &RestrictedRow) -> anyhow::Result<()>;
-    async fn write_command_complete(&mut self, status: CompleteStatus) -> anyhow::Result<()>;
-}
-
-#[async_trait]
-impl<W> ServerProto for W
+impl<W> ServerProtoWrapper<W>
 where
-    W: BackendProtoWriter + Send,
+    W: BackendProtoWriter + ProtoFlush + Send,
 {
-    async fn write_error_message(&mut self, error_message: &str) -> anyhow::Result<()> {
-        self.write_proto(
-            ErrorResponseData {
-                code: b'M',
-                message: format!("{error_message}\0").into(),
-            }
-            .into(),
-        )
-        .await?;
+    pub fn new(writer: W, throttle: Option<Duration>) -> Self {
+        Self(writer, throttle)
+    }
+
+    pub async fn write_error_message(&mut self, error_message: &str) -> anyhow::Result<()> {
+        self.0
+            .write_proto(
+                ErrorResponseData {
+                    code: b'M',
+                    message: format!("{error_message}\0").into(),
+                }
+                .into(),
+            )
+            .await?;
 
         Ok(())
     }
 
-    async fn write_ready_for_query(&mut self) -> anyhow::Result<()> {
-        self.write_proto(ReadyForQueryData { status: b'I' }.into())
+    pub async fn write_ready_for_query(&mut self) -> anyhow::Result<()> {
+        self.0
+            .write_proto(ReadyForQueryData { status: b'I' }.into())
             .await?;
         Ok(())
     }
+}
 
-    async fn write_empty_query(&mut self) -> anyhow::Result<()> {
-        self.write_proto(BackendMessage::EmptyQueryResponse).await?;
-        Ok(())
+#[async_trait]
+impl<W> ProtoFlush for ServerProtoWrapper<W>
+where
+    W: ProtoFlush + Send,
+{
+    async fn flush(&mut self) -> Result<(), Error> {
+        self.0.flush().await
     }
+}
 
+#[async_trait]
+impl<W> ResponseWriter for ServerProtoWrapper<W>
+where
+    W: BackendProtoWriter + ProtoFlush + Send,
+{
     async fn write_table_header(
         &mut self,
         table_schema: &TableSchema,
@@ -84,13 +71,14 @@ where
             .map(|column| column_to_description(table_schema, *column))
             .collect::<anyhow::Result<Vec<ColumnDescription>>>()?;
 
-        self.write_proto(
-            RowDescriptionData {
-                columns: columns.into(),
-            }
-            .into(),
-        )
-        .await?;
+        self.0
+            .write_proto(
+                RowDescriptionData {
+                    columns: columns.into(),
+                }
+                .into(),
+            )
+            .await?;
         Ok(())
     }
 
@@ -100,18 +88,26 @@ where
             .map(|(_, value)| value.as_text_bytes().into())
             .collect::<Vec<PgList<u8, i32>>>();
 
-        self.write_proto(BackendMessage::DataRow(DataRowData {
-            columns: values.into(),
-        }))
-        .await?;
+        self.0
+            .write_proto(BackendMessage::DataRow(DataRowData {
+                columns: values.into(),
+            }))
+            .await?;
+
+        if let Some(throttle) = self.1 {
+            self.0.flush().await?;
+            tokio::time::sleep(throttle).await;
+        }
+
         Ok(())
     }
 
     async fn write_command_complete(&mut self, status: CompleteStatus) -> anyhow::Result<()> {
-        self.write_proto(BackendMessage::CommandComplete(CommandCompleteData {
-            tag: status.to_string().into(),
-        }))
-        .await?;
+        self.0
+            .write_proto(BackendMessage::CommandComplete(CommandCompleteData {
+                tag: status.to_string().into(),
+            }))
+            .await?;
         Ok(())
     }
 }