Skip to main content

floz_orm/
migration.rs

1use sqlx::{AnyPool, Row};
2use floz_macros_core::model::snapshot::{SnapshotModel, SnapshotField};
3
4/// Structural Fingerprinting Engine entry point.
5/// Compares Rust Abstract Syntax Tree snapshots loaded from `.json` 
6/// iteratively against `information_schema.columns`.
7pub struct Migrator {
8    pool: AnyPool,
9}
10
11impl Migrator {
12    /// Initializes a new Migrator against the provided `sqlx` Database Pool.
13    pub fn new(pool: AnyPool) -> Self {
14        Self { pool }
15    }
16
17    /// Migrate a single table using the Structural Fingerprinting Engine.
18    pub async fn migrate_table(&self, table_name: &str, history: &[SnapshotModel]) -> Result<(), String> {
19        let live_model: Option<SnapshotModel> = self.introspect_table(table_name).await
20            .map_err(|e| format!("Introspection failed: {}", e))?;
21
22        if live_model.is_none() {
23            // Table doesn't exist, execute all from v0 to the end.
24            println!("Table {} does not exist. Creating from v0.", table_name);
25            return self.apply_transitions(None, history).await;
26        }
27        
28        let live_model = live_model.unwrap();
29        
30        // Fingerprinting Algorithm: Scan backwards
31        let mut target_index = None;
32        for (i, snapshot) in history.iter().enumerate().rev() {
33            if self.is_structural_match(&live_model, snapshot) {
34                target_index = Some(i);
35                break;
36            }
37        }
38
39        if let Some(idx) = target_index {
40            if idx == history.len() - 1 {
41                // We perfectly match the latest!
42                return Ok(());
43            }
44            // Transition up through history
45            let upcoming_history = &history[idx..]; 
46            return self.apply_transitions(Some(&live_model), upcoming_history).await;
47        } else {
48            return Err("Error: Database structure does not match any known historical state. Tampering detected.".to_string());
49        }
50    }
51
52    /// Rollback a single table by one exact structural fingerprint.
53    pub async fn rollback_table(&self, table_name: &str, history: &[SnapshotModel]) -> Result<(), String> {
54        let live_model: Option<SnapshotModel> = self.introspect_table(table_name).await
55            .map_err(|e| format!("Introspection failed: {}", e))?;
56
57        if live_model.is_none() {
58            return Err(format!("Table {} does not exist, cannot rollback.", table_name));
59        }
60        
61        let live_model = live_model.unwrap();
62        
63        let mut target_index = None;
64        for (i, snapshot) in history.iter().enumerate().rev() {
65            if self.is_structural_match(&live_model, snapshot) {
66                target_index = Some(i);
67                break;
68            }
69        }
70
71        if let Some(idx) = target_index {
72            if idx == 0 {
73                println!("Table {} is already at v0. Rollback would drop it (not implemented automatically for safety).", table_name);
74                return Ok(());
75            }
76            // Transition down 1 version
77            let target_version = &history[idx - 1];
78            return self.apply_transitions(Some(&live_model), std::slice::from_ref(target_version)).await;
79        } else {
80            return Err("Error: Database structure does not match any known historical state. Tampering detected.".to_string());
81        }
82    }
83
84    async fn introspect_table(&self, table_name: &str) -> sqlx::Result<Option<SnapshotModel>> {
85        // Query postgres information_schema
86        let rows = sqlx::query(
87            "SELECT column_name, data_type, is_nullable, character_maximum_length, column_default 
88             FROM information_schema.columns 
89             WHERE table_name = $1 AND table_schema = 'public'"
90        )
91        .bind(table_name)
92        .fetch_all(&self.pool)
93        .await?;
94
95        if rows.is_empty() {
96            return Ok(None);
97        }
98
99        let mut db_columns = Vec::new();
100        for row in rows {
101            let col_name: String = row.get("column_name");
102            let data_type: String = row.get("data_type");
103            // let is_nullable: String = row.get("is_nullable"); // YES / NO
104            
105            db_columns.push(SnapshotField {
106                rust_name: col_name.clone(),
107                column_name: col_name,
108                type_info: data_type, // Postgres data type
109                modifiers: Vec::new(),
110                renamed_from: None,
111            });
112        }
113
114        Ok(Some(SnapshotModel {
115            name: table_name.to_string(), // Dummy
116            table_name: table_name.to_string(),
117            db_columns,
118            relationships: Vec::new(),
119            constraints: Vec::new(),
120        }))
121    }
122
123    fn is_structural_match(&self, live: &SnapshotModel, snapshot: &SnapshotModel) -> bool {
124        // Simple heuristic: just check if column names match exactly.
125        // We can add data_type mapping later for robust type-change diffing.
126        if live.db_columns.len() != snapshot.db_columns.len() {
127            return false;
128        }
129
130        for live_col in &live.db_columns {
131            if !snapshot.db_columns.iter().any(|c| c.column_name == live_col.column_name) {
132                return false; // Live has a column snapshot doesn't
133            }
134        }
135        true
136    }
137
138    async fn apply_transitions(&self, current: Option<&SnapshotModel>, upcoming: &[SnapshotModel]) -> Result<(), String> {
139        let mut current_state = current;
140        let mut tx = self.pool.begin().await.map_err(|e| e.to_string())?;
141
142        for next_state in upcoming {
143            let sqls = self.diff_snapshots(current_state, next_state);
144            for sql in sqls {
145                println!("INFO: Executing: {}", sql);
146                sqlx::query(&sql).execute(&mut *tx).await.map_err(|e| format!("Migration failed: {}\nSQL: {}", e, sql))?;
147            }
148            current_state = Some(next_state);
149        }
150
151        tx.commit().await.map_err(|e| e.to_string())?;
152        Ok(())
153    }
154
155    /// Analytically computes exact directional `ALTER TABLE` / `CREATE TABLE` translations
156    /// between two structural states using the AST Snapshot models.
157    pub fn diff_snapshots(&self, a: Option<&SnapshotModel>, b: &SnapshotModel) -> Vec<String> {
158        let mut queries = Vec::new();
159        let table = &b.table_name;
160        
161        if let Some(old) = a {
162            // Renames
163            for new_col in &b.db_columns {
164                if let Some(ref old_name) = new_col.renamed_from {
165                    queries.push(format!("ALTER TABLE {} RENAME COLUMN {} TO {};", table, old_name, new_col.column_name));
166                }
167            }
168            
169            // Drops
170            for old_col in &old.db_columns {
171                if b.db_columns.iter().any(|c| c.renamed_from.as_ref() == Some(&old_col.column_name)) { continue; }
172                if !b.db_columns.iter().any(|c| c.column_name == old_col.column_name) {
173                    queries.push(format!("ALTER TABLE {} DROP COLUMN {};", table, old_col.column_name));
174                }
175            }
176            
177            // Adds
178            for new_col in &b.db_columns {
179                if new_col.renamed_from.is_some() { continue; }
180                if !old.db_columns.iter().any(|c| c.column_name == new_col.column_name) {
181                    let pg_type = self.map_type_to_pg(&new_col.type_info);
182                    let mut def = format!("ALTER TABLE {} ADD COLUMN {} {}", table, new_col.column_name, pg_type);
183                    if new_col.modifiers.contains(&"Primary".to_string()) { def.push_str(" PRIMARY KEY"); }
184                    if new_col.modifiers.contains(&"Unique".to_string()) { def.push_str(" UNIQUE"); }
185                    if !new_col.modifiers.contains(&"Nullable".to_string()) { def.push_str(" NOT NULL"); }
186                    queries.push(def); 
187                }
188            }
189        } else {
190            // Create Table
191            let mut defs = Vec::new();
192            for col in &b.db_columns {
193                let pg_type = self.map_type_to_pg(&col.type_info);
194                let mut def = format!("{} {}", col.column_name, pg_type);
195                
196                if col.modifiers.contains(&"Primary".to_string()) { def.push_str(" PRIMARY KEY"); }
197                if col.modifiers.contains(&"AutoIncrement".to_string()) { 
198                    def = def.replace("INTEGER", "SERIAL").replace("BIGINT", "BIGSERIAL"); 
199                }
200                if col.modifiers.contains(&"Unique".to_string()) { def.push_str(" UNIQUE"); }
201                if !col.modifiers.contains(&"Nullable".to_string()) { def.push_str(" NOT NULL"); }
202                defs.push(def);
203            }
204            queries.push(format!("CREATE TABLE {} (\n  {}\n);", table, defs.join(",\n  ")));
205        }
206        
207        queries
208    }
209
210    /// Fallbacks Rust AST macro-string definitions securely into Native PostgreSQL string syntaxes.
211    pub fn map_type_to_pg(&self, ti: &str) -> String {
212        if ti.starts_with("Varchar") || ti.starts_with("Decimal") {
213            ti.to_uppercase()
214        } else {
215            match ti {
216                "Integer" => "INTEGER".to_string(),
217                "BigInt" => "BIGINT".to_string(),
218                "Text" => "TEXT".to_string(),
219                "DateTime" => "TIMESTAMPTZ".to_string(),
220                "Uuid" => "UUID".to_string(),
221                "Bool" => "BOOLEAN".to_string(),
222                _ => "TEXT".to_string(),
223            }
224        }
225    }
226}