axion_db/
manager.rs

1// axion-db/src/manager.rs
2use crate::{
3    client::DbClient,
4    config::DbConfig,
5    error::DbResult,
6    introspection::{self, Introspector},
7    // IMPORTANT: Make RoutineKind accessible for matching
8    metadata::{DatabaseMetadata, RoutineKind},
9};
10use comfy_table::{presets::UTF8_FULL, Cell, CellAlignment, Table}; // Import comfy-table
11use owo_colors::OwoColorize; // Import the colorize trait
12use std::sync::Arc;
13use tracing::info;
14
15/// The ModelManager is the primary entry point for database introspection.
16/// It holds the complete database schema and provides methods to interact with it.
17#[derive(Clone)]
18pub struct ModelManager {
19    pub db_client: Arc<DbClient>,
20    pub metadata: Arc<DatabaseMetadata>,
21    introspector: Arc<dyn Introspector>,
22}
23
24impl ModelManager {
25    /// Creates a new ModelManager by connecting to the database and performing a full introspection.
26    pub async fn new(config: DbConfig) -> DbResult<Self> {
27        info!("Initializing ModelManager...");
28        let db_client = Arc::new(DbClient::new(config).await?);
29        let introspector = introspection::new_introspector(db_client.clone())?;
30
31        info!("Discovering user schemas...");
32        let schemas = introspector.list_user_schemas().await?;
33
34        info!("Performing full database introspection...");
35        let metadata = introspector.introspect(&schemas).await?;
36        info!(
37            "Introspection complete. Found {} schemas.",
38            metadata.schemas.len()
39        );
40
41        Ok(Self {
42            db_client,
43            metadata: Arc::new(metadata),
44            introspector: Arc::from(introspector),
45        })
46    }
47
48    // =================================================================================
49    //  DX: Pretty-Printing Methods (WITH THE NEW `display_summary`)
50    // =================================================================================
51
52    /// Prints a rich, colorized, table-based summary of the database metadata.
53    pub fn display_summary(&self) {
54        println!(); // Add a newline for spacing
55
56        let mut table = Table::new();
57
58        // THE FIX: Use the UTF8_BORDERS_ONLY preset.
59        // This preset creates the outer box and the line under the header,
60        // but no horizontal lines between the data rows, achieving the clean look.
61        table
62            .load_preset(comfy_table::presets::UTF8_BORDERS_ONLY)
63            .set_header(vec![
64                Cell::new("Schema").add_attribute(comfy_table::Attribute::Bold),
65                Cell::new("Tables").add_attribute(comfy_table::Attribute::Bold),
66                Cell::new("Views").add_attribute(comfy_table::Attribute::Bold),
67                Cell::new("Enums").add_attribute(comfy_table::Attribute::Bold),
68                Cell::new("Functions").add_attribute(comfy_table::Attribute::Bold),
69                Cell::new("Procedures").add_attribute(comfy_table::Attribute::Bold),
70                Cell::new("Triggers").add_attribute(comfy_table::Attribute::Bold),
71                Cell::new("Total").add_attribute(comfy_table::Attribute::Bold),
72            ]);
73
74        // --- Totals Initialization ---
75        let mut total_tables = 0;
76        let mut total_views = 0;
77        let mut total_enums = 0;
78        let mut total_functions = 0;
79        let mut total_procedures = 0;
80        let mut total_triggers = 0;
81
82        // --- Sort schemas for consistent output ---
83        let mut schemas: Vec<_> = self.metadata.schemas.keys().collect();
84        schemas.sort();
85
86        for schema_name in schemas {
87            if let Some(schema_data) = self.metadata.schemas.get(schema_name) {
88                // --- Per-schema Counts ---
89                let tables_count = schema_data.tables.len();
90                let views_count = schema_data.views.len();
91                let enums_count = schema_data.enums.len();
92
93                let mut functions_count = 0;
94                let mut procedures_count = 0;
95                let mut triggers_count = 0;
96                for func_meta in schema_data.functions.values() {
97                    match func_meta.kind {
98                        Some(RoutineKind::Function) => functions_count += 1,
99                        Some(RoutineKind::Procedure) => procedures_count += 1,
100                        Some(RoutineKind::Trigger) => triggers_count += 1,
101                        _ => {}
102                    }
103                }
104
105                let schema_total = tables_count + views_count + enums_count + functions_count + procedures_count + triggers_count;
106
107                // --- Add to Grand Totals ---
108                total_tables += tables_count;
109                total_views += views_count;
110                total_enums += enums_count;
111                total_functions += functions_count;
112                total_procedures += procedures_count;
113                total_triggers += triggers_count;
114
115                // --- Build and Add the Row ---
116                table.add_row(vec![
117                    Cell::new(schema_name).fg(comfy_table::Color::Cyan),
118                    Cell::new(tables_count).set_alignment(CellAlignment::Right).fg(comfy_table::Color::Blue),
119                    Cell::new(views_count).set_alignment(CellAlignment::Right).fg(comfy_table::Color::Green),
120                    Cell::new(enums_count).set_alignment(CellAlignment::Right).fg(comfy_table::Color::Magenta),
121                    Cell::new(functions_count).set_alignment(CellAlignment::Right).fg(comfy_table::Color::Red),
122                    Cell::new(procedures_count).set_alignment(CellAlignment::Right).fg(comfy_table::Color::Yellow),
123                    Cell::new(triggers_count).set_alignment(CellAlignment::Right).fg(comfy_table::Color::DarkYellow),
124                    Cell::new(schema_total).set_alignment(CellAlignment::Right).add_attribute(comfy_table::Attribute::Bold),
125                ]);
126            }
127        }
128
129        // --- Grand Total Calculation ---
130        let grand_total = total_tables + total_views + total_enums + total_functions + total_procedures + total_triggers;
131
132        // --- Add the TOTAL row which will act as the footer ---
133        // This row will have the bottom border of the table drawn after it.
134        table.add_row(vec![
135            Cell::new("TOTAL").add_attribute(comfy_table::Attribute::Bold),
136            Cell::new(total_tables).set_alignment(CellAlignment::Right).fg(comfy_table::Color::Blue).add_attribute(comfy_table::Attribute::Bold),
137            Cell::new(total_views).set_alignment(CellAlignment::Right).fg(comfy_table::Color::Green).add_attribute(comfy_table::Attribute::Bold),
138            Cell::new(total_enums).set_alignment(CellAlignment::Right).fg(comfy_table::Color::Magenta).add_attribute(comfy_table::Attribute::Bold),
139            Cell::new(total_functions).set_alignment(CellAlignment::Right).fg(comfy_table::Color::Red).add_attribute(comfy_table::Attribute::Bold),
140            Cell::new(total_procedures).set_alignment(CellAlignment::Right).fg(comfy_table::Color::Yellow).add_attribute(comfy_table::Attribute::Bold),
141            Cell::new(total_triggers).set_alignment(CellAlignment::Right).fg(comfy_table::Color::DarkYellow).add_attribute(comfy_table::Attribute::Bold),
142            Cell::new(grand_total).set_alignment(CellAlignment::Right).add_attribute(comfy_table::Attribute::Bold),
143        ]);
144
145        // Print the title and the final table
146        println!("{}", " ModelManager Statistics".green().bold().underline());
147        println!("{table}");
148    }
149
150    /// Prints a detailed, prism-py-like breakdown of tables for the specified schemas.
151    /// If `schemas` is empty, it displays all schemas.
152    pub fn display_tables(&self, schemas: &[&str]) {
153        println!("\n{:=<80}", "");
154        println!("           TABLES OVERVIEW");
155        println!("{:=<80}\n", "");
156
157        let schemas_to_display: Box<dyn Iterator<Item = &str>> = if schemas.is_empty() {
158            Box::new(self.metadata.schemas.keys().map(|s| s.as_str()))
159        } else {
160            Box::new(schemas.iter().copied())
161        };
162
163        for schema_name in schemas_to_display {
164            if let Some(schema_data) = self.metadata.schemas.get(schema_name) {
165                for table_data in schema_data.tables.values() {
166                    // This now uses the beautiful `Display` implementation we wrote for TableMetadata
167                    println!("{}\n", table_data);
168                }
169            }
170        }
171    }
172
173    /// Prints a detailed, prism-py-like breakdown of views for the specified schemas.
174    /// If `schemas` is empty, it displays all schemas.
175    pub fn display_views(&self, schemas: &[&str]) {
176        println!("\n{:=<80}", "");
177        println!("           VIEWS OVERVIEW");
178        println!("{:=<80}\n", "");
179
180        let schemas_to_display: Box<dyn Iterator<Item = &str>> = if schemas.is_empty() {
181            Box::new(self.metadata.schemas.keys().map(|s| s.as_str()))
182        } else {
183            Box::new(schemas.iter().copied())
184        };
185
186        for schema_name in schemas_to_display {
187            if let Some(schema_data) = self.metadata.schemas.get(schema_name) {
188                for view_data in schema_data.views.values() {
189                    // Uses the `Display` implementation for ViewMetadata
190                    println!("{}\n", view_data);
191                }
192            }
193        }
194    }
195
196    /// Prints a summary of all enums for the specified schemas with enhanced formatting.
197    /// If `schemas` is empty, it displays all schemas.
198    pub fn display_enums(&self, schemas: &[&str]) {
199        println!("\n{:=<80}", "");
200        println!("           ENUMS OVERVIEW");
201        println!("{:=<80}\n", "");
202
203        let schemas_to_display: Box<dyn Iterator<Item = &str>> = if schemas.is_empty() {
204            Box::new(self.metadata.schemas.keys().map(|s| s.as_str()))
205        } else {
206            Box::new(schemas.iter().copied())
207        };
208
209        for schema_name in schemas_to_display {
210            if let Some(schema_data) = self.metadata.schemas.get(schema_name) {
211                if !schema_data.enums.is_empty() {
212                    println!("Schema '{}':", schema_name.cyan().bold());
213                    for enum_data in schema_data.enums.values() {
214                        // Print the enum name, indented and in yellow.
215                        println!("  {}", enum_data.name.yellow());
216
217                        // Format the values string, indented further, and styled.
218                        let values_str = format!("({})", enum_data.values.join(", "));
219                        println!("    {}", values_str.dimmed().italic());
220
221                        // Add a blank line for spacing between enums.
222                        println!();
223                    }
224                }
225            }
226        }
227    }
228}
axion_db/manager.rs

axion_db/
manager.rs
