138 lines
3.7 KiB
Rust
138 lines
3.7 KiB
Rust
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::{
|
|
BackendMessage, ColumnDescription, CommandCompleteData, DataRowData, ErrorResponseData,
|
|
ReadyForQueryData, RowDescriptionData,
|
|
};
|
|
use proto::message::primitive::pglist::PgList;
|
|
use proto::message::primitive::pgoid::PgOid;
|
|
use proto::writer::backend::BackendProtoWriter;
|
|
use proto::writer::protowriter::ProtoFlush;
|
|
use std::io::Error;
|
|
use std::time::Duration;
|
|
|
|
pub struct ServerProtoWrapper<W>(W, Option<Duration>);
|
|
|
|
impl<W> ServerProtoWrapper<W>
|
|
where
|
|
W: BackendProtoWriter + ProtoFlush + Send,
|
|
{
|
|
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(())
|
|
}
|
|
|
|
pub async fn write_ready_for_query(&mut self) -> anyhow::Result<()> {
|
|
self.0
|
|
.write_proto(ReadyForQueryData { status: b'I' }.into())
|
|
.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,
|
|
columns: &ColumnSelection,
|
|
) -> anyhow::Result<()> {
|
|
let columns = columns
|
|
.iter()
|
|
.map(|column| column_to_description(table_schema, *column))
|
|
.collect::<anyhow::Result<Vec<ColumnDescription>>>()?;
|
|
|
|
self.0
|
|
.write_proto(
|
|
RowDescriptionData {
|
|
columns: columns.into(),
|
|
}
|
|
.into(),
|
|
)
|
|
.await?;
|
|
Ok(())
|
|
}
|
|
|
|
async fn write_table_row(&mut self, row: &RestrictedRow) -> anyhow::Result<()> {
|
|
let values = row
|
|
.iter()
|
|
.map(|(_, value)| value.as_text_bytes().into())
|
|
.collect::<Vec<PgList<u8, i32>>>();
|
|
|
|
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.0
|
|
.write_proto(BackendMessage::CommandComplete(CommandCompleteData {
|
|
tag: status.to_string().into(),
|
|
}))
|
|
.await?;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
fn column_to_description(
|
|
schema: &TableSchema,
|
|
column: Column,
|
|
) -> anyhow::Result<ColumnDescription> {
|
|
let table_name = schema.table_name();
|
|
let table_oid = PgOid::from_unique_name(table_name);
|
|
|
|
let column_type = schema.column_type(column);
|
|
let name = schema.column_name_from_column(column);
|
|
|
|
let column_index = column.try_into()?;
|
|
let type_oid = column_type.type_oid();
|
|
let type_size = column_type.type_size();
|
|
|
|
Ok(ColumnDescription {
|
|
name: name.to_string().into(),
|
|
table_oid,
|
|
column_index,
|
|
type_oid,
|
|
type_size,
|
|
type_modifier: -1,
|
|
format_code: 0, // text format
|
|
})
|
|
}
|