rust_orm_gen/
context.rs

1use crate::error::OrmError;
2use crate::metadata::{get_tables, get_columns};
3use crate::generator::generate_struct;
4use crate::crud::generate_crud_operations;
5use std::collections::HashMap;
6use std::fs;
7use std::path::Path;
8use log::{info, error};
9use crate::db::PostgresConnectionManager;
10use chrono::Utc;
11
12pub struct DbContext {
13    pub manager: PostgresConnectionManager,
14}
15
16impl DbContext {
17    pub async fn new(database_url: &str) -> Result<Self, OrmError> {
18        let manager = PostgresConnectionManager::new(database_url.to_string());
19        Ok(Self { manager })
20    }
21
22    pub async fn reverse_engineer(&self, output_dir: &str, author: &str, github_link: &str) -> Result<(), OrmError> {
23        info!("Reverse engineering the database schema");
24        let conn = self.manager.connect().await?;
25        let tables = get_tables(&conn).await?;
26        let date = Utc::now().date_naive();
27        for table in tables {
28            info!("Processing table: {}", table);
29            match get_columns(&conn, &table).await {
30                Ok(columns) => {
31                    let columns_map: HashMap<String, String> = columns.into_iter().collect();
32                    let struct_def = generate_struct(&table, columns_map.clone(), author, github_link, date);
33                    let crud_ops = generate_crud_operations(&table, columns_map, author, github_link, date);
34
35                    // Ensure output directory exists
36                    fs::create_dir_all(output_dir)?;
37
38                    // Write struct definition to file
39                    let struct_file_path = Path::new(output_dir).join(format!("{}.rs", table));
40                    fs::write(&struct_file_path, struct_def)
41                        .map_err(|e| OrmError::IoError(e))?;
42
43                    // Write CRUD operations to file
44                    let crud_file_path = Path::new(output_dir).join(format!("{}_crud.rs", table));
45                    fs::write(&crud_file_path, crud_ops)
46                        .map_err(|e| OrmError::IoError(e))?;
47
48                    info!("Completed processing table: {}", table);
49                }
50                Err(e) => error!("Failed to get columns for table {}: {}", table, e),
51            }
52        }
53        Ok(())
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60    use tokio;
61    use dotenv::dotenv;
62    use std::env;
63
64    #[tokio::test]
65    async fn test_reverse_engineer() {
66        dotenv().ok();
67        let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
68        let db_context = DbContext::new(&database_url).await.unwrap();
69        let result = db_context.reverse_engineer("db", "Tom Blanchard", "https://github.com/tomblanchard312/rust_orm_gen").await;
70        if let Err(e) = &result {
71            eprintln!("Reverse engineering failed: {:?}", e);
72        }
73        assert!(result.is_ok(), "Reverse engineering should succeed");
74    }
75}
rust_orm_gen/context.rs

rust_orm_gen/
context.rs
