fix ltree binary copyout parsing

Author: wangxiaoying

connectorx-python/connectorx/tests/test_postgres.py

@@ -450,7 +450,7 @@ def test_postgres_with_index_col(postgres_url: str) -> None:
 
 
 def test_postgres_types_binary(postgres_url: str) -> None:
-    query = "SELECT test_date, test_timestamp, test_timestamptz, test_int16, test_int64, test_float32, test_numeric, test_bpchar, test_char, test_varchar, test_uuid, test_time, test_json, test_jsonb, test_bytea, test_enum, test_f4array, test_f8array, test_narray, test_i2array, test_i4array, test_i8array, test_citext FROM test_types"
+    query = "SELECT test_date, test_timestamp, test_timestamptz, test_int16, test_int64, test_float32, test_numeric, test_bpchar, test_char, test_varchar, test_uuid, test_time, test_json, test_jsonb, test_bytea, test_enum, test_f4array, test_f8array, test_narray, test_i2array, test_i4array, test_i8array, test_citext, test_ltree FROM test_types"
     df = read_sql(postgres_url, query)
     expected = pd.DataFrame(
         index=range(4),
@@ -549,7 +549,7 @@ def test_postgres_types_binary(postgres_url: str) -> None:
                 dtype="object",
             ),
             "test_citext": pd.Series(["str_citext", "", "s", None], dtype="object"),
-            # "test_ltree": pd.Series(["A.B.C.D", "A.B.E", "A", None], dtype="object"), # waiting for https://github.com/sfackler/rust-postgres/issues/960
+            "test_ltree": pd.Series(["A.B.C.D", "A.B.E", "A", None], dtype="object"),
         },
     )
     assert_frame_equal(df, expected, check_names=True)

connectorx/src/sources/postgres/mod.rs

@@ -6,7 +6,7 @@ mod typesystem;
 
 pub use self::errors::PostgresSourceError;
 pub use connection::rewrite_tls_args;
-pub use typesystem::PostgresTypeSystem;
+pub use typesystem::{PostgresTypePairs, PostgresTypeSystem};
 
 use crate::constants::DB_BUFFER_SIZE;
 use crate::{
@@ -88,6 +88,7 @@ where
     queries: Vec<CXQuery<String>>,
     names: Vec<String>,
     schema: Vec<PostgresTypeSystem>,
+    pg_schema: Vec<postgres::types::Type>,
     _protocol: PhantomData<P>,
 }
 
@@ -109,6 +110,7 @@ where
             queries: vec![],
             names: vec![],
             schema: vec![],
+            pg_schema: vec![],
             _protocol: PhantomData,
         }
     }
@@ -153,19 +155,23 @@ where
 
         let stmt = conn.prepare(first_query.as_str())?;
 
-        let (names, types) = stmt
+        let (names, pg_types): (Vec<String>, Vec<postgres::types::Type>) = stmt
             .columns()
             .iter()
-            .map(|col| {
-                (
-                    col.name().to_string(),
-                    PostgresTypeSystem::from(col.type_()),
-                )
-            })
+            .map(|col| (col.name().to_string(), col.type_().clone()))
             .unzip();
 
         self.names = names;
-        self.schema = types;
+        self.schema = pg_types
+            .iter()
+            .map(|t| PostgresTypeSystem::from(t))
+            .collect();
+        self.pg_schema = self
+            .schema
+            .iter()
+            .zip(pg_types.iter())
+            .map(|(t1, t2)| PostgresTypePairs(t2, t1).into())
+            .collect();
     }
 
     #[throws(PostgresSourceError)]
@@ -199,6 +205,7 @@ where
                 conn,
                 &query,
                 &self.schema,
+                &self.pg_schema,
             ));
         }
         ret
@@ -215,6 +222,7 @@ where
     conn: PgConn<C>,
     query: CXQuery<String>,
     schema: Vec<PostgresTypeSystem>,
+    pg_schema: Vec<postgres::types::Type>,
     nrows: usize,
     ncols: usize,
     _protocol: PhantomData<P>,
@@ -227,11 +235,17 @@ where
     C::Stream: Send,
     <C::TlsConnect as TlsConnect<Socket>>::Future: Send,
 {
-    pub fn new(conn: PgConn<C>, query: &CXQuery<String>, schema: &[PostgresTypeSystem]) -> Self {
+    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,
@@ -259,8 +273,7 @@ where
     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 pg_schema: Vec<_> = self.schema.iter().map(|&dt| dt.into()).collect();
-        let iter = BinaryCopyOutIter::new(reader, &pg_schema);
+        let iter = BinaryCopyOutIter::new(reader, &self.pg_schema);
 
         PostgresBinarySourcePartitionParser::new(iter, &self.schema)
     }

connectorx/src/sources/postgres/typesystem.rs

@@ -103,38 +103,16 @@ impl<'a> From<&'a Type> for PostgresTypeSystem {
     }
 }
 
-// Link PostgresDTypes back to the one defined by the postgres crate.
-impl<'a> From<PostgresTypeSystem> for Type {
-    fn from(ty: PostgresTypeSystem) -> Type {
+pub struct PostgresTypePairs<'a>(pub &'a Type, pub &'a PostgresTypeSystem);
+
+// Link (postgres::Type, connectorx::PostgresTypes) back to the one defiend by the postgres crate.
+impl<'a> From<PostgresTypePairs<'a>> for Type {
+    fn from(ty: PostgresTypePairs) -> Type {
         use PostgresTypeSystem::*;
-        match ty {
-            Int2(_) => Type::INT2,
-            Int4(_) => Type::INT4,
-            Int8(_) => Type::INT8,
-            Float4(_) => Type::FLOAT4,
-            Float8(_) => Type::FLOAT8,
-            Numeric(_) => Type::NUMERIC,
-            Int2Array(_) => Type::INT2_ARRAY,
-            Int4Array(_) => Type::INT4_ARRAY,
-            Int8Array(_) => Type::INT8_ARRAY,
-            Float4Array(_) => Type::FLOAT4_ARRAY,
-            Float8Array(_) => Type::FLOAT8_ARRAY,
-            NumericArray(_) => Type::NUMERIC_ARRAY,
-            Bool(_) => Type::BOOL,
-            Text(_) => Type::TEXT,
-            BpChar(_) => Type::BPCHAR,
-            VarChar(_) => Type::VARCHAR,
-            Char(_) => Type::CHAR,
-            ByteA(_) => Type::BYTEA,
-            Date(_) => Type::DATE,
-            Time(_) => Type::TIME,
-            Timestamp(_) => Type::TIMESTAMP,
-            TimestampTz(_) => Type::TIMESTAMPTZ,
-            UUID(_) => Type::UUID,
-            JSON(_) => Type::JSON,
-            JSONB(_) => Type::JSONB,
+        match ty.1 {
             Enum(_) => Type::TEXT,
             HSTORE(_) => Type::TEXT, // hstore is not supported in binary protocol (since no corresponding inner TYPE)
+            _ => ty.0.clone(),
         }
     }
 }

docs/databases/postgres.md

@@ -41,7 +41,7 @@ cx.read_sql(conn, query)                                        # read data from
 | JSON            | object                    |                                    |
 | JSONB           | object                    |                                    |
 | ENUM            | object                    | need to convert enum column to text manually (`::text`) when using `csv` and `cursor` protocol |
-| ltree           | object                    | binary protocol returns with a hex char prefix. Check https://github.com/sfu-db/connector-x/pull/382 and https://github.com/sfackler/rust-postgres/issues/960 for status |
+| ltree           | object                    |                                    |
 | INT2[]          | object                    | list of i64                        |
 | INT4[]          | object                    | list of i64                        |
 | INT8[]          | object                    | list of i64                        |