1use sqlx::{AnyPool, Row};
2use floz_macros_core::model::snapshot::{SnapshotModel, SnapshotField};
3
4pub struct Migrator {
8 pool: AnyPool,
9}
10
11impl Migrator {
12 pub fn new(pool: AnyPool) -> Self {
14 Self { pool }
15 }
16
17 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 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 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 return Ok(());
43 }
44 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 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 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 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 db_columns.push(SnapshotField {
106 rust_name: col_name.clone(),
107 column_name: col_name,
108 type_info: data_type, modifiers: Vec::new(),
110 renamed_from: None,
111 });
112 }
113
114 Ok(Some(SnapshotModel {
115 name: table_name.to_string(), 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 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; }
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 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 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 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 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 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 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}