Generate enum types from Postgres and MySQL (#188)

* add enum in postgres

* hrmm

* chore: better message

* wip

* wip: working on enums

* postgres enums

* wip

* mysql handling

* works for mysql as wellg

* fix

* fix clippy

* clean up

* clippy fixes

* fmt

* fix

* hrm

* fix?

* fix

Author: JasonShin

.github/workflows/rust-test.yaml

@@ -51,6 +51,7 @@ jobs:
         env:
           MYSQL_VERSION: ${{ matrix.db.mysql }}
           PG_VERSION: ${{ matrix.db.postgres }}
+          MYSQL_MIGRATION_FILE: "${{ matrix.db.mysql == '5.6' && 'mysql_migration_5_6.sql' || 'mysql_migration.sql' }}"
 
       - uses: GuillaumeFalourd/wait-sleep-action@v1
         with:

docker-compose.yml

@@ -16,7 +16,7 @@ services:
     image: mysql:${MYSQL_VERSION:-8}
     restart: always
     volumes:
-      - ./playpen/db/mysql_migration.sql:/docker-entrypoint-initdb.d/mysql_migration.sql
+      - ./playpen/db/${MYSQL_MIGRATION_FILE:-mysql_migration.sql}:/docker-entrypoint-initdb.d/mysql_migration.sql
     ports:
       - 33306:3306
     environment:

playpen/db/mysql_migration.sql

@@ -65,6 +65,6 @@ CREATE TABLE random (
 
 	-- JSON types
 	json1 JSON
-	
+
 );
 

playpen/db/mysql_migration_5_6.sql

@@ -0,0 +1,66 @@
+CREATE TABLE tables (
+    id INTEGER NOT NULL AUTO_INCREMENT,
+    number INTEGER NOT NULL,
+    occupied BOOL NOT NULL DEFAULT FALSE,
+    PRIMARY KEY (id)
+);
+
+CREATE TABLE items (
+   id INTEGER NOT NULL,
+   food_type VARCHAR(30) NOT NULL,
+   time_takes_to_cook INTEGER NOT NULL,
+   table_id INTEGER NOT NULL,
+   points SMALLINT NOT NULL,
+   FOREIGN KEY (table_id) REFERENCES tables (id),
+   PRIMARY KEY (id)
+);
+
+INSERT INTO tables (number)
+VALUES
+    (1), (2), (3), (4), (5), (6), (7), (8), (9), (10);
+
+INSERT INTO items (id, food_type, time_takes_to_cook, table_id, points)
+VALUES
+  (1, 'korean', 10, 1, 2),
+  (2, 'chinese', 10, 1, 2),
+  (3, 'japanese', 10, 1, 2),
+  (4, 'italian', 10, 1, 2),
+  (5, 'french', 10, 1, 2);
+
+
+-- We can primarily use this table to check how a column in MySQL can be converted to a TsFieldType
+
+CREATE TABLE random (
+	-- numeric types
+    intz INT,
+    smallint1 SMALLINT,
+    tinyint1 TINYINT,
+	medium1 MEDIUMINT,
+	bigint1 BIGINT,
+	decimal1 DECIMAL(2, 2),
+	numeric1 NUMERIC(2, 2),
+	double_precision1 DOUBLE PRECISION(2, 2),
+	float1 FLOAT,
+	double1 DOUBLE,
+	bit1 BIT(2),
+	bool1 BOOL,
+	bool2 BOOLEAN,
+	
+	-- date and datetime types
+	date1 DATE,
+	datetime1 DATETIME,
+	timestamp1 TIMESTAMP,
+	year1 YEAR,
+	
+	-- string types
+	char1 CHAR,
+	varchar1 VARCHAR(20),
+	binary1 BINARY,
+	varbinary1 VARBINARY(2),
+	blob1 BLOB,
+	text1 TEXT,
+		-- ideally this one should be generated as a legit enum type
+	enum1 ENUM('x-small', 'small', 'medium', 'large', 'x-large'),
+	set1 SET('one', 'two')
+);
+

playpen/db/postgres_migration.sql

@@ -21,6 +21,8 @@ CREATE TABLE postgres.public.items (
     PRIMARY KEY (id)
 );
 
+CREATE TYPE sizes AS ENUM ('x-small', 'small', 'medium', 'large', 'x-large');
+
 -- A table of randomness, just to test various field types in PostgreSQL
 -- There is a pretty comprehensive list of data types available in Postgres
 -- found in https://www.geeksforgeeks.org/postgresql-data-types/ -> not the official Postgres doc
@@ -58,6 +60,8 @@ CREATE TABLE postgres.public.random (
 
    	-- UUID
    	uuid1 UUID,
+
+	enum1 sizes,
 
     -- Special data types
    	box1 BOX,

src/ts_generator/information_schema.rs

@@ -1,4 +1,5 @@
 use crate::common::errors::{DB_CONN_POOL_RETRIEVE_ERROR, DB_SCHEME_READ_ERROR};
+use crate::common::logger::*;
 use crate::core::connection::DBConn;
 use crate::core::mysql::pool::MySqlConnectionManager;
 use crate::core::postgres::pool::PostgresConnectionManager;
@@ -25,7 +26,12 @@ struct ColumnsQueryResultRow {
 }
 
 pub struct DBSchema {
+  // Holds cache details for table / columns of the target database
   tables_cache: HashMap<String, Fields>,
+  // Holds cache details for enums that exists in the target database
+  enums_cache: HashMap<String, HashMap<String, Vec<String>>>,
+  // A flag to track if we have already tried to fetch enum and cache it
+  has_cached_enums: bool,
 }
 
 impl Default for DBSchema {
@@ -38,6 +44,8 @@ impl DBSchema {
   pub fn new() -> DBSchema {
     DBSchema {
       tables_cache: HashMap::new(),
+      enums_cache: HashMap::new(),
+      has_cached_enums: false,
     }
   }
 
@@ -58,7 +66,7 @@ impl DBSchema {
 
     let result = match &conn {
       DBConn::MySQLPooledConn(conn) => Self::mysql_fetch_table(self, table_name, conn).await,
-      DBConn::PostgresConn(conn) => Self::postgres_fetch_table(self, table_name, conn).await,
+      DBConn::PostgresConn(conn) => Self::postgres_fetch_table(self, &"public".to_string(), table_name, conn).await,
     };
 
     if let Some(result) = &result {
@@ -70,6 +78,7 @@ impl DBSchema {
 
   async fn postgres_fetch_table(
     &self,
+    schema: &String,
     table_names: &Vec<&str>,
     conn: &Mutex<Pool<PostgresConnectionManager>>,
   ) -> Option<Fields> {
@@ -82,14 +91,24 @@ impl DBSchema {
     let query = format!(
       r"
         SELECT
-            COLUMN_NAME as column_name,
-            DATA_TYPE as data_type,
-            IS_NULLABLE as is_nulalble
-        FROM information_schema.COLUMNS
-        WHERE TABLE_SCHEMA = 'public'
-        AND TABLE_NAME IN ({})
+          COLUMN_NAME as column_name,
+          DATA_TYPE as data_type,
+          IS_NULLABLE as is_nulalble,
+          TABLE_NAME as table_name,
+          (
+            select string_agg(e.enumlabel, ',')
+          from pg_type t
+              join pg_enum e on t.oid = e.enumtypid
+              join pg_catalog.pg_namespace n ON n.oid = t.typnamespace
+          where n.nspname = '{}'
+          and t.typname = udt_name
+          group by n.nspname, t.typname
+          ) as enum_values
+      FROM information_schema.COLUMNS
+      WHERE TABLE_SCHEMA = '{}'
+      AND TABLE_NAME IN ({});
                 ",
-      table_names,
+      schema, schema, table_names,
     );
 
     let mut fields: HashMap<String, Field> = HashMap::new();
@@ -103,10 +122,27 @@ impl DBSchema {
         let field_name: String = row.get(0);
         let field_type: String = row.get(1);
         let is_nullable: String = row.get(2);
+        let table_name: String = row.get(3);
+        let enum_values: Option<Vec<String>> = row
+          .try_get(4)
+          .ok()
+          .map(|val: String| val.split(",").map(|x| x.to_string()).collect());
+
         let field = Field {
-          field_type: TsFieldType::get_ts_field_type_from_postgres_field_type(field_type.to_owned()),
+          field_type: TsFieldType::get_ts_field_type_from_postgres_field_type(
+            field_type.to_owned(),
+            field_name.to_owned(),
+            table_name,
+            enum_values,
+          ),
           is_nullable: is_nullable == "YES",
         };
+        if field.field_type == TsFieldType::Any {
+          let message = format!(
+            "The column {field_name} of type {field_type} will be translated any as it isn't supported by sqlx-ts"
+          );
+          info!(message);
+        }
         fields.insert(field_name.to_owned(), field);
       }
 
@@ -131,8 +167,22 @@ impl DBSchema {
         SELECT
             COLUMN_NAME as column_name,
             DATA_TYPE as data_type,
-            IS_NULLABLE as is_nulalble
-        FROM information_schema.COLUMNS
+            IS_NULLABLE as is_nulalble,
+            TABLE_NAME,
+            (
+              SELECT REPLACE(
+                  TRIM(TRAILING ')' FROM
+                  TRIM(LEADING '(' from
+                  TRIM(LEADING 'enum' FROM COLUMN_TYPE)))
+                , '\''
+                , ''
+              )
+              FROM information_schema.COLUMNS subcols
+              WHERE subcols.TABLE_SCHEMA = (SELECT DATABASE())
+                AND subcols.TABLE_NAME = C.TABLE_NAME
+                AND subcols.COLUMN_NAME = C.COLUMN_NAME
+            ) AS enums
+        FROM information_schema.COLUMNS C
         WHERE TABLE_SCHEMA = (SELECT DATABASE())
         AND TABLE_NAME IN ({})
                 ",
@@ -149,8 +199,22 @@ impl DBSchema {
         let field_name: String = row.clone().take(0).expect(DB_SCHEME_READ_ERROR);
         let field_type: String = row.clone().take(1).expect(DB_SCHEME_READ_ERROR);
         let is_nullable: String = row.clone().take(2).expect(DB_SCHEME_READ_ERROR);
+        let table_name: String = row.clone().take(3).expect(DB_SCHEME_READ_ERROR);
+
+        let enum_values: Option<Vec<String>> = if field_type == "enum" {
+          let enums: String = row.clone().take(4).expect(DB_SCHEME_READ_ERROR);
+          let enum_values: Vec<String> = enums.split(",").map(|x| x.to_string()).collect();
+          Some(enum_values)
+        } else {
+          None
+        };
         let field = Field {
-          field_type: TsFieldType::get_ts_field_type_from_mysql_field_type(field_type.to_owned()),
+          field_type: TsFieldType::get_ts_field_type_from_mysql_field_type(
+            field_type.to_owned(),
+            table_name.to_owned(),
+            field_name.to_owned(),
+            enum_values.to_owned(),
+          ),
           is_nullable: is_nullable == "YES",
         };
         fields.insert(field_name.to_owned(), field);

src/ts_generator/types/ts_query.rs

@@ -1,23 +1,24 @@
+use crate::common::logger::*;
 use color_eyre::eyre::Result;
 use convert_case::{Case, Casing};
 use regex::Regex;
 use std::collections::{BTreeMap, HashMap};
 use std::fmt::{self};
 
 use crate::common::lazy::CONFIG;
-use crate::common::logger::*;
 use crate::ts_generator::errors::TsGeneratorError;
 
 type Array2DContent = Vec<Vec<TsFieldType>>;
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub enum TsFieldType {
   String,
   Number,
   Boolean,
   Object,
   Date,
   Null,
+  Enum(Vec<String>),
   Any,
   Array2D(Array2DContent),
   Array(Box<TsFieldType>),
@@ -57,6 +58,11 @@ impl fmt::Display for TsFieldType {
 
         write!(f, "{result}")
       }
+      TsFieldType::Enum(values) => {
+        let enums: Vec<String> = values.iter().map(|x| format!("'{x}'")).collect();
+        let joined_enums = enums.join(" | ");
+        write!(f, "{joined_enums}")
+      }
     }
   }
 }
@@ -69,30 +75,53 @@ impl TsFieldType {
   /// @examples
   /// get_ts_field_type_from_postgres_field_type("integer") -> TsFieldType::Number
   /// get_ts_field_type_from_postgres_field_type("smallint") -> TsFieldType::Number
+  /// get_ts_field_type_from_postgres_field_type("character varying",  ",,")
   ///
-  pub fn get_ts_field_type_from_postgres_field_type(field_type: String) -> Self {
+  pub fn get_ts_field_type_from_postgres_field_type(
+    field_type: String,
+    field_name: String,
+    table_name: String,
+    enum_values: Option<Vec<String>>,
+  ) -> Self {
     match field_type.as_str() {
       "smallint" | "integer" | "real" | "double precision" | "numeric" => Self::Number,
       "character" | "character varying" | "bytea" | "uuid" | "text" => Self::String,
       "boolean" => Self::Boolean,
       "json" | "jsonb" => Self::Object,
-      "ARRAY" | "array" => {
-        info!(
-          "Currently we cannot figure out the type information for an array, the feature will be added in the future"
-        );
+      "ARRAY" | "array" => Self::Any,
+      "date" => Self::Date,
+      "USER-DEFINED" => {
+        if let Some(enum_values) = enum_values {
+          return Self::Enum(enum_values);
+        }
+        let warning_message = format!("Failed to find enum values for field {field_name} of table {table_name}");
+        warning!(warning_message);
         Self::Any
       }
-      "date" => Self::Date,
       _ => Self::Any,
     }
   }
 
-  pub fn get_ts_field_type_from_mysql_field_type(mysql_field_type: String) -> Self {
+  pub fn get_ts_field_type_from_mysql_field_type(
+    mysql_field_type: String,
+    table_name: String,
+    field_name: String,
+    enum_values: Option<Vec<String>>,
+  ) -> Self {
     match mysql_field_type.as_str() {
       "bigint" | "decimal" | "double" | "float" | "int" | "mediumint" | "smallint" | "year" => Self::Number,
       "binary" | "bit" | "blob" | "char" | "text" | "varbinary" | "varchar" => Self::String,
       "tinyint" => Self::Boolean,
       "date" | "datetime" | "timestamp" => Self::Date,
+      "enum" => {
+        if let Some(enum_values) = enum_values {
+          return Self::Enum(enum_values);
+        }
+
+        let warning_message = format!("Failed to find enum values for field {field_name} of table {table_name}");
+        warning!(warning_message);
+        Self::Any
+      }
       _ => Self::Any,
     }
   }