Parse & write Postgres array datatypes (#83)

SeaQL · Oct 21, 2022 · 9c91fd4 · 9c91fd4
1 parent 0753c75
commit 9c91fd4
Show file tree

Hide file tree

Showing 4 changed files with 192 additions and 255 deletions.
diff --git a/src/postgres/def/types.rs b/src/postgres/def/types.rs
@@ -1,3 +1,4 @@
+use sea_query::SeaRc;
 #[cfg(feature = "with-serde")]
 use serde::{Deserialize, Serialize};
 
@@ -110,7 +111,7 @@ pub enum Type {
     JsonBinary,
 
     /// Variable-length multidimensional array
-    Array,
+    Array(SeaRc<Type>),
 
     // TODO:
     // /// The structure of a row or record; a list of field names and types
@@ -148,61 +149,66 @@ impl Type {
     // TODO: Support more types
     #[allow(clippy::should_implement_trait)]
     pub fn from_str(name: &str) -> Type {
-        match name.to_lowercase().as_str() {
-            "smallint" | "int2" => Type::SmallInt,
-            "integer" | "int" | "int4" => Type::Integer,
-            "bigint" | "int8" => Type::BigInt,
-            "decimal" => Type::Decimal(ArbitraryPrecisionNumericAttr::default()),
-            "numeric" => Type::Numeric(ArbitraryPrecisionNumericAttr::default()),
-            "real" | "float4" => Type::Real,
-            "double precision" | "double" | "float8" => Type::DoublePrecision,
-            "smallserial" | "serial2" => Type::SmallSerial,
-            "serial" | "serial4" => Type::Serial,
-            "bigserial" | "serial8" => Type::BigSerial,
-            "money" => Type::Money,
-            "character varying" | "varchar" => Type::Varchar(StringAttr::default()),
-            "character" | "char" => Type::Char(StringAttr::default()),
-            "text" => Type::Text,
-            "bytea" => Type::Bytea,
-            "timestamp" | "timestamp without time zone" => Type::Timestamp(TimeAttr::default()),
-            "timestamp with time zone" => Type::TimestampWithTimeZone(TimeAttr::default()),
-            "date" => Type::Date,
-            "time" | "time without time zone" => Type::Time(TimeAttr::default()),
-            "time with time zone" => Type::TimeWithTimeZone(TimeAttr::default()),
-            "interval" => Type::Interval(IntervalAttr::default()),
-            "boolean" => Type::Boolean,
-            "point" => Type::Point,
-            "line" => Type::Line,
-            "lseg" => Type::Lseg,
-            "box" => Type::Box,
-            "path" => Type::Path,
-            "polygon" => Type::Polygon,
-            "circle" => Type::Circle,
-            "cidr" => Type::Cidr,
-            "inet" => Type::Inet,
-            "macaddr" => Type::MacAddr,
-            "macaddr8" => Type::MacAddr8,
-            "bit" => Type::Bit(BitAttr::default()),
-            "tsvector" => Type::TsVector,
-            "tsquery" => Type::TsQuery,
-            "uuid" => Type::Uuid,
-            "xml" => Type::Xml,
-            "json" => Type::Json,
-            "jsonb" => Type::JsonBinary,
-            "array" => Type::Array,
-            // "" => Type::Composite,
-            "int4range" => Type::Int4Range,
-            "int8range" => Type::Int8Range,
-            "numrange" => Type::NumRange,
-            "tsrange" => Type::TsRange,
-            "tstzrange" => Type::TsTzRange,
-            "daterange" => Type::DateRange,
-            // "" => Type::Domain,
-            "pg_lsn" => Type::PgLsn,
-            "user-defined" => Type::Enum(EnumDef::default()),
-
-            _ => Type::Unknown(name.to_owned()),
+        fn parse_type(name: &str) -> Type {
+            match name.to_lowercase().as_str() {
+                "smallint" | "int2" => Type::SmallInt,
+                "integer" | "int" | "int4" => Type::Integer,
+                "bigint" | "int8" => Type::BigInt,
+                "decimal" => Type::Decimal(ArbitraryPrecisionNumericAttr::default()),
+                "numeric" => Type::Numeric(ArbitraryPrecisionNumericAttr::default()),
+                "real" | "float4" => Type::Real,
+                "double precision" | "double" | "float8" => Type::DoublePrecision,
+                "smallserial" | "serial2" => Type::SmallSerial,
+                "serial" | "serial4" => Type::Serial,
+                "bigserial" | "serial8" => Type::BigSerial,
+                "money" => Type::Money,
+                "character varying" | "varchar" => Type::Varchar(StringAttr::default()),
+                "character" | "char" => Type::Char(StringAttr::default()),
+                "text" => Type::Text,
+                "bytea" => Type::Bytea,
+                "timestamp" | "timestamp without time zone" => Type::Timestamp(TimeAttr::default()),
+                "timestamp with time zone" => Type::TimestampWithTimeZone(TimeAttr::default()),
+                "date" => Type::Date,
+                "time" | "time without time zone" => Type::Time(TimeAttr::default()),
+                "time with time zone" => Type::TimeWithTimeZone(TimeAttr::default()),
+                "interval" => Type::Interval(IntervalAttr::default()),
+                "boolean" => Type::Boolean,
+                "point" => Type::Point,
+                "line" => Type::Line,
+                "lseg" => Type::Lseg,
+                "box" => Type::Box,
+                "path" => Type::Path,
+                "polygon" => Type::Polygon,
+                "circle" => Type::Circle,
+                "cidr" => Type::Cidr,
+                "inet" => Type::Inet,
+                "macaddr" => Type::MacAddr,
+                "macaddr8" => Type::MacAddr8,
+                "bit" => Type::Bit(BitAttr::default()),
+                "tsvector" => Type::TsVector,
+                "tsquery" => Type::TsQuery,
+                "uuid" => Type::Uuid,
+                "xml" => Type::Xml,
+                "json" => Type::Json,
+                "jsonb" => Type::JsonBinary,
+                // "" => Type::Composite,
+                "int4range" => Type::Int4Range,
+                "int8range" => Type::Int8Range,
+                "numrange" => Type::NumRange,
+                "tsrange" => Type::TsRange,
+                "tstzrange" => Type::TsTzRange,
+                "daterange" => Type::DateRange,
+                // "" => Type::Domain,
+                "pg_lsn" => Type::PgLsn,
+                "user-defined" => Type::Enum(EnumDef::default()),
+                _ if name.ends_with("[]") => {
+                    let col_type = parse_type(&name.replacen("[]", "", 1));
+                    Type::Array(SeaRc::new(col_type))
+                }
+                _ => Type::Unknown(name.to_owned()),
+            }
         }
+        parse_type(name)
     }
 }
 

diff --git a/src/postgres/query/column.rs b/src/postgres/query/column.rs
@@ -1,6 +1,6 @@
 use super::{InformationSchema, SchemaQueryBuilder};
 use crate::sqlx_types::postgres::PgRow;
-use sea_query::{Expr, Iden, Query, SeaRc, SelectStatement};
+use sea_query::{Alias, BinOper, Expr, Iden, Query, SeaRc, SelectStatement};
 
 #[derive(Debug, sea_query::Iden)]
 /// Ref: https://www.postgresql.org/docs/13/infoschema-columns.html
@@ -75,9 +75,12 @@ impl SchemaQueryBuilder {
         table: SeaRc<dyn Iden>,
     ) -> SelectStatement {
         Query::select()
-            .columns(vec![
-                ColumnsField::ColumnName,
-                ColumnsField::DataType,
+            .column(ColumnsField::ColumnName)
+            .expr(
+                Expr::expr(Expr::cust("udt_name::regtype").cast_as(Alias::new("text")))
+                    .binary(BinOper::As, Expr::col(ColumnsField::DataType)),
+            )
+            .columns([
                 ColumnsField::ColumnDefault,
                 ColumnsField::GenerationExpression,
                 ColumnsField::IsNullable,