update

Cargo.toml

@@ -9,6 +9,7 @@ color-eyre = "0.6.3"
 comment-parser = "0.1.0"
 figment = { version = "0.10.19", features = ["yaml"] }
 heck = "0.5.0"
+indicatif = "0.17.11"
 inquire = "0.7.5"
 prettytable = "0.10.0"
 quote = "1.0.40"

src/generate.rs

@@ -1,10 +1,10 @@
 use core::time;
 
-const HEADER: &str = r#"AUTOGENERATED DATA"#;
+const HEADER: &str = r#"== Schema Information"#;
 const COMMENTHEAD: &str = r#"/*"#;
 const COMMENTBODY: &str = r#" *"#;
 const COMMENTTAIL: &str = r#"*/"#;
-use crate::Config;
+use crate::{config::DateTimeCrate, Config};
 use color_eyre::{
     eyre::{eyre, ContextCompat, Report},
     Result,
@@ -12,7 +12,7 @@ use color_eyre::{
 use comment_parser::{CommentParser, Event};
 use prettytable::{format, row, Table};
 use sea_orm_codegen::OutputFile;
-use sea_schema::sea_query::{self, ColumnSpec, TableCreateStatement};
+use sea_schema::sea_query::{self, ColumnSpec, ColumnType, StringLen, TableCreateStatement};
 use tokio::{fs, task::JoinSet};
 use url::Url;
 
@@ -128,6 +128,7 @@ pub async fn get_tables(
             tracing::info!("Discovering schema");
             let schema_discovery = SchemaDiscovery::new(connection, schema);
             let schema = schema_discovery.discover().await?;
+            tracing::info!(?schema);
             let table_stmts = schema
                 .tables
                 .into_iter()
@@ -177,7 +178,7 @@ pub async fn generate_models(
     tables: Vec<TableCreateStatement>,
     config: Config,
 ) -> Result<Vec<OutputFile>> {
-    tracing::info!(?tables);
+    tracing::debug!(?tables);
     let output_path = config.output.path.clone();
     let files = tables
         .into_iter()
@@ -200,7 +201,7 @@ pub async fn generate_models(
                     None => return Err(eyre!("Table name not found")),
                 };
                 let table_str = generate_table(table, config.clone()).await?;
-                tracing::info!(?table_str);
+                tracing::debug!(?table_str);
                 let filename = format!("{}.rs", table_name);
                 let file_path = output_path.join(&filename);
                 let exists = file_path.exists();
@@ -211,11 +212,11 @@ pub async fn generate_models(
                     let mut found = false;
                     for comment in comments {
                         if let Event::BlockComment(a, b) = comment {
-                            tracing::info!(?a, ?b);
+                            tracing::debug!(?a, ?b);
                             if b.contains(HEADER) {
                                 found = true;
                                 file_content = file_content.replace(a, &table_str);
-                                tracing::info!("Found header");
+                                tracing::debug!("Found header");
                                 break;
                             }
                         }
@@ -223,10 +224,11 @@ pub async fn generate_models(
                     if found {
                         Ok::<String, Report>(file_content)
                     } else {
-                        todo!()
+                        let merged_content = format!("{}\n\n{}", table_str, file_content);
+                        Ok::<String, Report>(merged_content)
                     }
                 } else {
-                    todo!()
+                    Ok::<String, Report>(table_str)
                 }?;
 
                 Ok(OutputFile {
@@ -261,24 +263,35 @@ async fn generate_table(table: TableCreateStatement, config: Config) -> Result<S
         None => return Err(eyre!("Table name not found")),
     };
     let mut inner = String::new();
-    inner.push_str(format!("{COMMENTBODY}\n\n  Table name: {}\n\n", table_name).as_str());
+    inner.push_str(format!("{COMMENTBODY}\n Table name: {table_name}\n\n\n").as_str());
     inner = inner.strip_suffix('\n').unwrap().to_string();
     let mut ptable = Table::new();
     let format = format::FormatBuilder::new()
         .column_separator(' ')
         .borders(' ')
         .separators(
-            &[format::LinePosition::Top, format::LinePosition::Bottom],
+            &[
-            format::LineSeparator::new(' ', ' ', ' ', ' '),
+                format::LinePosition::Bottom,
+                format::LinePosition::Title,
+                // format::LinePosition::Top,
+            ],
+            format::LineSeparator::default(),
         )
         .padding(1, 1)
         .build();
     ptable.set_format(format);
+    ptable.set_titles(row!["Name", "Type", "RustType", "Attributes"]);
+    let indexes = table.get_indexes();
+    tracing::info!(?indexes);
     for column in table.get_columns() {
         let name = column.get_column_name();
         if let Some(column_type) = column.get_column_type() {
+            let column_type_rust =
+                type_to_rust_string(column_type, config.sea_orm.entity.date_time_crate.clone());
+            let column_type =
+                type_to_string(column_type, config.sea_orm.entity.date_time_crate.clone());
             let attrs = attrs_to_string(column.get_column_spec());
-            ptable.add_row(row![name, format!("{:?}", column_type), attrs]);
+            ptable.add_row(row![name, column_type, column_type_rust, attrs]);
         }
     }
     inner.push_str(ptable.to_string().as_str());
@@ -292,12 +305,41 @@ async fn generate_table(table: TableCreateStatement, config: Config) -> Result<S
     Ok(string)
 }
 
+fn type_to_string(column: &ColumnType, date_time_create: DateTimeCrate) -> String {
+    tracing::debug!(?column, ?date_time_create);
+    match column {
+        ColumnType::Char(Some(max)) => format!("Char({}max)", max),
+        ColumnType::Char(None) => "Char".to_owned(),
+        ColumnType::String(StringLen::None) => "String".to_owned(),
+        ColumnType::String(StringLen::Max) => "String(Max)".to_owned(),
+        ColumnType::String(StringLen::N(len)) => format!("String({}max)", len),
+        ColumnType::Text => "Text".to_owned(),
+        ColumnType::Integer => "Integer".to_owned(),
+
+        ColumnType::Custom(_) => "String".to_owned(),
+        _ => unimplemented!(),
+    }
+}
+fn type_to_rust_string(column: &ColumnType, date_time_create: DateTimeCrate) -> String {
+    tracing::debug!(?column, ?date_time_create);
+    match column {
+        ColumnType::Char(_) | ColumnType::String(_) | ColumnType::Text | ColumnType::Custom(_) => {
+            "String".to_owned()
+        }
+        ColumnType::Integer => "i32".to_owned(),
+        _ => unimplemented!(),
+    }
+}
+
 fn attrs_to_string(column: &Vec<ColumnSpec>) -> String {
+    tracing::debug!(?column);
     column
         .iter()
         .filter_map(|c| match c {
             ColumnSpec::PrimaryKey => Some("primary key"),
+            ColumnSpec::UniqueKey => Some("unique key"),
             ColumnSpec::AutoIncrement => Some("autoincrement"),
+            ColumnSpec::NotNull => Some("not null"),
             _ => None,
         })
         .map(|s| s.to_string())

src/main.rs

@@ -9,6 +9,7 @@ use figment::{
     providers::{Format, Serialized, Yaml},
     Figment,
 };
+use indicatif::{ProgressBar, ProgressStyle};
 use sea_orm_codegen::{
     DateTimeCrate as CodegenDateTimeCrate, EntityTransformer, EntityWriterContext, OutputFile,
     WithPrelude, WithSerde,
@@ -53,18 +54,32 @@ async fn main() -> Result<()> {
     files.extend(generate_files);
     tracing::info!("Generated {} files", files.len());
     fs::create_dir_all(&output_internal_entities).await?;
+    let progress_bar = ProgressBar::new((files.len() * 2) as u64)
+        .with_style(
+            ProgressStyle::default_bar()
+                .template("[{elapsed_precise}] {bar:40.cyan/blue} {pos}/{len} {msg}")?,
+        )
+        .with_message("Writing files");
+
     for (file_path, content) in files.iter() {
-        tracing::info!(?file_path, "Writing file");
+        progress_bar.set_message(format!("Writing file: {:?}", file_path));
+        tracing::debug!(?file_path, "Writing file");
         let mut file = fs::File::create(&file_path).await?;
         file.write_all(content.as_bytes()).await?;
+        progress_bar.inc(1);
     }
+    progress_bar.set_message("Running rustfmt");
+
     for (file_path, ..) in files.iter() {
-        tracing::info!(?file_path, "Running rustfmt");
+        tracing::debug!(?file_path, "Running rustfmt");
+        progress_bar.set_message(format!("Running rustfmt: {:?}", file_path));
         let exit_status = Command::new("rustfmt").arg(file_path).status().await?; // Get the status code
         if !exit_status.success() {
             // Propagate the error if any
             return Err(eyre!("Failed to run rustfmt"));
         }
+        progress_bar.inc(1);
     }
+    progress_bar.finish();
     Ok(())
 }