rrdb 0.0.2

ready
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

use bytes::BytesMut;
use tokio_util::codec::Encoder;

use crate::pgwire::protocol::{
    backend::RowDescription, ConnectionCodec, FormatCode, ProtocolError,
};

use super::data_row_writer::DataRowWriter;

/// Supports batched rows for e.g. returning portal result sets.
///
/// NB: this struct only performs limited validation of column consistency across rows.
pub struct DataRowBatch {
    pub(crate) format_code: FormatCode,
    pub(crate) num_cols: usize,
    pub(crate) num_rows: usize,
    pub(crate) data: BytesMut,
    pub(crate) row: BytesMut,
}

impl DataRowBatch {
    /// Creates a new row batch using the given format code, requiring a certain number of columns per row.
    pub fn new(format_code: FormatCode, num_cols: usize) -> Self {
        Self {
            format_code,
            num_cols,
            num_rows: 0,
            data: BytesMut::new(),
            row: BytesMut::new(),
        }
    }

    /// Creates a [DataRowBatch] from the given [RowDescription].
    pub fn from_row_desc(desc: &RowDescription) -> Self {
        Self::new(desc.format_code, desc.fields.len())
    }

    /// Starts writing a new row.
    ///
    /// Returns a [DataRowWriter] that is responsible for the actual value encoding.
    pub fn create_row(&mut self) -> DataRowWriter {
        self.num_rows += 1;
        DataRowWriter::new(self)
    }

    /// Returns the number of rows currently written to this batch.
    pub fn num_rows(&self) -> usize {
        self.num_rows
    }
}

impl Encoder<DataRowBatch> for ConnectionCodec {
    type Error = ProtocolError;

    fn encode(&mut self, item: DataRowBatch, dst: &mut BytesMut) -> Result<(), Self::Error> {
        dst.extend(item.data);
        Ok(())
    }
}