update

src/main.rs

@@ -1,93 +1,21 @@
 mod config;
 mod generate;
-use std::{fs, path::PathBuf, str::FromStr};
+use std::{path::PathBuf, str::FromStr};
 
 use clap::Parser;
-use color_eyre::{Report, Result};
+use color_eyre::{eyre::eyre, Report, Result};
 use config::Config;
 use figment::{
     providers::{Format, Serialized, Yaml},
     Figment,
 };
 use sea_orm_codegen::{
-    DateTimeCrate as CodegenDateTimeCrate, EntityTransformer, EntityWriterContext, WithPrelude,
-    WithSerde,
+    DateTimeCrate as CodegenDateTimeCrate, EntityTransformer, EntityWriterContext, OutputFile,
+    WithPrelude, WithSerde,
 };
 use serde::{Deserialize, Serialize};
+use tokio::{fs, io::AsyncWriteExt, process::Command};
 
-// #[derive(Deserialize, Serialize, Debug, Clone)]
-// struct Config {
-//     sea_orm: SeaOrmConfig,
-// }
-
-// #[derive(Deserialize, Serialize, Debug, Clone)]
-// struct SeaOrmConfig {
-//     expanded_format: bool,
-//     table: SeaOrmTableConfig,
-//     with_serde: String,
-//     with_prelude: String,
-//     with_copy_enums: bool,
-//     serde_skip_deserializing_primary_key: bool,
-//     serde_skip_hidden_column: bool,
-//     max_connections: u32,
-//     acquire_timeout: u64,
-//     database_schema: Option<String>,
-//     date_format: DateTimeCrate,
-// }
-// #[derive(Deserialize, Serialize, Debug, Clone)]
-// struct SeaOrmTableConfig {
-//     include_hidden: bool,
-//     only: Vec<String>,
-//     exclude: Vec<String>,
-// }
-//
-// impl Default for Config {
-//     fn default() -> Self {
-//         Self {
-//             sea_orm: SeaOrmConfig {
-//                 table: SeaOrmTableConfig {
-//                     include_hidden: false,
-//                     only: Vec::new(),
-//                     exclude: Vec::new(),
-//                 },
-//                 database_schema: None,
-//                 max_connections: 10,
-//                 acquire_timeout: 30,
-//                 expanded_format: false,
-//                 with_copy_enums: false,
-//                 with_serde: String::from("none"),
-//                 with_prelude: String::from("none"),
-//                 serde_skip_hidden_column: false,
-//                 serde_skip_deserializing_primary_key: false,
-//                 date_format: DateTimeCrate::Time,
-//             },
-//         }
-//     }
-// }
-
-// impl TryInto<EntityWriterContext> for Config {
-//     type Error = Report;
-//     fn try_into(self) -> Result<EntityWriterContext> {
-//         Ok(EntityWriterContext::new(
-//             self.sea_orm.expanded_format,
-//             WithPrelude::from_str(&self.sea_orm.with_prelude)?,
-//             WithSerde::from_str(&self.sea_orm.with_serde)?,
-//             self.sea_orm.with_copy_enums,
-//             self.sea_orm.date_format.into(),
-//             self.sea_orm.database_schema,
-//             false,
-//             self.sea_orm.serde_skip_deserializing_primary_key,
-//             self.sea_orm.serde_skip_hidden_column,
-//             Vec::new(),
-//             Vec::new(),
-//             Vec::new(),
-//             Vec::new(),
-//             false,
-//             false,
-//         ))
-//     }
-// }
-//
 #[derive(Parser, Debug)]
 struct Args {
     #[clap(short, long, default_value = "generator.yml")]
@@ -102,13 +30,41 @@ async fn main() -> Result<()> {
     let args = Args::parse();
 
     let config: Config = Figment::new()
-        // .merge(Serialized::defaults(Config::default()))
+        .merge(Serialized::defaults(Config::default()))
         .merge(Yaml::file(&args.config))
         .extract()?;
     tracing::info!(?config);
     tracing::info!(?args);
-    // let (_, table_stmts) = generate::get_tables(args.database_url, &config).await?;
-    // let writer_context = config.clone().try_into()?;
-    // let output = EntityTransformer::transform(table_stmts)?.generate(&writer_context);
+    let output_dir = &config.output.path;
+    let output_internal_entities = output_dir.join("_entities");
+    let (_, table_stmts) = generate::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<_>>();
+    let generate_files = generate::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?;
+    for (file_path, content) in files.iter() {
+        tracing::info!(?file_path, "Writing file");
+        let mut file = fs::File::create(&file_path).await?;
+        file.write_all(content.as_bytes()).await?;
+    }
+    for (file_path, ..) in files.iter() {
+        tracing::info!(?file_path, "Running rustfmt");
+        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"));
+        }
+    }
     Ok(())
 }