From 2cd6120867b49aff56719e639eb3929c19ad58df Mon Sep 17 00:00:00 2001 From: Nikkuss Date: Wed, 2 Apr 2025 17:11:58 +0400 Subject: [PATCH] update generator --- Cargo.toml | 3 ++ src/generate.rs | 110 +++++++++++++++++++++++++++++++++++++++++++----- 2 files changed, 103 insertions(+), 10 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 4638b6d..9c61034 100644 --- a/Cargo.toml +++ b/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"] } diff --git a/src/generate.rs b/src/generate.rs index ef2ad0c..c020adb 100644 --- a/src/generate.rs +++ b/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> { 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 { - true => { - // let file_content = fs::read_to_string(path) + let content = if exists { + let mut file_content = fs::read_to_string(file_path).await?; + 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::(file_content) + } else { + todo!() + } + } else { + todo!() + }?; Ok(OutputFile { - name: format!("{}.rs", table_name), - content: String::new(), + name: filename, + content, }) } }) @@ -215,4 +243,66 @@ pub async fn generate_models( .collect::>>()?; Ok(files) } + +async fn generate_table(table: TableCreateStatement, config: Config) -> Result { + 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) -> String { + column + .iter() + .filter_map(|c| match c { + ColumnSpec::PrimaryKey => Some("primary key"), + ColumnSpec::AutoIncrement => Some("autoincrement"), + _ => None, + }) + .map(|s| s.to_string()) + .collect::>() + .join(", ") +} + //