Infinite Postgres Bug

connectorx/src/sources/postgres/mod.rs

@@ -601,7 +601,7 @@ macro_rules! impl_csv_produce {
     };
 }
 
-impl_csv_produce!(i8, i16, i32, i64, f32, f64, Decimal, Uuid,);
+impl_csv_produce!(i8, i16, i32, i64, f32, f64, Uuid,);
 
 macro_rules! impl_csv_vec_produce {
     ($($t: ty,)+) => {
@@ -754,6 +754,40 @@ impl<'r, 'a> Produce<'r, Option<Vec<bool>>> for PostgresCSVSourceParser<'a> {
     }
 }
 
+
+impl<'r, 'a> Produce<'r, Decimal> for PostgresCSVSourceParser<'a> {
+    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, 'a> Produce<'r, Option<Decimal>> for PostgresCSVSourceParser<'a> {
+    type Error = PostgresSourceError;
+
+    #[throws(PostgresSourceError)]
+    fn produce(&'r mut self) -> Option<Decimal> {
+        let (ridx, cidx) = self.next_loc()?;
+        match &self.rowbuf[ridx][cidx][..] {
+            "" => None,
+            v => Some(v.parse().map_err(|_| {
+                ConnectorXError::cannot_produce::<Decimal>(Some(v.into()))
+            })?),
+        }
+    }
+}
+
+
 impl<'r, 'a> Produce<'r, DateTime<Utc>> for PostgresCSVSourceParser<'a> {
     type Error = PostgresSourceError;
 
@@ -792,9 +826,13 @@ impl<'r, 'a> Produce<'r, NaiveDate> for PostgresCSVSourceParser<'a> {
     #[throws(PostgresSourceError)]
     fn produce(&mut self) -> NaiveDate {
         let (ridx, cidx) = self.next_loc()?;
-        NaiveDate::parse_from_str(&self.rowbuf[ridx][cidx], "%Y-%m-%d").map_err(|_| {
-            ConnectorXError::cannot_produce::<NaiveDate>(Some(self.rowbuf[ridx][cidx].into()))
-        })?
+        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()))
+            })?
+        }
     }
 }
 
@@ -820,13 +858,17 @@ impl<'r, 'a> Produce<'r, NaiveDateTime> for PostgresCSVSourceParser<'a> {
     #[throws(PostgresSourceError)]
     fn produce(&mut self) -> NaiveDateTime {
         let (ridx, cidx) = self.next_loc()?;
-        NaiveDateTime::parse_from_str(&self.rowbuf[ridx][cidx], "%Y-%m-%d %H:%M:%S").map_err(
-            |_| {
-                ConnectorXError::cannot_produce::<NaiveDateTime>(Some(
-                    self.rowbuf[ridx][cidx].into(),
-                ))
-            },
-        )?
+        match &self.rowbuf[ridx][cidx] {
+            "infinity" => NaiveDateTime::MAX,
+            "-infinity" => NaiveDateTime::MIN,
+            v => NaiveDateTime::parse_from_str(v, "%Y-%m-%d %H:%M:%S").map_err(
+                |_| {
+                    ConnectorXError::cannot_produce::<NaiveDateTime>(Some(
+                        v.into(),
+                    ))
+                },
+            )?
+        }
     }
 }
 

connectorx/tests/test_postgres.rs

@@ -1,3 +1,5 @@
+use chrono::{NaiveDate, NaiveDateTime};
+use rust_decimal::Decimal;
 use arrow::{
     array::{BooleanArray, Float64Array, Int64Array, StringArray},
     record_batch::RecordBatch,
@@ -149,6 +151,55 @@ fn test_postgres() {
     verify_arrow_results(result);
 }
 
+#[test]
+fn test_csv_infinite_values() {
+
+    let _ = env_logger::builder().is_test(true).try_init();
+
+    let dburl = env::var("POSTGRES_URL").unwrap();
+    #[derive(Debug, PartialEq)]
+    struct Row(i32, NaiveDate, NaiveDateTime, Decimal);
+
+    let url = Url::parse(dburl.as_str()).unwrap();
+    let (config, _tls) = rewrite_tls_args(&url).unwrap();
+    let mut source = PostgresSource::<CSVProtocol, NoTls>::new(config, NoTls, 1).unwrap();
+    source.set_queries(&[CXQuery::naked("select * from test_infinite_values")]);
+    source.fetch_metadata().unwrap();
+
+    let mut partitions = source.partition().unwrap();
+    assert!(partitions.len() == 1);
+    let mut partition = partitions.remove(0);
+    partition.result_rows().expect("run query");
+
+    assert_eq!(2, partition.nrows());
+    assert_eq!(4, partition.ncols());
+
+    let mut parser = partition.parser().unwrap();
+
+    let mut rows: Vec<Row> = Vec::new();
+    loop {
+        let (n, is_last) = parser.fetch_next().unwrap();
+        for _i in 0..n {
+            rows.push(Row(
+                parser.produce().unwrap(),
+                parser.produce().unwrap(),
+                parser.produce().unwrap(),
+                parser.produce().unwrap(),
+            ));
+        }
+        if is_last {
+            break;
+        }
+    }
+    assert_eq!(
+        vec![
+            Row(1, NaiveDate::MAX, NaiveDateTime::MAX, Decimal::MAX),
+            Row(2, NaiveDate::MIN, NaiveDateTime::MIN, Decimal::MIN),
+        ],
+        rows
+    );
+}
+
 #[test]
 fn test_postgres_csv() {
     let _ = env_logger::builder().is_test(true).try_init();

scripts/postgres.sql

@@ -1,6 +1,7 @@
 DROP TABLE IF EXISTS test_table;
 DROP TABLE IF EXISTS test_str;
 DROP TABLE IF EXISTS test_types;
+DROP TABLE IF EXISTS test_infinite_values;
 DROP TYPE IF EXISTS happiness;
 DROP EXTENSION IF EXISTS citext;
 DROP EXTENSION IF EXISTS ltree;
@@ -20,6 +21,16 @@ INSERT INTO test_table VALUES (3, 7, 'b', 3, FALSE);
 INSERT INTO test_table VALUES (4, 9, 'c', 7.8, NULL);
 INSERT INTO test_table VALUES (1314, 2, NULL, -10, TRUE);
 
+CREATE TABLE IF NOT EXISTS test_infinite_values(
+    test_int INTEGER NOT NULL,
+    test_date DATE,
+	test_timestamp TIMESTAMP,
+	test_real REAL
+);
+
+INSERT INTO test_infinite_values VALUES (1, 'infinity'::DATE, 'infinity'::TIMESTAMP, 'infinity'::REAL);
+INSERT INTO test_infinite_values VALUES (2, '-infinity'::DATE, '-infinity'::TIMESTAMP, '-infinity'::REAL);
+
 CREATE TABLE IF NOT EXISTS test_str(
     id INTEGER NOT NULL,
     test_language TEXT,