update generator

Cargo.toml

@@ -6,9 +6,12 @@ edition = "2021"
 [dependencies]
 clap = { version = "4.5.32", features = ["derive", "env"] }
 color-eyre = "0.6.3"
+comment-parser = "0.1.0"
 figment = { version = "0.10.19", features = ["yaml"] }
 heck = "0.5.0"
 inquire = "0.7.5"
+prettytable = "0.10.0"
+quote = "1.0.40"
 sea-orm-codegen = "1.1.8"
 sea-schema = { version = "0.16.1", features = ["sqlx-all"] }
 serde = { version = "1.0.219", features = ["derive"] }

src/generate.rs

@@ -1,12 +1,18 @@
 use core::time;
 
+const HEADER: &str = r#"AUTOGENERATED DATA"#;
+const COMMENTHEAD: &str = r#"/*"#;
+const COMMENTBODY: &str = r#" *"#;
+const COMMENTTAIL: &str = r#"*/"#;
 use crate::Config;
 use color_eyre::{
     eyre::{eyre, ContextCompat, Report},
     Result,
 };
+use comment_parser::{CommentParser, Event};
+use prettytable::{format, row, Table};
 use sea_orm_codegen::OutputFile;
-use sea_schema::sea_query::{self, TableCreateStatement};
+use sea_schema::sea_query::{self, ColumnSpec, TableCreateStatement};
 use tokio::{fs, task::JoinSet};
 use url::Url;
 
@@ -172,11 +178,12 @@ pub async fn generate_models(
     config: Config,
 ) -> Result<Vec<OutputFile>> {
     tracing::info!(?tables);
-    let output_path = config.output.path;
+    let output_path = config.output.path.clone();
     let files = tables
         .into_iter()
         .map(|table| {
             let output_path = output_path.clone();
+            let config = config.clone();
             async move {
                 let table_name = match table.get_table_name() {
                     Some(table_ref) => match table_ref {
@@ -192,18 +199,39 @@ pub async fn generate_models(
                     },
                     None => return Err(eyre!("Table name not found")),
                 };
-                let file_path = output_path.join(&table_name);
+                let table_str = generate_table(table, config.clone()).await?;
+                tracing::info!(?table_str);
+                let filename = format!("{}.rs", table_name);
+                let file_path = output_path.join(&filename);
                 let exists = file_path.exists();
-                let content = match exists {
+                let content = if exists {
-                    true => {
+                    let mut file_content = fs::read_to_string(file_path).await?;
-                        // let file_content = fs::read_to_string(path)
+                    let rules = comment_parser::get_syntax("rust").unwrap();
+                    let comments = CommentParser::new(&file_content, rules);
+                    let mut found = false;
+                    for comment in comments {
+                        if let Event::BlockComment(a, b) = comment {
+                            tracing::info!(?a, ?b);
+                            if b.contains(HEADER) {
+                                found = true;
+                                file_content = file_content.replace(a, &table_str);
+                                tracing::info!("Found header");
+                                break;
+                            }
+                        }
                     }
-                    false => {}
+                    if found {
-                };
+                        Ok::<String, Report>(file_content)
+                    } else {
+                        todo!()
+                    }
+                } else {
+                    todo!()
+                }?;
 
                 Ok(OutputFile {
-                    name: format!("{}.rs", table_name),
+                    name: filename,
-                    content: String::new(),
+                    content,
                 })
             }
         })
@@ -215,4 +243,66 @@ pub async fn generate_models(
         .collect::<Result<Vec<OutputFile>>>()?;
     Ok(files)
 }
+
+async fn generate_table(table: TableCreateStatement, config: Config) -> Result<String> {
+    let mut string = String::new();
+    string.push_str(format!("{COMMENTHEAD} {HEADER}\n").as_str());
+
+    let table_name = match table.get_table_name() {
+        Some(table_ref) => match table_ref {
+            sea_query::TableRef::Table(t)
+            | sea_query::TableRef::SchemaTable(_, t)
+            | sea_query::TableRef::DatabaseSchemaTable(_, _, t)
+            | sea_query::TableRef::TableAlias(t, _)
+            | sea_query::TableRef::SchemaTableAlias(_, t, _)
+            | sea_query::TableRef::DatabaseSchemaTableAlias(_, _, t, _) => t.to_string(),
+            _ => unimplemented!(),
+        },
+        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 = 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(' ', ' ', ' ', ' '),
+        )
+        .padding(1, 1)
+        .build();
+    ptable.set_format(format);
+    for column in table.get_columns() {
+        let name = column.get_column_name();
+        if let Some(column_type) = column.get_column_type() {
+            let attrs = attrs_to_string(column.get_column_spec());
+            ptable.add_row(row![name, format!("{:?}", column_type), attrs]);
+        }
+    }
+    inner.push_str(ptable.to_string().as_str());
+
+    string.push_str(
+        inner
+            .replace("\n", format!("\n{}", COMMENTBODY).as_str())
+            .as_str(),
+    );
+    string.push_str(format!("\n{COMMENTTAIL}\n").as_str());
+    Ok(string)
+}
+
+fn attrs_to_string(column: &Vec<ColumnSpec>) -> String {
+    column
+        .iter()
+        .filter_map(|c| match c {
+            ColumnSpec::PrimaryKey => Some("primary key"),
+            ColumnSpec::AutoIncrement => Some("autoincrement"),
+            _ => None,
+        })
+        .map(|s| s.to_string())
+        .collect::<Vec<String>>()
+        .join(", ")
+}
+
 //