postgres: write bit and varbit column types

SeaQL · Jan 31, 2024 · 7cdadb2 · 7cdadb2
1 parent 22a4ae8
commit 7cdadb2
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 14 deletions.
diff --git a/src/postgres/def/types.rs b/src/postgres/def/types.rs
@@ -91,6 +91,9 @@ pub enum Type {
     /// Fixed length bit string
     Bit(BitAttr),
 
+    /// Variable length bit string
+    VarBit(BitAttr),
+
     // Text search types
     /// A sorted list of distinct lexemes which are words that have been normalized to merge different
     /// variants of the same word
@@ -171,7 +174,7 @@ impl Type {
             "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,
+            "boolean" | "bool" => Type::Boolean,
             "point" => Type::Point,
             "line" => Type::Line,
             "lseg" => Type::Lseg,
@@ -184,6 +187,7 @@ impl Type {
             "macaddr" => Type::MacAddr,
             "macaddr8" => Type::MacAddr8,
             "bit" => Type::Bit(BitAttr::default()),
+            "bit varying" | "varbit" => Type::VarBit(BitAttr::default()),
             "tsvector" => Type::TsVector,
             "tsquery" => Type::TsQuery,
             "uuid" => Type::Uuid,
@@ -288,7 +292,7 @@ impl Type {
     }
 
     pub fn has_bit_attr(&self) -> bool {
-        matches!(self, Type::Bit(_))
+        matches!(self, Type::Bit(_) | Type::VarBit(_))
     }
 
     pub fn has_enum_attr(&self) -> bool {

diff --git a/src/postgres/parser/column.rs b/src/postgres/parser/column.rs
@@ -176,6 +176,15 @@ pub fn parse_bit_attributes(
                 },
             };
         }
+        Type::VarBit(ref mut attr) => {
+            attr.length = match character_maximum_length {
+                None => None,
+                Some(num) => match u16::try_from(num) {
+                    Ok(num) => Some(num),
+                    Err(_) => None,
+                },
+            };
+        }
         _ => panic!("parse_bit_attributes(_) received a type that does not have BitAttr"),
     };
 

diff --git a/src/postgres/writer/column.rs b/src/postgres/writer/column.rs
@@ -106,18 +106,8 @@ impl ColumnInfo {
                 Type::Inet => ColumnType::Custom(Alias::new("inet").into_iden()),
                 Type::MacAddr => ColumnType::Custom(Alias::new("macaddr").into_iden()),
                 Type::MacAddr8 => ColumnType::Custom(Alias::new("macaddr8").into_iden()),
-                Type::Bit(bit_attr) => {
-                    let mut str = String::new();
-                    write!(str, "bit").unwrap();
-                    if bit_attr.length.is_some() {
-                        write!(str, "(").unwrap();
-                        if let Some(length) = bit_attr.length {
-                            write!(str, "{}", length).unwrap();
-                        }
-                        write!(str, ")").unwrap();
-                    }
-                    ColumnType::Custom(Alias::new(&str).into_iden())
-                }
+                Type::Bit(bit_attr) => ColumnType::Bit(bit_attr.length.map(Into::into)),
+                Type::VarBit(bit_attr) => ColumnType::VarBit(bit_attr.length.unwrap_or(1).into()),
                 Type::TsVector => ColumnType::Custom(Alias::new("tsvector").into_iden()),
                 Type::TsQuery => ColumnType::Custom(Alias::new("tsquery").into_iden()),
                 Type::Uuid => ColumnType::Uuid,

diff --git a/tests/live/postgres/src/main.rs b/tests/live/postgres/src/main.rs
@@ -54,6 +54,7 @@ async fn main() {
         create_collection_table(),
         create_parent_table(),
         create_child_table(),
+        create_db_types_table(),
     ];
 
     for tbl_create_stmt in tbl_create_stmts.iter() {
@@ -424,3 +425,34 @@ fn create_child_table() -> TableCreateStatement {
         )
         .to_owned()
 }
+
+fn create_db_types_table() -> TableCreateStatement {
+    Table::create()
+        .table(Alias::new("db_types"))
+        .col(
+            ColumnDef::new(Alias::new("id"))
+                .integer()
+                .not_null()
+                .auto_increment(),
+        )
+        .col(ColumnDef::new(Alias::new("binary_1")).binary())
+        .col(ColumnDef::new(Alias::new("binary_2")).binary_len(1))
+        .col(ColumnDef::new(Alias::new("binary_3")).binary_len(16))
+        .col(ColumnDef::new(Alias::new("var_binary_1")).var_binary(1))
+        .col(ColumnDef::new(Alias::new("var_binary_2")).var_binary(16))
+        .col(ColumnDef::new(Alias::new("var_binary_3")).var_binary(32))
+        .col(ColumnDef::new(Alias::new("bit_1")).bit(Some(1)))
+        .col(ColumnDef::new(Alias::new("bit_2")).bit(Some(16)))
+        .col(ColumnDef::new(Alias::new("bit_3")).bit(Some(32)))
+        .col(ColumnDef::new(Alias::new("var_bit_1")).varbit(1))
+        .col(ColumnDef::new(Alias::new("var_bit_2")).varbit(16))
+        .col(ColumnDef::new(Alias::new("var_bit_3")).varbit(32))
+        .col(ColumnDef::new(Alias::new("bool")).boolean())
+        .primary_key(
+            Index::create()
+                .primary()
+                .name("db_types_pkey")
+                .col(Alias::new("id")),
+        )
+        .to_owned()
+}