This commit is contained in:
2025-04-02 22:02:26 +04:00
parent 2cd6120867
commit 556ae626ae
3 changed files with 74 additions and 16 deletions

View File

@@ -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"

View File

@@ -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,43 +263,83 @@ 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());
string.push_str(
inner
.replace("\n", format!("\n{}", COMMENTBODY).as_str())
.replace("\n", format!("\n{} ", COMMENTBODY).as_str())
.as_str(),
);
string.push_str(format!("\n{COMMENTTAIL}\n").as_str());
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())

View File

@@ -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(())
}