89 lines
3.1 KiB
Rust
89 lines
3.1 KiB
Rust
mod config;
|
|
mod generator;
|
|
|
|
use clap::Parser;
|
|
use color_eyre::{eyre::eyre, Report, Result};
|
|
use config::Config;
|
|
use figment::{
|
|
providers::{Format, Serialized, Yaml},
|
|
Figment,
|
|
};
|
|
use indicatif::{ProgressBar, ProgressStyle};
|
|
use sea_orm_codegen::{
|
|
DateTimeCrate as CodegenDateTimeCrate, EntityTransformer, EntityWriterContext, OutputFile,
|
|
WithPrelude, WithSerde,
|
|
};
|
|
use serde::{Deserialize, Serialize};
|
|
use tokio::{fs, io::AsyncWriteExt, process::Command};
|
|
|
|
#[derive(Parser, Debug)]
|
|
struct Args {
|
|
#[clap(short, long, default_value = "generator.yml")]
|
|
config: String,
|
|
#[clap(short, long, env = "DATABASE_URL")]
|
|
database_url: String,
|
|
}
|
|
#[tokio::main]
|
|
async fn main() -> Result<()> {
|
|
color_eyre::install()?;
|
|
tracing_subscriber::fmt::init();
|
|
let args = Args::parse();
|
|
|
|
let config: Config = Figment::new()
|
|
.merge(Serialized::defaults(Config::default()))
|
|
.merge(Yaml::file(&args.config))
|
|
.extract()?;
|
|
tracing::info!(?config);
|
|
tracing::info!(?args);
|
|
|
|
let mut handlebars = Handlebars::new();
|
|
|
|
let output_dir = &config.output.path;
|
|
let output_internal_entities = output_dir.join("_entities");
|
|
let (_, table_stmts) = generator::discover::get_tables(args.database_url, &config).await?;
|
|
let writer_context = config.clone().into();
|
|
let output = EntityTransformer::transform(table_stmts.clone())?.generate(&writer_context);
|
|
let mut files = output
|
|
.files
|
|
.into_iter()
|
|
.map(|OutputFile { name, content }| (output_internal_entities.join(name), content))
|
|
.collect::<Vec<_>>();
|
|
generator::update_files(table_stmts, config.clone(), &handlebars).await?;
|
|
// let generate_files = generator::generate_models(table_stmts, config.clone())
|
|
// .await?
|
|
// .into_iter()
|
|
// .map(|OutputFile { name, content }| (output_dir.join(name), content))
|
|
// .collect::<Vec<_>>();
|
|
// 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() {
|
|
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::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(())
|
|
}
|