connectorx/sources/postgres/

mod.rs

//! Source implementation for Postgres database, including the TLS support (client only).

mod connection;
mod errors;
mod typesystem;

pub use self::errors::PostgresSourceError;
pub use cidr_02::IpInet;
pub use connection::rewrite_tls_args;
pub use pgvector::{Bit, HalfVector, SparseVector, Vector};
pub use typesystem::{PostgresTypePairs, PostgresTypeSystem};

use crate::constants::DB_BUFFER_SIZE;
use crate::{
    data_order::DataOrder,
    errors::ConnectorXError,
    sources::{PartitionParser, Produce, Source, SourcePartition},
    sql::{count_query, CXQuery},
};
use anyhow::anyhow;
use chrono::{DateTime, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime, Utc};
use csv::{ReaderBuilder, StringRecord, StringRecordsIntoIter};
use fehler::{throw, throws};
use hex::decode;
use postgres::{
    binary_copy::{BinaryCopyOutIter, BinaryCopyOutRow},
    fallible_iterator::FallibleIterator,
    tls::{MakeTlsConnect, TlsConnect},
    Config, CopyOutReader, Row, RowIter, SimpleQueryMessage, Socket,
};
use r2d2::{Pool, PooledConnection};
use r2d2_postgres::PostgresConnectionManager;
use rust_decimal::Decimal;
use serde_json::{from_str, Value};
use sqlparser::dialect::PostgreSqlDialect;
use std::collections::HashMap;
use std::convert::TryFrom;
use std::marker::PhantomData;
use uuid::Uuid;

/// Protocol - Binary based bulk load
pub enum BinaryProtocol {}

/// Protocol - CSV based bulk load
pub enum CSVProtocol {}

/// Protocol - use Cursor
pub enum CursorProtocol {}

/// Protocol - use Simple Query
pub enum SimpleProtocol {}

type PgManager<C> = PostgresConnectionManager<C>;
type PgConn<C> = PooledConnection<PgManager<C>>;

macro_rules! impl_produce_unimplemented {
    ($(($protocol: ty, $t: ty, $msg: expr),)+) => {
        $(
            impl<'r> Produce<'r, $t> for $protocol {
                type Error = PostgresSourceError;

                #[throws(PostgresSourceError)]
                fn produce(&'r mut self) -> $t {
                   unimplemented!($msg);
                }
            }

            impl<'r> Produce<'r, Option<$t>> for $protocol {
                type Error = PostgresSourceError;

                #[throws(PostgresSourceError)]
                fn produce(&'r mut self) -> Option<$t> {
                   unimplemented!($msg);
                }
            }
        )+
    };
}

impl_produce_unimplemented!(
    (PostgresCSVSourceParser<'_>, HashMap<String, Option<String>>, "Please use `cursor` protocol for hstore type"),
    (PostgresCSVSourceParser<'_>, Vector, "Please use `binary` protocol for vector type"),
    (PostgresCSVSourceParser<'_>, HalfVector, "Please use `binary` protocol for halfvector type"),
    (PostgresCSVSourceParser<'_>, Bit, "Please use `binary` protocol for bit type"),
    (PostgresCSVSourceParser<'_>, SparseVector, "Please use `binary` protocol for sparsevector type"),


    (PostgresSimpleSourceParser,HashMap<String, Option<String>>, "unimplemented"),
    (PostgresSimpleSourceParser,Value, "unimplemented"),
    (PostgresSimpleSourceParser, Vector, "Please use `binary` protocol for vector type"),
    (PostgresSimpleSourceParser, HalfVector, "Please use `binary` protocol for halfvector type"),
    (PostgresSimpleSourceParser, Bit, "Please use `binary` protocol for bit type"),
    (PostgresSimpleSourceParser, SparseVector, "Please use `binary` protocol for sparsevector type"),

);

/// If any columns are range types, wrap the query in a subquery that casts
/// range columns to `::text`. For example:
///   SELECT * FROM orders WHERE id < 2
/// becomes:
///   SELECT "id", "period"::text, "name" FROM (SELECT * FROM orders WHERE id < 2) AS _cx_sub
///
/// This forces PG to serialize ranges as text, which makes binary and cursor
/// protocols work without needing range-specific FromSql impls.
/// For CSV and simple protocols the rewrite is harmless (text cast on text is a no-op).
fn maybe_rewrite_range_query(
    query: &CXQuery<String>,
    names: &[String],
    schema: &[PostgresTypeSystem],
) -> CXQuery<String> {
    let range_mask: Vec<bool> = schema
        .iter()
        .map(|ts| matches!(ts, PostgresTypeSystem::Range(_)))
        .collect();

    if !range_mask.iter().any(|&r| r) {
        return query.clone();
    }

    let cols: String = names
        .iter()
        .zip(range_mask.iter())
        .map(|(name, is_range)| {
            let quoted = quote_ident(name);
            if *is_range {
                format!("{}::text", quoted)
            } else {
                quoted
            }
        })
        .collect::<Vec<_>>()
        .join(", ");

    let rewritten = format!("SELECT {} FROM ({}) AS _cx_sub", cols, query.as_str());
    CXQuery::Wrapped(rewritten)
}

fn quote_ident(ident: &str) -> String {
    format!("\"{}\"", ident.replace('\"', "\"\""))
}

#[cfg(test)]
mod range_rewrite_tests {
    use super::{maybe_rewrite_range_query, quote_ident, PostgresTypeSystem};
    use crate::sql::CXQuery;

    #[test]
    fn quote_ident_escapes_embedded_quotes() {
        assert_eq!(quote_ident("a\"b"), "\"a\"\"b\"");
    }

    #[test]
    fn rewrite_escapes_column_names_and_casts_only_ranges() {
        let q = CXQuery::Naked("SELECT 1".to_string());
        let names = vec!["plain".to_string(), "a\"b".to_string()];
        let schema = vec![
            PostgresTypeSystem::Int4(true),
            PostgresTypeSystem::Range(true),
        ];
        let rewritten = maybe_rewrite_range_query(&q, &names, &schema);
        assert_eq!(
            rewritten.as_str(),
            "SELECT \"plain\", \"a\"\"b\"::text FROM (SELECT 1) AS _cx_sub"
        );
    }
}

// take a row and unwrap the interior field from column 0
fn convert_row<'b, R: TryFrom<usize> + postgres::types::FromSql<'b> + Clone>(row: &'b Row) -> R {
    let nrows: Option<R> = row.get(0);
    nrows.expect("Could not parse int result from count_query")
}

#[throws(PostgresSourceError)]
fn get_total_rows<C>(conn: &mut PgConn<C>, query: &CXQuery<String>) -> usize
where
    C: MakeTlsConnect<Socket> + Clone + 'static + Sync + Send,
    C::TlsConnect: Send,
    C::Stream: Send,
    <C::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
    let dialect = PostgreSqlDialect {};

    let row = conn.query_one(count_query(query, &dialect)?.as_str(), &[])?;
    let col_type = PostgresTypeSystem::from(row.columns()[0].type_());
    match col_type {
        PostgresTypeSystem::Int2(_) => convert_row::<i16>(&row) as usize,
        PostgresTypeSystem::Int4(_) => convert_row::<i32>(&row) as usize,
        PostgresTypeSystem::Int8(_) => convert_row::<i64>(&row) as usize,
        _ => throw!(anyhow!(
            "The result of the count query was not an int, aborting."
        )),
    }
}

pub struct PostgresSource<P, C>
where
    C: MakeTlsConnect<Socket> + Clone + 'static + Sync + Send,
    C::TlsConnect: Send,
    C::Stream: Send,
    <C::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
    pool: Pool<PgManager<C>>,
    origin_query: Option<String>,
    queries: Vec<CXQuery<String>>,
    names: Vec<String>,
    schema: Vec<PostgresTypeSystem>,
    pg_schema: Vec<postgres::types::Type>,
    pre_execution_queries: Option<Vec<String>>,
    _protocol: PhantomData<P>,
}

impl<P, C> PostgresSource<P, C>
where
    C: MakeTlsConnect<Socket> + Clone + 'static + Sync + Send,
    C::TlsConnect: Send,
    C::Stream: Send,
    <C::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
    #[throws(PostgresSourceError)]
    pub fn new(config: Config, tls: C, nconn: usize) -> Self {
        let manager = PostgresConnectionManager::new(config, tls);
        let pool = Pool::builder().max_size(nconn as u32).build(manager)?;

        Self {
            pool,
            origin_query: None,
            queries: vec![],
            names: vec![],
            schema: vec![],
            pg_schema: vec![],
            pre_execution_queries: None,
            _protocol: PhantomData,
        }
    }
}

impl<P, C> Source for PostgresSource<P, C>
where
    PostgresSourcePartition<P, C>:
        SourcePartition<TypeSystem = PostgresTypeSystem, Error = PostgresSourceError>,
    P: Send,
    C: MakeTlsConnect<Socket> + Clone + 'static + Sync + Send,
    C::TlsConnect: Send,
    C::Stream: Send,
    <C::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
    const DATA_ORDERS: &'static [DataOrder] = &[DataOrder::RowMajor];
    type Partition = PostgresSourcePartition<P, C>;
    type TypeSystem = PostgresTypeSystem;
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn set_data_order(&mut self, data_order: DataOrder) {
        if !matches!(data_order, DataOrder::RowMajor) {
            throw!(ConnectorXError::UnsupportedDataOrder(data_order));
        }
    }

    fn set_queries<Q: ToString>(&mut self, queries: &[CXQuery<Q>]) {
        self.queries = queries.iter().map(|q| q.map(Q::to_string)).collect();
    }

    fn set_origin_query(&mut self, query: Option<String>) {
        self.origin_query = query;
    }

    fn set_pre_execution_queries(&mut self, pre_execution_queries: Option<&[String]>) {
        self.pre_execution_queries = pre_execution_queries.map(|s| s.to_vec());
    }

    #[throws(PostgresSourceError)]
    fn fetch_metadata(&mut self) {
        assert!(!self.queries.is_empty());

        let mut conn = self.pool.get()?;
        let first_query = &self.queries[0];

        let stmt = conn.prepare(first_query.as_str())?;

        let (names, pg_types): (Vec<String>, Vec<postgres::types::Type>) = stmt
            .columns()
            .iter()
            .map(|col| (col.name().to_string(), col.type_().clone()))
            .unzip();

        self.names = names;
        self.schema = pg_types.iter().map(PostgresTypeSystem::from).collect();
        self.pg_schema = self
            .schema
            .iter()
            .zip(pg_types.iter())
            .map(|(t1, t2)| PostgresTypePairs(t2, t1).into())
            .collect();
    }

    #[throws(PostgresSourceError)]
    fn result_rows(&mut self) -> Option<usize> {
        match &self.origin_query {
            Some(q) => {
                let cxq = CXQuery::Naked(q.clone());
                let mut conn = self.pool.get()?;
                let nrows = get_total_rows(&mut conn, &cxq)?;
                Some(nrows)
            }
            None => None,
        }
    }

    fn names(&self) -> Vec<String> {
        self.names.clone()
    }

    fn schema(&self) -> Vec<Self::TypeSystem> {
        self.schema.clone()
    }

    #[throws(PostgresSourceError)]
    fn partition(self) -> Vec<Self::Partition> {
        let mut ret = vec![];
        for query in self.queries {
            let mut conn = self.pool.get()?;
            let rewritten = maybe_rewrite_range_query(&query, &self.names, &self.schema);

            if let Some(pre_queries) = &self.pre_execution_queries {
                for pre_query in pre_queries {
                    conn.query(pre_query, &[])?;
                }
            }

            ret.push(PostgresSourcePartition::<P, C>::new(
                conn,
                &rewritten,
                &self.schema,
                &self.pg_schema,
            ));
        }
        ret
    }
}

pub struct PostgresSourcePartition<P, C>
where
    C: MakeTlsConnect<Socket> + Clone + 'static + Sync + Send,
    C::TlsConnect: Send,
    C::Stream: Send,
    <C::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
    conn: PgConn<C>,
    query: CXQuery<String>,
    schema: Vec<PostgresTypeSystem>,
    pg_schema: Vec<postgres::types::Type>,
    nrows: usize,
    ncols: usize,
    _protocol: PhantomData<P>,
}

impl<P, C> PostgresSourcePartition<P, C>
where
    C: MakeTlsConnect<Socket> + Clone + 'static + Sync + Send,
    C::TlsConnect: Send,
    C::Stream: Send,
    <C::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
    pub fn new(
        conn: PgConn<C>,
        query: &CXQuery<String>,
        schema: &[PostgresTypeSystem],
        pg_schema: &[postgres::types::Type],
    ) -> Self {
        Self {
            conn,
            query: query.clone(),
            schema: schema.to_vec(),
            pg_schema: pg_schema.to_vec(),
            nrows: 0,
            ncols: schema.len(),
            _protocol: PhantomData,
        }
    }
}

impl<C> SourcePartition for PostgresSourcePartition<BinaryProtocol, C>
where
    C: MakeTlsConnect<Socket> + Clone + 'static + Sync + Send,
    C::TlsConnect: Send,
    C::Stream: Send,
    <C::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
    type TypeSystem = PostgresTypeSystem;
    type Parser<'a> = PostgresBinarySourcePartitionParser<'a>;
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn result_rows(&mut self) -> () {
        self.nrows = get_total_rows(&mut self.conn, &self.query)?;
    }

    #[throws(PostgresSourceError)]
    fn parser(&mut self) -> Self::Parser<'_> {
        let query = format!("COPY ({}) TO STDOUT WITH BINARY", self.query);
        let reader = self.conn.copy_out(&*query)?; // unless reading the data, it seems like issue the query is fast
        let iter = BinaryCopyOutIter::new(reader, &self.pg_schema);

        PostgresBinarySourcePartitionParser::new(iter, &self.schema)
    }

    fn nrows(&self) -> usize {
        self.nrows
    }

    fn ncols(&self) -> usize {
        self.ncols
    }
}

impl<C> SourcePartition for PostgresSourcePartition<CSVProtocol, C>
where
    C: MakeTlsConnect<Socket> + Clone + 'static + Sync + Send,
    C::TlsConnect: Send,
    C::Stream: Send,
    <C::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
    type TypeSystem = PostgresTypeSystem;
    type Parser<'a> = PostgresCSVSourceParser<'a>;
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn result_rows(&mut self) {
        self.nrows = get_total_rows(&mut self.conn, &self.query)?;
    }

    #[throws(PostgresSourceError)]
    fn parser(&mut self) -> Self::Parser<'_> {
        let query = format!("COPY ({}) TO STDOUT WITH CSV", self.query);
        let reader = self.conn.copy_out(&*query)?; // unless reading the data, it seems like issue the query is fast
        let iter = ReaderBuilder::new()
            .has_headers(false)
            .from_reader(reader)
            .into_records();

        PostgresCSVSourceParser::new(iter, &self.schema)
    }

    fn nrows(&self) -> usize {
        self.nrows
    }

    fn ncols(&self) -> usize {
        self.ncols
    }
}

impl<C> SourcePartition for PostgresSourcePartition<CursorProtocol, C>
where
    C: MakeTlsConnect<Socket> + Clone + 'static + Sync + Send,
    C::TlsConnect: Send,
    C::Stream: Send,
    <C::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
    type TypeSystem = PostgresTypeSystem;
    type Parser<'a> = PostgresRawSourceParser<'a>;
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn result_rows(&mut self) {
        self.nrows = get_total_rows(&mut self.conn, &self.query)?;
    }

    #[throws(PostgresSourceError)]
    fn parser(&mut self) -> Self::Parser<'_> {
        let iter = self
            .conn
            .query_raw::<_, bool, _>(self.query.as_str(), vec![])?; // unless reading the data, it seems like issue the query is fast
        PostgresRawSourceParser::new(iter, &self.schema)
    }

    fn nrows(&self) -> usize {
        self.nrows
    }

    fn ncols(&self) -> usize {
        self.ncols
    }
}
pub struct PostgresBinarySourcePartitionParser<'a> {
    iter: BinaryCopyOutIter<'a>,
    rowbuf: Vec<BinaryCopyOutRow>,
    ncols: usize,
    current_col: usize,
    current_row: usize,
    is_finished: bool,
}

impl<'a> PostgresBinarySourcePartitionParser<'a> {
    pub fn new(iter: BinaryCopyOutIter<'a>, schema: &[PostgresTypeSystem]) -> Self {
        Self {
            iter,
            rowbuf: Vec::with_capacity(DB_BUFFER_SIZE),
            ncols: schema.len(),
            current_row: 0,
            current_col: 0,
            is_finished: false,
        }
    }

    #[throws(PostgresSourceError)]
    fn next_loc(&mut self) -> (usize, usize) {
        let ret = (self.current_row, self.current_col);
        self.current_row += (self.current_col + 1) / self.ncols;
        self.current_col = (self.current_col + 1) % self.ncols;
        ret
    }
}

impl<'a> PartitionParser<'a> for PostgresBinarySourcePartitionParser<'a> {
    type TypeSystem = PostgresTypeSystem;
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn fetch_next(&mut self) -> (usize, bool) {
        assert!(self.current_col == 0);
        let remaining_rows = self.rowbuf.len() - self.current_row;
        if remaining_rows > 0 {
            return (remaining_rows, self.is_finished);
        } else if self.is_finished {
            return (0, self.is_finished);
        }

        // clear the buffer
        if !self.rowbuf.is_empty() {
            self.rowbuf.drain(..);
        }
        for _ in 0..DB_BUFFER_SIZE {
            match self.iter.next()? {
                Some(row) => {
                    self.rowbuf.push(row);
                }
                None => {
                    self.is_finished = true;
                    break;
                }
            }
        }

        // reset current cursor positions
        self.current_row = 0;
        self.current_col = 0;

        (self.rowbuf.len(), self.is_finished)
    }
}

macro_rules! impl_produce {
    ($($t: ty,)+) => {
        $(
            impl<'r, 'a> Produce<'r, $t> for PostgresBinarySourcePartitionParser<'a> {
                type Error = PostgresSourceError;

                #[throws(PostgresSourceError)]
                fn produce(&'r mut self) -> $t {
                    let (ridx, cidx) = self.next_loc()?;
                    let row = &self.rowbuf[ridx];
                    let val = row.try_get(cidx)?;
                    val
                }
            }

            impl<'r, 'a> Produce<'r, Option<$t>> for PostgresBinarySourcePartitionParser<'a> {
                type Error = PostgresSourceError;

                #[throws(PostgresSourceError)]
                fn produce(&'r mut self) -> Option<$t> {
                    let (ridx, cidx) = self.next_loc()?;
                    let row = &self.rowbuf[ridx];
                    let val = row.try_get(cidx)?;
                    val
                }
            }
        )+
    };
}

impl_produce!(
    i8,
    i16,
    i32,
    i64,
    u32,
    f32,
    f64,
    Decimal,
    bool,
    &'r str,
    Vec<u8>,
    NaiveTime,
    Uuid,
    Value,
    IpInet,
    Vector,
    HalfVector,
    Bit,
    SparseVector,
    Vec<Option<bool>>,
    Vec<Option<i16>>,
    Vec<Option<i32>>,
    Vec<Option<i64>>,
    Vec<Option<Decimal>>,
    Vec<Option<f32>>,
    Vec<Option<f64>>,
    Vec<Option<String>>,
);

impl<'r> Produce<'r, NaiveDateTime> for PostgresBinarySourcePartitionParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> NaiveDateTime {
        let (ridx, cidx) = self.next_loc()?;
        let row = &self.rowbuf[ridx];
        let val = row.try_get(cidx)?;
        match val {
            postgres::types::Timestamp::PosInfinity => NaiveDateTime::MAX,
            postgres::types::Timestamp::NegInfinity => NaiveDateTime::MIN,
            postgres::types::Timestamp::Value(t) => t,
        }
    }
}

impl<'r> Produce<'r, Option<NaiveDateTime>> for PostgresBinarySourcePartitionParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<NaiveDateTime> {
        let (ridx, cidx) = self.next_loc()?;
        let row = &self.rowbuf[ridx];
        let val = row.try_get(cidx)?;
        match val {
            Some(postgres::types::Timestamp::PosInfinity) => Some(NaiveDateTime::MAX),
            Some(postgres::types::Timestamp::NegInfinity) => Some(NaiveDateTime::MIN),
            Some(postgres::types::Timestamp::Value(t)) => t,
            None => None,
        }
    }
}

impl<'r> Produce<'r, DateTime<Utc>> for PostgresBinarySourcePartitionParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> DateTime<Utc> {
        let (ridx, cidx) = self.next_loc()?;
        let row = &self.rowbuf[ridx];
        let val = row.try_get(cidx)?;
        match val {
            postgres::types::Timestamp::PosInfinity => DateTime::<Utc>::MAX_UTC,
            postgres::types::Timestamp::NegInfinity => DateTime::<Utc>::MIN_UTC,
            postgres::types::Timestamp::Value(t) => t,
        }
    }
}

impl<'r> Produce<'r, Option<DateTime<Utc>>> for PostgresBinarySourcePartitionParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<DateTime<Utc>> {
        let (ridx, cidx) = self.next_loc()?;
        let row = &self.rowbuf[ridx];
        let val = row.try_get(cidx)?;
        match val {
            Some(postgres::types::Timestamp::PosInfinity) => Some(DateTime::<Utc>::MAX_UTC),
            Some(postgres::types::Timestamp::NegInfinity) => Some(DateTime::<Utc>::MIN_UTC),
            Some(postgres::types::Timestamp::Value(t)) => t,
            None => None,
        }
    }
}

impl<'r> Produce<'r, NaiveDate> for PostgresBinarySourcePartitionParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> NaiveDate {
        let (ridx, cidx) = self.next_loc()?;
        let row = &self.rowbuf[ridx];
        let val = row.try_get(cidx)?;
        match val {
            postgres::types::Date::PosInfinity => NaiveDate::MAX,
            postgres::types::Date::NegInfinity => NaiveDate::MIN,
            postgres::types::Date::Value(t) => t,
        }
    }
}

impl<'r> Produce<'r, Option<NaiveDate>> for PostgresBinarySourcePartitionParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<NaiveDate> {
        let (ridx, cidx) = self.next_loc()?;
        let row = &self.rowbuf[ridx];
        let val = row.try_get(cidx)?;
        match val {
            Some(postgres::types::Date::PosInfinity) => Some(NaiveDate::MAX),
            Some(postgres::types::Date::NegInfinity) => Some(NaiveDate::MIN),
            Some(postgres::types::Date::Value(t)) => t,
            None => None,
        }
    }
}

impl Produce<'_, HashMap<String, Option<String>>> for PostgresBinarySourcePartitionParser<'_> {
    type Error = PostgresSourceError;
    #[throws(PostgresSourceError)]
    fn produce(&mut self) -> HashMap<String, Option<String>> {
        unimplemented!("Please use `cursor` protocol for hstore type");
    }
}

impl Produce<'_, Option<HashMap<String, Option<String>>>>
    for PostgresBinarySourcePartitionParser<'_>
{
    type Error = PostgresSourceError;
    #[throws(PostgresSourceError)]
    fn produce(&mut self) -> Option<HashMap<String, Option<String>>> {
        unimplemented!("Please use `cursor` protocol for hstore type");
    }
}

pub struct PostgresCSVSourceParser<'a> {
    iter: StringRecordsIntoIter<CopyOutReader<'a>>,
    rowbuf: Vec<StringRecord>,
    ncols: usize,
    current_col: usize,
    current_row: usize,
    is_finished: bool,
}

impl<'a> PostgresCSVSourceParser<'a> {
    pub fn new(
        iter: StringRecordsIntoIter<CopyOutReader<'a>>,
        schema: &[PostgresTypeSystem],
    ) -> Self {
        Self {
            iter,
            rowbuf: Vec::with_capacity(DB_BUFFER_SIZE),
            ncols: schema.len(),
            current_row: 0,
            current_col: 0,
            is_finished: false,
        }
    }

    #[throws(PostgresSourceError)]
    fn next_loc(&mut self) -> (usize, usize) {
        let ret = (self.current_row, self.current_col);
        self.current_row += (self.current_col + 1) / self.ncols;
        self.current_col = (self.current_col + 1) % self.ncols;
        ret
    }
}

impl<'a> PartitionParser<'a> for PostgresCSVSourceParser<'a> {
    type Error = PostgresSourceError;
    type TypeSystem = PostgresTypeSystem;

    #[throws(PostgresSourceError)]
    fn fetch_next(&mut self) -> (usize, bool) {
        assert!(self.current_col == 0);
        let remaining_rows = self.rowbuf.len() - self.current_row;
        if remaining_rows > 0 {
            return (remaining_rows, self.is_finished);
        } else if self.is_finished {
            return (0, self.is_finished);
        }

        if !self.rowbuf.is_empty() {
            self.rowbuf.drain(..);
        }
        for _ in 0..DB_BUFFER_SIZE {
            if let Some(row) = self.iter.next() {
                self.rowbuf.push(row?);
            } else {
                self.is_finished = true;
                break;
            }
        }
        self.current_row = 0;
        self.current_col = 0;
        (self.rowbuf.len(), self.is_finished)
    }
}

macro_rules! impl_csv_produce {
    ($($t: ty,)+) => {
        $(
            impl<'r, 'a> Produce<'r, $t> for PostgresCSVSourceParser<'a> {
                type Error = PostgresSourceError;

                #[throws(PostgresSourceError)]
                fn produce(&'r mut self) -> $t {
                    let (ridx, cidx) = self.next_loc()?;
                    self.rowbuf[ridx][cidx].parse().map_err(|_| {
                        ConnectorXError::cannot_produce::<$t>(Some(self.rowbuf[ridx][cidx].into()))
                    })?
                }
            }

            impl<'r, 'a> Produce<'r, Option<$t>> for PostgresCSVSourceParser<'a> {
                type Error = PostgresSourceError;

                #[throws(PostgresSourceError)]
                fn produce(&'r mut self) -> Option<$t> {
                    let (ridx, cidx) = self.next_loc()?;
                    match &self.rowbuf[ridx][cidx][..] {
                        "" => None,
                        v => Some(v.parse().map_err(|_| {
                            ConnectorXError::cannot_produce::<$t>(Some(self.rowbuf[ridx][cidx].into()))
                        })?),
                    }
                }
            }
        )+
    };
}

impl_csv_produce!(i8, i16, i32, i64, u32, f32, f64, Uuid, IpInet,);

macro_rules! impl_csv_vec_produce {
    ($($t: ty,)+) => {
        $(
            impl<'r, 'a> Produce<'r, Vec<Option<$t>>> for PostgresCSVSourceParser<'a> {
                type Error = PostgresSourceError;

                #[throws(PostgresSourceError)]
                fn produce(&mut self) -> Vec<Option<$t>> {
                    let (ridx, cidx) = self.next_loc()?;
                    let s = &self.rowbuf[ridx][cidx][..];
                    match s {
                        "{}" => vec![],
                        _ if s.len() < 3 => throw!(ConnectorXError::cannot_produce::<$t>(Some(s.into()))),
                        s => s[1..s.len() - 1]
                            .split(",")
                            .map(|v| {
                                if v == "NULL" {
                                    Ok(None)
                                } else {
                                    match v.parse() {
                                        Ok(v) => Ok(Some(v)),
                                        Err(e) => Err(e).map_err(|_| ConnectorXError::cannot_produce::<$t>(Some(s.into())))
                                    }
                                }
                            })
                            .collect::<Result<Vec<Option<$t>>, ConnectorXError>>()?,
                    }
                }
            }

            impl<'r, 'a> Produce<'r, Option<Vec<Option<$t>>>> for PostgresCSVSourceParser<'a> {
                type Error = PostgresSourceError;

                #[throws(PostgresSourceError)]
                fn produce(&mut self) -> Option<Vec<Option<$t>>> {
                    let (ridx, cidx) = self.next_loc()?;
                    let s = &self.rowbuf[ridx][cidx][..];
                    match s {
                        "" => None,
                        "{}" => Some(vec![]),
                        _ if s.len() < 3 => throw!(ConnectorXError::cannot_produce::<$t>(Some(s.into()))),
                        s => Some(
                            s[1..s.len() - 1]
                                .split(",")
                                .map(|v| {
                                    if v == "NULL" {
                                        Ok(None)
                                    } else {
                                        match v.parse() {
                                            Ok(v) => Ok(Some(v)),
                                            Err(e) => Err(e).map_err(|_| ConnectorXError::cannot_produce::<$t>(Some(s.into())))
                                        }
                                    }
                                })
                                .collect::<Result<Vec<Option<$t>>, ConnectorXError>>()?,
                        ),
                    }
                }
            }
        )+
    };
}

impl_csv_vec_produce!(i8, i16, i32, i64, f32, f64, Decimal, String,);

impl Produce<'_, bool> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&mut self) -> bool {
        let (ridx, cidx) = self.next_loc()?;
        let ret = match &self.rowbuf[ridx][cidx][..] {
            "t" => true,
            "f" => false,
            _ => throw!(ConnectorXError::cannot_produce::<bool>(Some(
                self.rowbuf[ridx][cidx].into()
            ))),
        };
        ret
    }
}

impl Produce<'_, Option<bool>> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&mut self) -> Option<bool> {
        let (ridx, cidx) = self.next_loc()?;
        let ret = match &self.rowbuf[ridx][cidx][..] {
            "" => None,
            "t" => Some(true),
            "f" => Some(false),
            _ => throw!(ConnectorXError::cannot_produce::<bool>(Some(
                self.rowbuf[ridx][cidx].into()
            ))),
        };
        ret
    }
}

impl Produce<'_, Vec<Option<bool>>> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&mut self) -> Vec<Option<bool>> {
        let (ridx, cidx) = self.next_loc()?;
        let s = &self.rowbuf[ridx][cidx][..];
        match s {
            "{}" => vec![],
            _ if s.len() < 3 => throw!(ConnectorXError::cannot_produce::<bool>(Some(s.into()))),
            s => s[1..s.len() - 1]
                .split(',')
                .map(|v| match v {
                    "NULL" => Ok(None),
                    "t" => Ok(Some(true)),
                    "f" => Ok(Some(false)),
                    _ => throw!(ConnectorXError::cannot_produce::<bool>(Some(s.into()))),
                })
                .collect::<Result<Vec<Option<bool>>, ConnectorXError>>()?,
        }
    }
}

impl Produce<'_, Option<Vec<Option<bool>>>> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&mut self) -> Option<Vec<Option<bool>>> {
        let (ridx, cidx) = self.next_loc()?;
        let s = &self.rowbuf[ridx][cidx][..];
        match s {
            "" => None,
            "{}" => Some(vec![]),
            _ if s.len() < 3 => throw!(ConnectorXError::cannot_produce::<bool>(Some(s.into()))),
            s => Some(
                s[1..s.len() - 1]
                    .split(',')
                    .map(|v| match v {
                        "NULL" => Ok(None),
                        "t" => Ok(Some(true)),
                        "f" => Ok(Some(false)),
                        _ => throw!(ConnectorXError::cannot_produce::<bool>(Some(s.into()))),
                    })
                    .collect::<Result<Vec<Option<bool>>, ConnectorXError>>()?,
            ),
        }
    }
}

impl<'r> Produce<'r, Decimal> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Decimal {
        let (ridx, cidx) = self.next_loc()?;
        match &self.rowbuf[ridx][cidx][..] {
            "Infinity" => Decimal::MAX,
            "-Infinity" => Decimal::MIN,
            v => v
                .parse()
                .map_err(|_| ConnectorXError::cannot_produce::<Decimal>(Some(v.into())))?,
        }
    }
}

impl<'r> Produce<'r, Option<Decimal>> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<Decimal> {
        let (ridx, cidx) = self.next_loc()?;
        match &self.rowbuf[ridx][cidx][..] {
            "" => None,
            "Infinity" => Some(Decimal::MAX),
            "-Infinity" => Some(Decimal::MIN),
            v => Some(
                v.parse()
                    .map_err(|_| ConnectorXError::cannot_produce::<Decimal>(Some(v.into())))?,
            ),
        }
    }
}

impl Produce<'_, DateTime<Utc>> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&mut self) -> DateTime<Utc> {
        let (ridx, cidx) = self.next_loc()?;
        match &self.rowbuf[ridx][cidx][..] {
            "infinity" => DateTime::<Utc>::MAX_UTC,
            "-infinity" => DateTime::<Utc>::MIN_UTC,
            // postgres csv return example: 1970-01-01 00:00:01+00
            v => format!("{}:00", v)
                .parse()
                .map_err(|_| ConnectorXError::cannot_produce::<DateTime<Utc>>(Some(v.into())))?,
        }
    }
}

impl Produce<'_, Option<DateTime<Utc>>> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&mut self) -> Option<DateTime<Utc>> {
        let (ridx, cidx) = self.next_loc()?;
        match &self.rowbuf[ridx][cidx][..] {
            "" => None,
            "infinity" => Some(DateTime::<Utc>::MAX_UTC),
            "-infinity" => Some(DateTime::<Utc>::MIN_UTC),
            v => {
                // postgres csv return example: 1970-01-01 00:00:01+00
                Some(format!("{}:00", v).parse().map_err(|_| {
                    ConnectorXError::cannot_produce::<DateTime<Utc>>(Some(v.into()))
                })?)
            }
        }
    }
}

impl Produce<'_, NaiveDate> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&mut self) -> NaiveDate {
        let (ridx, cidx) = self.next_loc()?;
        match &self.rowbuf[ridx][cidx][..] {
            "infinity" => NaiveDate::MAX,
            "-infinity" => NaiveDate::MIN,
            v => NaiveDate::parse_from_str(v, "%Y-%m-%d")
                .map_err(|_| ConnectorXError::cannot_produce::<NaiveDate>(Some(v.into())))?,
        }
    }
}

impl Produce<'_, Option<NaiveDate>> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&mut self) -> Option<NaiveDate> {
        let (ridx, cidx) = self.next_loc()?;
        match &self.rowbuf[ridx][cidx][..] {
            "" => None,
            "infinity" => Some(NaiveDate::MAX),
            "-infinity" => Some(NaiveDate::MIN),
            v => Some(
                NaiveDate::parse_from_str(v, "%Y-%m-%d")
                    .map_err(|_| ConnectorXError::cannot_produce::<NaiveDate>(Some(v.into())))?,
            ),
        }
    }
}

impl Produce<'_, NaiveDateTime> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&mut self) -> NaiveDateTime {
        let (ridx, cidx) = self.next_loc()?;
        match &self.rowbuf[ridx][cidx] {
            "infinity" => NaiveDateTime::MAX,
            "-infinity" => NaiveDateTime::MIN,
            v => NaiveDateTime::parse_from_str(v, "%Y-%m-%d %H:%M:%S%.f")
                .map_err(|_| ConnectorXError::cannot_produce::<NaiveDateTime>(Some(v.into())))?,
        }
    }
}

impl Produce<'_, Option<NaiveDateTime>> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&mut self) -> Option<NaiveDateTime> {
        let (ridx, cidx) = self.next_loc()?;
        match &self.rowbuf[ridx][cidx][..] {
            "" => None,
            "infinity" => Some(NaiveDateTime::MAX),
            "-infinity" => Some(NaiveDateTime::MIN),
            v => Some(
                NaiveDateTime::parse_from_str(v, "%Y-%m-%d %H:%M:%S%.f").map_err(|_| {
                    ConnectorXError::cannot_produce::<NaiveDateTime>(Some(v.into()))
                })?,
            ),
        }
    }
}

impl Produce<'_, NaiveTime> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&mut self) -> NaiveTime {
        let (ridx, cidx) = self.next_loc()?;
        NaiveTime::parse_from_str(&self.rowbuf[ridx][cidx], "%H:%M:%S%.f").map_err(|_| {
            ConnectorXError::cannot_produce::<NaiveTime>(Some(self.rowbuf[ridx][cidx].into()))
        })?
    }
}

impl Produce<'_, Option<NaiveTime>> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&mut self) -> Option<NaiveTime> {
        let (ridx, cidx) = self.next_loc()?;
        match &self.rowbuf[ridx][cidx][..] {
            "" => None,
            v => Some(
                NaiveTime::parse_from_str(v, "%H:%M:%S%.f")
                    .map_err(|_| ConnectorXError::cannot_produce::<NaiveTime>(Some(v.into())))?,
            ),
        }
    }
}

impl<'r> Produce<'r, &'r str> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> &'r str {
        let (ridx, cidx) = self.next_loc()?;
        &self.rowbuf[ridx][cidx]
    }
}

impl<'r> Produce<'r, Option<&'r str>> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<&'r str> {
        let (ridx, cidx) = self.next_loc()?;
        match &self.rowbuf[ridx][cidx][..] {
            "" => None,
            v => Some(v),
        }
    }
}

impl<'r> Produce<'r, Vec<u8>> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Vec<u8> {
        let (ridx, cidx) = self.next_loc()?;
        decode(&self.rowbuf[ridx][cidx][2..])? // escape \x in the beginning
    }
}

impl<'r> Produce<'r, Option<Vec<u8>>> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<Vec<u8>> {
        let (ridx, cidx) = self.next_loc()?;
        match &self.rowbuf[ridx][cidx] {
            // escape \x in the beginning, empty if None
            "" => None,
            v => Some(decode(&v[2..])?),
        }
    }
}

impl<'r> Produce<'r, Value> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Value {
        let (ridx, cidx) = self.next_loc()?;
        let v = &self.rowbuf[ridx][cidx];
        from_str(v).map_err(|_| ConnectorXError::cannot_produce::<Value>(Some(v.into())))?
    }
}

impl<'r> Produce<'r, Option<Value>> for PostgresCSVSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<Value> {
        let (ridx, cidx) = self.next_loc()?;

        match &self.rowbuf[ridx][cidx][..] {
            "" => None,
            v => {
                from_str(v).map_err(|_| ConnectorXError::cannot_produce::<Value>(Some(v.into())))?
            }
        }
    }
}

pub struct PostgresRawSourceParser<'a> {
    iter: RowIter<'a>,
    rowbuf: Vec<Row>,
    ncols: usize,
    current_col: usize,
    current_row: usize,
    is_finished: bool,
}

impl<'a> PostgresRawSourceParser<'a> {
    pub fn new(iter: RowIter<'a>, schema: &[PostgresTypeSystem]) -> Self {
        Self {
            iter,
            rowbuf: Vec::with_capacity(DB_BUFFER_SIZE),
            ncols: schema.len(),
            current_row: 0,
            current_col: 0,
            is_finished: false,
        }
    }

    #[throws(PostgresSourceError)]
    fn next_loc(&mut self) -> (usize, usize) {
        let ret = (self.current_row, self.current_col);
        self.current_row += (self.current_col + 1) / self.ncols;
        self.current_col = (self.current_col + 1) % self.ncols;
        ret
    }
}

impl<'a> PartitionParser<'a> for PostgresRawSourceParser<'a> {
    type TypeSystem = PostgresTypeSystem;
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn fetch_next(&mut self) -> (usize, bool) {
        assert!(self.current_col == 0);
        let remaining_rows = self.rowbuf.len() - self.current_row;
        if remaining_rows > 0 {
            return (remaining_rows, self.is_finished);
        } else if self.is_finished {
            return (0, self.is_finished);
        }

        if !self.rowbuf.is_empty() {
            self.rowbuf.drain(..);
        }
        for _ in 0..DB_BUFFER_SIZE {
            if let Some(row) = self.iter.next()? {
                self.rowbuf.push(row);
            } else {
                self.is_finished = true;
                break;
            }
        }
        self.current_row = 0;
        self.current_col = 0;
        (self.rowbuf.len(), self.is_finished)
    }
}

macro_rules! impl_produce {
    ($($t: ty,)+) => {
        $(
            impl<'r, 'a> Produce<'r, $t> for PostgresRawSourceParser<'a> {
                type Error = PostgresSourceError;

                #[throws(PostgresSourceError)]
                fn produce(&'r mut self) -> $t {
                    let (ridx, cidx) = self.next_loc()?;
                    let row = &self.rowbuf[ridx];
                    let val = row.try_get(cidx)?;
                    val
                }
            }

            impl<'r, 'a> Produce<'r, Option<$t>> for PostgresRawSourceParser<'a> {
                type Error = PostgresSourceError;

                #[throws(PostgresSourceError)]
                fn produce(&'r mut self) -> Option<$t> {
                    let (ridx, cidx) = self.next_loc()?;
                    let row = &self.rowbuf[ridx];
                    let val = row.try_get(cidx)?;
                    val
                }
            }
        )+
    };
}

impl_produce!(
    i8,
    i16,
    i32,
    i64,
    u32,
    f32,
    f64,
    Decimal,
    bool,
    &'r str,
    Vec<u8>,
    NaiveTime,
    Uuid,
    Value,
    IpInet,
    Vector,
    HalfVector,
    Bit,
    SparseVector,
    HashMap<String, Option<String>>,
    Vec<Option<bool>>,
    Vec<Option<String>>,
    Vec<Option<i16>>,
    Vec<Option<i32>>,
    Vec<Option<i64>>,
    Vec<Option<f32>>,
    Vec<Option<f64>>,
    Vec<Option<Decimal>>,
);

impl<'r> Produce<'r, DateTime<Utc>> for PostgresRawSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> DateTime<Utc> {
        let (ridx, cidx) = self.next_loc()?;
        let row = &self.rowbuf[ridx];
        let val: postgres::types::Timestamp<DateTime<Utc>> = row.try_get(cidx)?;
        match val {
            postgres::types::Timestamp::PosInfinity => DateTime::<Utc>::MAX_UTC,
            postgres::types::Timestamp::NegInfinity => DateTime::<Utc>::MIN_UTC,
            postgres::types::Timestamp::Value(t) => t,
        }
    }
}

impl<'r> Produce<'r, Option<DateTime<Utc>>> for PostgresRawSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<DateTime<Utc>> {
        let (ridx, cidx) = self.next_loc()?;
        let row = &self.rowbuf[ridx];
        let val = row.try_get(cidx)?;
        match val {
            Some(postgres::types::Timestamp::PosInfinity) => Some(DateTime::<Utc>::MAX_UTC),
            Some(postgres::types::Timestamp::NegInfinity) => Some(DateTime::<Utc>::MIN_UTC),
            Some(postgres::types::Timestamp::Value(t)) => t,
            None => None,
        }
    }
}

impl<'r> Produce<'r, NaiveDateTime> for PostgresRawSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> NaiveDateTime {
        let (ridx, cidx) = self.next_loc()?;
        let row = &self.rowbuf[ridx];
        let val: postgres::types::Timestamp<NaiveDateTime> = row.try_get(cidx)?;
        match val {
            postgres::types::Timestamp::PosInfinity => NaiveDateTime::MAX,
            postgres::types::Timestamp::NegInfinity => NaiveDateTime::MIN,
            postgres::types::Timestamp::Value(t) => t,
        }
    }
}

impl<'r> Produce<'r, Option<NaiveDateTime>> for PostgresRawSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<NaiveDateTime> {
        let (ridx, cidx) = self.next_loc()?;
        let row = &self.rowbuf[ridx];
        let val = row.try_get(cidx)?;
        match val {
            Some(postgres::types::Timestamp::PosInfinity) => Some(NaiveDateTime::MAX),
            Some(postgres::types::Timestamp::NegInfinity) => Some(NaiveDateTime::MIN),
            Some(postgres::types::Timestamp::Value(t)) => t,
            None => None,
        }
    }
}

impl<'r> Produce<'r, NaiveDate> for PostgresRawSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> NaiveDate {
        let (ridx, cidx) = self.next_loc()?;
        let row = &self.rowbuf[ridx];
        let val: postgres::types::Date<NaiveDate> = row.try_get(cidx)?;
        match val {
            postgres::types::Date::PosInfinity => NaiveDate::MAX,
            postgres::types::Date::NegInfinity => NaiveDate::MIN,
            postgres::types::Date::Value(t) => t,
        }
    }
}

impl<'r> Produce<'r, Option<NaiveDate>> for PostgresRawSourceParser<'_> {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<NaiveDate> {
        let (ridx, cidx) = self.next_loc()?;
        let row = &self.rowbuf[ridx];
        let val = row.try_get(cidx)?;
        match val {
            Some(postgres::types::Date::PosInfinity) => Some(NaiveDate::MAX),
            Some(postgres::types::Date::NegInfinity) => Some(NaiveDate::MIN),
            Some(postgres::types::Date::Value(t)) => t,
            None => None,
        }
    }
}

impl<C> SourcePartition for PostgresSourcePartition<SimpleProtocol, C>
where
    C: MakeTlsConnect<Socket> + Clone + 'static + Sync + Send,
    C::TlsConnect: Send,
    C::Stream: Send,
    <C::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
    type TypeSystem = PostgresTypeSystem;
    type Parser<'a> = PostgresSimpleSourceParser;
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn result_rows(&mut self) {
        self.nrows = get_total_rows(&mut self.conn, &self.query)?;
    }

    #[throws(PostgresSourceError)]
    fn parser(&mut self) -> Self::Parser<'_> {
        let rows = self.conn.simple_query(self.query.as_str())?; // unless reading the data, it seems like issue the query is fast
        PostgresSimpleSourceParser::new(rows, &self.schema)
    }

    fn nrows(&self) -> usize {
        self.nrows
    }

    fn ncols(&self) -> usize {
        self.ncols
    }
}

pub struct PostgresSimpleSourceParser {
    rows: Vec<SimpleQueryMessage>,
    ncols: usize,
    current_col: usize,
    current_row: usize,
}
impl PostgresSimpleSourceParser {
    pub fn new(rows: Vec<SimpleQueryMessage>, schema: &[PostgresTypeSystem]) -> Self {
        Self {
            rows,
            ncols: schema.len(),
            current_row: 0,
            current_col: 0,
        }
    }

    #[throws(PostgresSourceError)]
    fn next_loc(&mut self) -> (usize, usize) {
        let ret = (self.current_row, self.current_col);
        self.current_row += (self.current_col + 1) / self.ncols;
        self.current_col = (self.current_col + 1) % self.ncols;
        ret
    }
}

impl PartitionParser<'_> for PostgresSimpleSourceParser {
    type TypeSystem = PostgresTypeSystem;
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn fetch_next(&mut self) -> (usize, bool) {
        self.current_row = 0;
        self.current_col = 0;
        if !self.rows.is_empty() {
            if let SimpleQueryMessage::RowDescription(_) = &self.rows[0] {
                self.current_row = 1;
            }
        }

        (self.rows.len() - 1 - self.current_row, true) // last message is command complete
    }
}

macro_rules! impl_simple_produce {
    ($($t: ty,)+) => {
        $(
            impl<'r> Produce<'r, $t> for PostgresSimpleSourceParser {
                type Error = PostgresSourceError;

                #[throws(PostgresSourceError)]
                fn produce(&'r mut self) -> $t {
                    let (ridx, cidx) = self.next_loc()?;
                    let val = match &self.rows[ridx] {
                        SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                            Some(s) => s
                                .parse()
                                .map_err(|_| ConnectorXError::cannot_produce::<$t>(Some(s.into())))?,
                            None => throw!(anyhow!(
                                "Cannot parse NULL in NOT NULL column."
                            )),
                        },
                        SimpleQueryMessage::CommandComplete(c) => {
                            panic!("get command: {}", c);
                        }
                        _ => {
                            panic!("what?");
                        }
                    };
                    val
                }
            }

            impl<'r, 'a> Produce<'r, Option<$t>> for PostgresSimpleSourceParser {
                type Error = PostgresSourceError;

                #[throws(PostgresSourceError)]
                fn produce(&'r mut self) -> Option<$t> {
                    let (ridx, cidx) = self.next_loc()?;
                    let val = match &self.rows[ridx] {
                        SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                            Some(s) => Some(
                                s.parse()
                                    .map_err(|_| ConnectorXError::cannot_produce::<$t>(Some(s.into())))?,
                            ),
                            None => None,
                        },
                        SimpleQueryMessage::CommandComplete(c) => {
                            panic!("get command: {}", c);
                        }
                        _ => {
                            panic!("what?");
                        }
                    };
                    val
                }
            }
        )+
    };
}

impl_simple_produce!(i8, i16, i32, i64, u32, f32, f64, Uuid, IpInet,);

impl<'r> Produce<'r, bool> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> bool {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                Some(s) => match s {
                    "t" => true,
                    "f" => false,
                    _ => throw!(ConnectorXError::cannot_produce::<bool>(Some(s.into()))),
                },
                None => throw!(anyhow!("Cannot parse NULL in non-NULL column.")),
            },
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}

impl<'r> Produce<'r, Option<bool>> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<bool> {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                Some(s) => match s {
                    "t" => Some(true),
                    "f" => Some(false),
                    _ => throw!(ConnectorXError::cannot_produce::<bool>(Some(s.into()))),
                },
                None => None,
            },
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}

impl<'r> Produce<'r, Decimal> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Decimal {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                Some("Infinity") => Decimal::MAX,
                Some("-Infinity") => Decimal::MIN,
                Some(s) => s
                    .parse()
                    .map_err(|_| ConnectorXError::cannot_produce::<Decimal>(Some(s.into())))?,
                None => throw!(anyhow!("Cannot parse NULL in NOT NULL column.")),
            },
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}

impl<'r> Produce<'r, Option<Decimal>> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<Decimal> {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                Some("Infinity") => Some(Decimal::MAX),
                Some("-Infinity") => Some(Decimal::MIN),
                Some(s) => Some(
                    s.parse()
                        .map_err(|_| ConnectorXError::cannot_produce::<Decimal>(Some(s.into())))?,
                ),
                None => None,
            },
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}

impl<'r> Produce<'r, &'r str> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> &'r str {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                Some(s) => s,
                None => throw!(anyhow!("Cannot parse NULL in non-NULL column.")),
            },
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}

impl<'r> Produce<'r, Option<&'r str>> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<&'r str> {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => row.try_get(cidx)?,
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}

impl<'r> Produce<'r, Vec<u8>> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Vec<u8> {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                Some(s) => {
                    let mut res = s.chars();
                    res.next();
                    res.next();
                    decode(
                        res.enumerate()
                            .fold(String::new(), |acc, (_i, c)| format!("{}{}", acc, c))
                            .chars()
                            .map(|c| c as u8)
                            .collect::<Vec<u8>>(),
                    )?
                }
                None => throw!(anyhow!("Cannot parse NULL in non-NULL column.")),
            },
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}

impl<'r> Produce<'r, Option<Vec<u8>>> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<Vec<u8>> {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                Some(s) => {
                    let mut res = s.chars();
                    res.next();
                    res.next();
                    Some(decode(
                        res.enumerate()
                            .fold(String::new(), |acc, (_i, c)| format!("{}{}", acc, c))
                            .chars()
                            .map(|c| c as u8)
                            .collect::<Vec<u8>>(),
                    )?)
                }
                None => None,
            },
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}

fn rem_first_and_last(value: &str) -> &str {
    let mut chars = value.chars();
    chars.next();
    chars.next_back();
    chars.as_str()
}

macro_rules! impl_simple_vec_produce {
    ($($t: ty,)+) => {
        $(
            impl<'r> Produce<'r, Vec<Option<$t>>> for PostgresSimpleSourceParser {
                type Error = PostgresSourceError;

                #[throws(PostgresSourceError)]
                fn produce(&'r mut self) -> Vec<Option<$t>> {
                    let (ridx, cidx) = self.next_loc()?;
                    let val = match &self.rows[ridx] {
                        SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                            Some(s) => match s{
                                "" => throw!(anyhow!("Cannot parse NULL in non-NULL column.")),
                                "{}" => vec![],
                                _ => rem_first_and_last(s).split(",").map(|v| {
                                    if v == "NULL" {
                                        Ok(None)
                                    } else {
                                        match v.parse() {
                                            Ok(v) => Ok(Some(v)),
                                            Err(e) => Err(e).map_err(|_| ConnectorXError::cannot_produce::<Vec<$t>>(Some(s.into())))
                                        }
                                    }
                                }).collect::<Result<Vec<Option<$t>>, ConnectorXError>>()?
                            },
                            None => throw!(anyhow!("Cannot parse NULL in non-NULL column.")),
                        },
                        SimpleQueryMessage::CommandComplete(c) => {
                            panic!("get command: {}", c);
                        }
                        _ => {
                            panic!("what?");
                        }
                    };
                    val
                }
            }

            impl<'r, 'a> Produce<'r, Option<Vec<Option<$t>>>> for PostgresSimpleSourceParser {
                type Error = PostgresSourceError;

                #[throws(PostgresSourceError)]
                fn produce(&'r mut self) -> Option<Vec<Option<$t>>> {
                    let (ridx, cidx) = self.next_loc()?;
                    let val = match &self.rows[ridx] {

                        SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                            Some(s) => match s{
                                "" => None,
                                "{}" => Some(vec![]),
                                _ => Some(rem_first_and_last(s).split(",").map(|v| {
                                    if v == "NULL" {
                                        Ok(None)
                                    } else {
                                        match v.parse() {
                                            Ok(v) => Ok(Some(v)),
                                            Err(e) => Err(e).map_err(|_| ConnectorXError::cannot_produce::<Vec<$t>>(Some(s.into())))
                                        }
                                    }
                                }).collect::<Result<Vec<Option<$t>>, ConnectorXError>>()?)
                            },
                            None => None,
                        },

                        SimpleQueryMessage::CommandComplete(c) => {
                            panic!("get command: {}", c);
                        }
                        _ => {
                            panic!("what?");
                        }
                    };
                    val
                }
            }
        )+
    };
}
impl_simple_vec_produce!(i16, i32, i64, f32, f64, Decimal, String,);

impl<'r> Produce<'r, Vec<Option<bool>>> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Vec<Option<bool>> {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                Some(s) => match s {
                    "" => throw!(anyhow!("Cannot parse NULL in non-NULL column.")),
                    "{}" => vec![],
                    _ => rem_first_and_last(s)
                        .split(',')
                        .map(|token| match token {
                            "NULL" => Ok(None),
                            "t" => Ok(Some(true)),
                            "f" => Ok(Some(false)),
                            _ => {
                                throw!(ConnectorXError::cannot_produce::<Vec<bool>>(Some(s.into())))
                            }
                        })
                        .collect::<Result<Vec<Option<bool>>, ConnectorXError>>()?,
                },
                None => throw!(anyhow!("Cannot parse NULL in non-NULL column.")),
            },
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}

impl<'r> Produce<'r, Option<Vec<Option<bool>>>> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<Vec<Option<bool>>> {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                Some(s) => match s {
                    "" => None,
                    "{}" => Some(vec![]),
                    _ => Some(
                        rem_first_and_last(s)
                            .split(',')
                            .map(|token| match token {
                                "NULL" => Ok(None),
                                "t" => Ok(Some(true)),
                                "f" => Ok(Some(false)),
                                _ => {
                                    throw!(ConnectorXError::cannot_produce::<Vec<bool>>(Some(
                                        s.into()
                                    )))
                                }
                            })
                            .collect::<Result<Vec<Option<bool>>, ConnectorXError>>()?,
                    ),
                },
                None => None,
            },
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}

impl<'r> Produce<'r, NaiveDate> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> NaiveDate {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                Some(s) => match s {
                    "infinity" => NaiveDate::MAX,
                    "-infinity" => NaiveDate::MIN,
                    s => NaiveDate::parse_from_str(s, "%Y-%m-%d").map_err(|_| {
                        ConnectorXError::cannot_produce::<NaiveDate>(Some(s.into()))
                    })?,
                },
                None => throw!(anyhow!("Cannot parse NULL in non-NULL column.")),
            },
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}

impl<'r> Produce<'r, Option<NaiveDate>> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<NaiveDate> {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                Some(s) => match s {
                    "infinity" => Some(NaiveDate::MAX),
                    "-infinity" => Some(NaiveDate::MIN),
                    s => Some(NaiveDate::parse_from_str(s, "%Y-%m-%d").map_err(|_| {
                        ConnectorXError::cannot_produce::<Option<NaiveDate>>(Some(s.into()))
                    })?),
                },
                None => None,
            },
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}

impl<'r> Produce<'r, NaiveTime> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> NaiveTime {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                Some(s) => NaiveTime::parse_from_str(s, "%H:%M:%S%.f")
                    .map_err(|_| ConnectorXError::cannot_produce::<NaiveTime>(Some(s.into())))?,
                None => throw!(anyhow!("Cannot parse NULL in non-NULL column.")),
            },
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}

impl<'r> Produce<'r, Option<NaiveTime>> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<NaiveTime> {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                Some(s) => Some(NaiveTime::parse_from_str(s, "%H:%M:%S%.f").map_err(|_| {
                    ConnectorXError::cannot_produce::<Option<NaiveTime>>(Some(s.into()))
                })?),
                None => None,
            },
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}

impl<'r> Produce<'r, NaiveDateTime> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> NaiveDateTime {
        let (ridx, cidx) = self.next_loc()?;
        let val =
            match &self.rows[ridx] {
                SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                    Some(s) => match s {
                        "infinity" => NaiveDateTime::MAX,
                        "-infinity" => NaiveDateTime::MIN,
                        s => NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f").map_err(
                            |_| ConnectorXError::cannot_produce::<NaiveDateTime>(Some(s.into())),
                        )?,
                    },
                    None => throw!(anyhow!("Cannot parse NULL in non-NULL column.")),
                },
                SimpleQueryMessage::CommandComplete(c) => {
                    panic!("get command: {}", c);
                }
                _ => {
                    panic!("what?");
                }
            };
        val
    }
}

impl<'r> Produce<'r, Option<NaiveDateTime>> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<NaiveDateTime> {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                Some(s) => match s {
                    "infinity" => Some(NaiveDateTime::MAX),
                    "-infinity" => Some(NaiveDateTime::MIN),
                    s => Some(
                        NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f").map_err(|_| {
                            ConnectorXError::cannot_produce::<Option<NaiveDateTime>>(Some(s.into()))
                        })?,
                    ),
                },
                None => None,
            },
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}

impl<'r> Produce<'r, DateTime<Utc>> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> DateTime<Utc> {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                Some("infinity") => DateTime::<Utc>::MAX_UTC,
                Some("-infinity") => DateTime::<Utc>::MIN_UTC,
                Some(s) => {
                    let time_string = format!("{}:00", s).to_owned();
                    let slice: &str = &time_string[..];
                    let time: DateTime<FixedOffset> =
                        DateTime::parse_from_str(slice, "%Y-%m-%d %H:%M:%S%.f%:z").unwrap();

                    time.with_timezone(&Utc)
                }
                None => throw!(anyhow!("Cannot parse NULL in non-NULL column.")),
            },
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}

impl<'r> Produce<'r, Option<DateTime<Utc>>> for PostgresSimpleSourceParser {
    type Error = PostgresSourceError;

    #[throws(PostgresSourceError)]
    fn produce(&'r mut self) -> Option<DateTime<Utc>> {
        let (ridx, cidx) = self.next_loc()?;
        let val = match &self.rows[ridx] {
            SimpleQueryMessage::Row(row) => match row.try_get(cidx)? {
                Some("infinity") => Some(DateTime::<Utc>::MAX_UTC),
                Some("-infinity") => Some(DateTime::<Utc>::MIN_UTC),
                Some(s) => {
                    let time_string = format!("{}:00", s).to_owned();
                    let slice: &str = &time_string[..];
                    let time: DateTime<FixedOffset> =
                        DateTime::parse_from_str(slice, "%Y-%m-%d %H:%M:%S%.f%:z").unwrap();

                    Some(time.with_timezone(&Utc))
                }
                None => None,
            },
            SimpleQueryMessage::CommandComplete(c) => {
                panic!("get command: {}", c);
            }
            _ => {
                panic!("what?");
            }
        };
        val
    }
}