feat: New Functionalities

Cargo.toml

@@ -16,3 +16,4 @@ tokio = { version = "1", features = ["full"] }
 dotenv = "0.15.0"
 testcontainers = { version ="0.15.0" }
 testcontainers-modules = { version = "0.3.5", features = ["postgres"] }
+once_cell = "1.19.0"

src/db_queries.rs

@@ -1,6 +1,6 @@
 use sqlx::PgPool;
 
-use crate::models::TableColumn;
+use crate::models::{TableColumn, UserDefinedEnums};
 
 pub async fn get_table_columns(
     pool: &PgPool,
@@ -98,3 +98,28 @@ ORDER BY
         .await?;
     Ok(rows)
 }
+
+pub async fn get_user_defined_enums(
+    udt_names: &Vec<String>,
+    pool: &PgPool,
+) -> sqlx::Result<Vec<UserDefinedEnums>> {
+    let query = "
+        SELECT
+            t.typname AS enum_name,
+            e.enumlabel AS enum_value
+        FROM
+            pg_type t
+            JOIN pg_enum e ON t.oid = e.enumtypid
+        WHERE
+            t.typname = ANY($1)
+        ORDER BY
+            t.typname,
+            e.enumsortorder;
+            ";
+
+    let rows = sqlx::query_as::<_, UserDefinedEnums>(query)
+        .bind(udt_names)
+        .fetch_all(pool)
+        .await?;
+    Ok(rows)
+}

src/generate.rs

@@ -3,19 +3,24 @@ use sqlx::PgPool;
 use std::fs;
 use std::path::Path;
 
-use crate::db_queries::get_table_columns;
+use crate::db_queries::{get_table_columns, get_user_defined_enums};
 use crate::models::TableColumn;
-use crate::utils::{generate_struct_code, to_pascal_case, to_snake_case};
+use crate::utils::{generate_enum_code, generate_struct_code, to_pascal_case, to_snake_case};
 
 use crate::query_generate::generate_query_code;
+use crate::utils::{DateTimeLib, SqlGenState};
+use crate::STATE;
 
 pub async fn generate(
+    enable_serde: bool,
     output_folder: &str,
     database_url: &str,
     context: Option<&str>,
     force: bool,
     include_tables: Option<Vec<&str>>,
+    exclude_tables: Vec<String>,
     schemas: Option<Vec<&str>>,
+    date_time_lib: DateTimeLib,
 ) -> Result<(), Box<dyn std::error::Error>> {
     // Connect to the Postgres database
     let pool = PgPoolOptions::new()
@@ -28,18 +33,56 @@ pub async fn generate(
 
     let default_schema = "public";
     let rows = get_table_columns(&pool, schemas.unwrap_or(vec![default_schema]), None).await?;
+    let user_defined = rows
+        .iter()
+        .filter_map(|e| {
+            if e.data_type.as_str() == "USER-DEFINED" && e.udt_name.as_str() != "geometry" {
+                Some(e.udt_name.clone())
+            } else {
+                None
+            }
+        })
+        .collect::<Vec<String>>();
 
+    let enum_rows = get_user_defined_enums(&user_defined, &pool).await?;
+    let mut unique_enums = std::collections::BTreeSet::new();
+    for row in &enum_rows {
+        unique_enums.insert(row.enum_name.clone());
+    }
+    let enums = unique_enums.into_iter().collect::<Vec<String>>();
     // Create the output folder if it doesn't exist
     fs::create_dir_all(output_folder)?;
 
     let mut unique = std::collections::BTreeSet::new();
     for row in &rows {
         unique.insert(row.table_name.clone());
     }
-    let tables: Vec<String> = unique.into_iter().collect::<Vec<String>>();
-
+    let tables: Vec<String> = unique
+        .into_iter()
+        .collect::<Vec<String>>()
+        .into_iter()
+        .filter(|e| !exclude_tables.contains(e))
+        .collect();
+
+    if !enums.is_empty() {
+        println!("Outputting user defined enums: {:?}", enums);
+    }
     println!("Outputting tables: {:?}", tables);
 
+    STATE
+        .set(SqlGenState {
+            user_defined: enums.clone(),
+            date_time_lib,
+        })
+        .expect("Unable to set state");
+
+    let mut rs_enums = Vec::new();
+
+    for user_enum in enums {
+        let enum_code = generate_enum_code(&user_enum, &enum_rows, enable_serde);
+        rs_enums.push(enum_code);
+    }
+
     // Generate structs and queries for each table
     for table in &tables {
         if let Some(ts) = include_tables.clone() {
@@ -48,7 +91,7 @@ pub async fn generate(
             }
         }
         // Generate the struct code based on the row
-        let struct_code = generate_struct_code(&table, &rows);
+        let struct_code = generate_struct_code(&table, &rows, enable_serde);
 
         // Generate the query code based on the row
         let query_code = generate_query_code(&table, &rows);
@@ -75,18 +118,28 @@ pub async fn generate(
         }
     }
 
-    let context_code = generate_db_context(context.unwrap_or(&database_name), &tables, &rows);
+    let context_code =
+        generate_db_context(context.unwrap_or(&database_name), &rs_enums, &tables, &rows);
     let context_file_path = format!("{}/mod.rs", output_folder);
     fs::write(context_file_path, context_code)?;
     Ok(())
 }
 
-fn generate_db_context(database_name: &str, tables: &[String], _rows: &[TableColumn]) -> String {
+fn generate_db_context(
+    database_name: &str,
+    enums: &[String],
+    tables: &[String],
+    _rows: &[TableColumn],
+) -> String {
     let mut db_context_code = String::new();
 
     db_context_code.push_str("#![allow(dead_code)]\n");
     db_context_code
         .push_str("// Generated with sql-gen\n//https://github.com/jayy-lmao/sql-gen\n\n");
+    for enum_item in enums {
+        db_context_code.push_str(enum_item);
+        db_context_code.push_str("\n\n");
+    }
     for table in tables {
         db_context_code.push_str(&format!("pub mod {};\n", to_snake_case(table)));
         db_context_code.push_str(&format!(

src/main.rs

@@ -1,4 +1,6 @@
 use clap::{App, Arg, SubCommand};
+use once_cell::sync::OnceCell;
+use utils::{DateTimeLib, SqlGenState};
 
 mod db_queries;
 mod generate;
@@ -7,6 +9,8 @@ mod models;
 mod query_generate;
 mod utils;
 
+pub static STATE: OnceCell<SqlGenState> = OnceCell::new();
+
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     dotenv::dotenv().ok();
@@ -22,6 +26,14 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 .help("Sets the output folder for generated structs")
                 .takes_value(true),
         )
+        .arg(
+            Arg::with_name("serde")
+                .long("serde")
+                .default_value("true")
+                .value_name("SQLGEN_ENABLE_SERDE")
+                .help("Adds Serde derices to created structs")
+                .takes_value(false),
+        )
         .arg(
             Arg::with_name("migrations")
                 .short('m')
@@ -68,6 +80,16 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 .use_delimiter(true)
                 .help("Specify the table name(s)"),
         )
+        .arg(
+            Arg::with_name("exclude")
+                .short('e')
+                .long("exclude")
+                .takes_value(true)
+                .value_name("SQLGEN_EXCLUDE")
+                .multiple(true)
+                .use_delimiter(true)
+                .help("Specify the excluded table name(s)"),
+        )
         .arg(
             Arg::new("force")
                 .short('f')
@@ -76,7 +98,16 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 .takes_value(false)
                 .required(false)
                 .help("Overwrites existing files sharing names in that folder"),
-        );
+        )
+        .arg(
+                Arg::with_name("datetime-lib")
+                    .long("datetime-lib")
+                    .default_value("chrono")
+                    .possible_values(&["chrono", "time"])
+                    .value_name("SQLGEN_DATETIME_LIB")
+                    .help("Specifies the library to use for date and time handling")
+                    .takes_value(true),
+            );
 
     let migrate_subcommand = SubCommand::with_name("migrate")
         .about("Generate SQL migrations based on struct differences")
@@ -195,7 +226,40 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         let schemas: Option<Vec<&str>> =
             matches.values_of("schema").map(|schemas| schemas.collect());
         let force = matches.is_present("force");
-        generate::generate(output_folder, database_url, context, force, None, schemas).await?;
+        let include_tables = matches.values_of("table").map(|v| v.collect::<Vec<&str>>());
+        let exclude_tables = matches
+            .values_of("exclude")
+            .map(|v| {
+                v.into_iter()
+                    .map(|e| e.to_string())
+                    .collect::<Vec<String>>()
+            })
+            .unwrap_or(vec![]);
+
+        if !exclude_tables.is_empty() {
+            println!("Excluding tables: {:?}", exclude_tables);
+        }
+
+        let enable_serde = matches.is_present("serde");
+
+        let date_time_lib = matches
+            .value_of("datetime-lib")
+            .map(|e| e.to_string())
+            .unwrap();
+        let date_time_lib = DateTimeLib::from(date_time_lib);
+
+        generate::generate(
+            enable_serde,
+            output_folder,
+            database_url,
+            context,
+            force,
+            include_tables,
+            exclude_tables,
+            schemas,
+            date_time_lib,
+        )
+        .await?;
     } else if let Some(matches) = matches.subcommand_matches("migrate") {
         let input_migrations_folder = matches.value_of("migrations").unwrap_or("./migrations");
         println!(

src/models.rs

@@ -1,4 +1,4 @@
-#[derive(sqlx::FromRow)]
+#[derive(sqlx::FromRow, Clone)]
 pub struct TableColumn {
     pub(crate) table_name: String,
     pub(crate) column_name: String,
@@ -12,3 +12,9 @@ pub struct TableColumn {
     // #todo
     pub(crate) table_schema: String,
 }
+
+#[derive(sqlx::FromRow, Clone)]
+pub struct UserDefinedEnums {
+    pub(crate) enum_name: String,
+    pub(crate) enum_value: String,
+}