1use serde::{Deserialize, Serialize};
7use serde_json::{Value, json};
8use tokio_postgres::Client;
9
10use crate::error::ToolError;
11use crate::sql::is_safe_ident;
12
13#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
14#[serde(rename_all = "camelCase")]
15pub struct SchemaSnapshot {
16 pub tables: Vec<TableSnapshot>,
17}
18
19#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
20#[serde(rename_all = "camelCase")]
21pub struct TableSnapshot {
22 pub name: String,
23 pub schema: String,
24 pub columns: Vec<ColumnSnapshot>,
25 pub constraints: Vec<ConstraintSnapshot>,
26 pub indexes: Vec<IndexSnapshot>,
27}
28
29#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
30pub struct ColumnSnapshot {
31 pub column_name: String,
32 pub data_type: String,
33 pub not_null: bool,
34 pub default_value: Option<String>,
35 pub ordinal: i32,
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
39pub struct ConstraintSnapshot {
40 pub name: String,
41 #[serde(rename = "type")]
42 pub type_: String,
43 pub definition: String,
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
47pub struct IndexSnapshot {
48 pub name: String,
49 pub definition: String,
50 pub is_unique: bool,
51 pub is_primary: bool,
52}
53
54#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
55#[serde(rename_all = "lowercase")]
56pub enum DiffStatus {
57 Added,
58 Removed,
59 Changed,
60 Unchanged,
61}
62
63#[derive(Debug, Clone, Serialize, Deserialize)]
64#[serde(rename_all = "camelCase")]
65pub struct TableDiff {
66 pub name: String,
67 pub status: DiffStatus,
68 pub column_diffs: Vec<ColumnDiff>,
69 pub constraint_diffs: Vec<ConstraintDiff>,
70 pub index_diffs: Vec<IndexDiff>,
71}
72
73#[derive(Debug, Clone, Serialize, Deserialize)]
74#[serde(rename_all = "camelCase")]
75pub struct ColumnDiff {
76 pub name: String,
77 pub status: DiffStatus,
78 #[serde(skip_serializing_if = "Option::is_none")]
79 pub before: Option<ColumnSnapshot>,
80 #[serde(skip_serializing_if = "Option::is_none")]
81 pub after: Option<ColumnSnapshot>,
82}
83
84#[derive(Debug, Clone, Serialize, Deserialize)]
85#[serde(rename_all = "camelCase")]
86pub struct ConstraintDiff {
87 pub name: String,
88 pub status: DiffStatus,
89 #[serde(skip_serializing_if = "Option::is_none")]
90 pub before: Option<ConstraintSnapshot>,
91 #[serde(skip_serializing_if = "Option::is_none")]
92 pub after: Option<ConstraintSnapshot>,
93}
94
95#[derive(Debug, Clone, Serialize, Deserialize)]
96#[serde(rename_all = "camelCase")]
97pub struct IndexDiff {
98 pub name: String,
99 pub status: DiffStatus,
100 #[serde(skip_serializing_if = "Option::is_none")]
101 pub before: Option<IndexSnapshot>,
102 #[serde(skip_serializing_if = "Option::is_none")]
103 pub after: Option<IndexSnapshot>,
104}
105
106pub fn require_safe_schema(schema: &str) -> Result<(), ToolError> {
107 if !is_safe_ident(schema) {
108 return Err(ToolError::InvalidArgs(format!(
109 "Invalid schema name \"{schema}\""
110 )));
111 }
112 Ok(())
113}
114
115pub async fn load_schema_snapshot(
116 client: &Client,
117 schema: &str,
118) -> Result<SchemaSnapshot, ToolError> {
119 require_safe_schema(schema)?;
120 let table_rows = client
121 .query(
122 r#"
123 SELECT c.relname AS name
124 FROM pg_class c
125 JOIN pg_namespace n ON n.oid = c.relnamespace
126 WHERE n.nspname = $1
127 AND c.relkind = 'r'
128 AND NOT c.relispartition
129 ORDER BY c.relname
130 "#,
131 &[&schema],
132 )
133 .await?;
134
135 let mut tables = Vec::with_capacity(table_rows.len());
136 for row in &table_rows {
137 let name: String = row.get("name");
138 let columns = load_columns(client, schema, &name).await?;
139 let constraints = load_constraints(client, schema, &name).await?;
140 let indexes = load_indexes(client, schema, &name).await?;
141 tables.push(TableSnapshot {
142 name,
143 schema: schema.to_owned(),
144 columns,
145 constraints,
146 indexes,
147 });
148 }
149 Ok(SchemaSnapshot { tables })
150}
151
152async fn load_columns(
153 client: &Client,
154 schema: &str,
155 table: &str,
156) -> Result<Vec<ColumnSnapshot>, ToolError> {
157 let rows = client
158 .query(
159 r#"
160 SELECT
161 a.attname AS column_name,
162 pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type,
163 a.attnotnull AS not_null,
164 pg_get_expr(ad.adbin, ad.adrelid) AS default_value,
165 a.attnum AS ordinal
166 FROM pg_attribute a
167 JOIN pg_class c ON c.oid = a.attrelid
168 JOIN pg_namespace n ON n.oid = c.relnamespace
169 LEFT JOIN pg_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum
170 WHERE n.nspname = $1
171 AND c.relname = $2
172 AND a.attnum > 0
173 AND NOT a.attisdropped
174 ORDER BY a.attnum
175 "#,
176 &[&schema, &table],
177 )
178 .await?;
179 Ok(rows
180 .iter()
181 .map(|r| ColumnSnapshot {
182 column_name: r.get("column_name"),
183 data_type: r.get("data_type"),
184 not_null: r.get("not_null"),
185 default_value: r.get("default_value"),
186 ordinal: r.get::<_, i16>("ordinal") as i32,
187 })
188 .collect())
189}
190
191async fn load_constraints(
192 client: &Client,
193 schema: &str,
194 table: &str,
195) -> Result<Vec<ConstraintSnapshot>, ToolError> {
196 let rows = client
197 .query(
198 r#"
199 SELECT
200 con.conname AS name,
201 con.contype::text AS type,
202 pg_get_constraintdef(con.oid, true) AS definition
203 FROM pg_constraint con
204 JOIN pg_class c ON c.oid = con.conrelid
205 JOIN pg_namespace n ON n.oid = c.relnamespace
206 WHERE n.nspname = $1
207 AND c.relname = $2
208 ORDER BY con.conname
209 "#,
210 &[&schema, &table],
211 )
212 .await?;
213 Ok(rows
214 .iter()
215 .map(|r| ConstraintSnapshot {
216 name: r.get("name"),
217 type_: r.get("type"),
218 definition: r.get("definition"),
219 })
220 .collect())
221}
222
223async fn load_indexes(
224 client: &Client,
225 schema: &str,
226 table: &str,
227) -> Result<Vec<IndexSnapshot>, ToolError> {
228 let rows = client
229 .query(
230 r#"
231 SELECT
232 i.relname AS name,
233 pg_get_indexdef(i.oid) AS definition,
234 ix.indisunique AS is_unique,
235 ix.indisprimary AS is_primary
236 FROM pg_index ix
237 JOIN pg_class t ON t.oid = ix.indrelid
238 JOIN pg_namespace n ON n.oid = t.relnamespace
239 JOIN pg_class i ON i.oid = ix.indexrelid
240 WHERE n.nspname = $1
241 AND t.relname = $2
242 ORDER BY i.relname
243 "#,
244 &[&schema, &table],
245 )
246 .await?;
247 Ok(rows
248 .iter()
249 .map(|r| IndexSnapshot {
250 name: r.get("name"),
251 definition: r.get("definition"),
252 is_unique: r.get("is_unique"),
253 is_primary: r.get("is_primary"),
254 })
255 .collect())
256}
257
258pub fn compute_schema_diff(source: &SchemaSnapshot, target: &SchemaSnapshot) -> Vec<TableDiff> {
259 let source_map: std::collections::HashMap<&str, &TableSnapshot> =
260 source.tables.iter().map(|t| (t.name.as_str(), t)).collect();
261 let target_map: std::collections::HashMap<&str, &TableSnapshot> =
262 target.tables.iter().map(|t| (t.name.as_str(), t)).collect();
263
264 let mut names: Vec<&str> = source_map
265 .keys()
266 .chain(target_map.keys())
267 .copied()
268 .collect();
269 names.sort();
270 names.dedup();
271
272 let mut diffs = Vec::new();
273 for table_name in names {
274 let src = source_map.get(table_name).copied();
275 let tgt = target_map.get(table_name).copied();
276 match (src, tgt) {
277 (None, Some(tgt_table)) => diffs.push(TableDiff {
278 name: table_name.to_owned(),
279 status: DiffStatus::Added,
280 column_diffs: tgt_table
281 .columns
282 .iter()
283 .map(|c| ColumnDiff {
284 name: c.column_name.clone(),
285 status: DiffStatus::Added,
286 before: None,
287 after: Some(c.clone()),
288 })
289 .collect(),
290 constraint_diffs: tgt_table
291 .constraints
292 .iter()
293 .map(|c| ConstraintDiff {
294 name: c.name.clone(),
295 status: DiffStatus::Added,
296 before: None,
297 after: Some(c.clone()),
298 })
299 .collect(),
300 index_diffs: tgt_table
301 .indexes
302 .iter()
303 .map(|i| IndexDiff {
304 name: i.name.clone(),
305 status: DiffStatus::Added,
306 before: None,
307 after: Some(i.clone()),
308 })
309 .collect(),
310 }),
311 (Some(src_table), None) => diffs.push(TableDiff {
312 name: table_name.to_owned(),
313 status: DiffStatus::Removed,
314 column_diffs: src_table
315 .columns
316 .iter()
317 .map(|c| ColumnDiff {
318 name: c.column_name.clone(),
319 status: DiffStatus::Removed,
320 before: Some(c.clone()),
321 after: None,
322 })
323 .collect(),
324 constraint_diffs: src_table
325 .constraints
326 .iter()
327 .map(|c| ConstraintDiff {
328 name: c.name.clone(),
329 status: DiffStatus::Removed,
330 before: Some(c.clone()),
331 after: None,
332 })
333 .collect(),
334 index_diffs: src_table
335 .indexes
336 .iter()
337 .map(|i| IndexDiff {
338 name: i.name.clone(),
339 status: DiffStatus::Removed,
340 before: Some(i.clone()),
341 after: None,
342 })
343 .collect(),
344 }),
345 (Some(src_table), Some(tgt_table)) => {
346 let column_diffs = diff_columns(&src_table.columns, &tgt_table.columns);
347 let constraint_diffs =
348 diff_constraints(&src_table.constraints, &tgt_table.constraints);
349 let index_diffs = diff_indexes(&src_table.indexes, &tgt_table.indexes);
350 let has_changes = column_diffs
351 .iter()
352 .any(|d| d.status != DiffStatus::Unchanged)
353 || constraint_diffs
354 .iter()
355 .any(|d| d.status != DiffStatus::Unchanged)
356 || index_diffs
357 .iter()
358 .any(|d| d.status != DiffStatus::Unchanged);
359 diffs.push(TableDiff {
360 name: table_name.to_owned(),
361 status: if has_changes {
362 DiffStatus::Changed
363 } else {
364 DiffStatus::Unchanged
365 },
366 column_diffs,
367 constraint_diffs,
368 index_diffs,
369 });
370 }
371 (None, None) => {}
372 }
373 }
374
375 let order = |s: DiffStatus| match s {
376 DiffStatus::Changed => 0,
377 DiffStatus::Added => 1,
378 DiffStatus::Removed => 2,
379 DiffStatus::Unchanged => 3,
380 };
381 diffs.sort_by_key(|d| order(d.status));
382 diffs
383}
384
385fn diff_columns(src: &[ColumnSnapshot], tgt: &[ColumnSnapshot]) -> Vec<ColumnDiff> {
386 let src_map: std::collections::HashMap<&str, &ColumnSnapshot> =
387 src.iter().map(|c| (c.column_name.as_str(), c)).collect();
388 let tgt_map: std::collections::HashMap<&str, &ColumnSnapshot> =
389 tgt.iter().map(|c| (c.column_name.as_str(), c)).collect();
390 let mut diffs = Vec::new();
391 for (name, src_col) in &src_map {
392 match tgt_map.get(name) {
393 None => diffs.push(ColumnDiff {
394 name: (*name).to_owned(),
395 status: DiffStatus::Removed,
396 before: Some((*src_col).clone()),
397 after: None,
398 }),
399 Some(tgt_col) => {
400 let changed = src_col.data_type != tgt_col.data_type
401 || src_col.not_null != tgt_col.not_null
402 || src_col.default_value.as_deref().unwrap_or("")
403 != tgt_col.default_value.as_deref().unwrap_or("");
404 diffs.push(ColumnDiff {
405 name: (*name).to_owned(),
406 status: if changed {
407 DiffStatus::Changed
408 } else {
409 DiffStatus::Unchanged
410 },
411 before: Some((*src_col).clone()),
412 after: Some((*tgt_col).clone()),
413 });
414 }
415 }
416 }
417 for (name, tgt_col) in &tgt_map {
418 if !src_map.contains_key(name) {
419 diffs.push(ColumnDiff {
420 name: (*name).to_owned(),
421 status: DiffStatus::Added,
422 before: None,
423 after: Some((*tgt_col).clone()),
424 });
425 }
426 }
427 diffs
428}
429
430fn diff_constraints(src: &[ConstraintSnapshot], tgt: &[ConstraintSnapshot]) -> Vec<ConstraintDiff> {
431 let src_map: std::collections::HashMap<&str, &ConstraintSnapshot> =
432 src.iter().map(|c| (c.name.as_str(), c)).collect();
433 let tgt_map: std::collections::HashMap<&str, &ConstraintSnapshot> =
434 tgt.iter().map(|c| (c.name.as_str(), c)).collect();
435 let mut diffs = Vec::new();
436 for (name, src_con) in &src_map {
437 match tgt_map.get(name) {
438 None => diffs.push(ConstraintDiff {
439 name: (*name).to_owned(),
440 status: DiffStatus::Removed,
441 before: Some((*src_con).clone()),
442 after: None,
443 }),
444 Some(tgt_con) => {
445 let changed = src_con.definition != tgt_con.definition;
446 diffs.push(ConstraintDiff {
447 name: (*name).to_owned(),
448 status: if changed {
449 DiffStatus::Changed
450 } else {
451 DiffStatus::Unchanged
452 },
453 before: Some((*src_con).clone()),
454 after: Some((*tgt_con).clone()),
455 });
456 }
457 }
458 }
459 for (name, tgt_con) in &tgt_map {
460 if !src_map.contains_key(name) {
461 diffs.push(ConstraintDiff {
462 name: (*name).to_owned(),
463 status: DiffStatus::Added,
464 before: None,
465 after: Some((*tgt_con).clone()),
466 });
467 }
468 }
469 diffs
470}
471
472fn diff_indexes(src: &[IndexSnapshot], tgt: &[IndexSnapshot]) -> Vec<IndexDiff> {
473 let src_map: std::collections::HashMap<&str, &IndexSnapshot> =
474 src.iter().map(|i| (i.name.as_str(), i)).collect();
475 let tgt_map: std::collections::HashMap<&str, &IndexSnapshot> =
476 tgt.iter().map(|i| (i.name.as_str(), i)).collect();
477 let mut diffs = Vec::new();
478 for (name, src_idx) in &src_map {
479 match tgt_map.get(name) {
480 None => diffs.push(IndexDiff {
481 name: (*name).to_owned(),
482 status: DiffStatus::Removed,
483 before: Some((*src_idx).clone()),
484 after: None,
485 }),
486 Some(tgt_idx) => {
487 let changed = src_idx.definition != tgt_idx.definition;
488 diffs.push(IndexDiff {
489 name: (*name).to_owned(),
490 status: if changed {
491 DiffStatus::Changed
492 } else {
493 DiffStatus::Unchanged
494 },
495 before: Some((*src_idx).clone()),
496 after: Some((*tgt_idx).clone()),
497 });
498 }
499 }
500 }
501 for (name, tgt_idx) in &tgt_map {
502 if !src_map.contains_key(name) {
503 diffs.push(IndexDiff {
504 name: (*name).to_owned(),
505 status: DiffStatus::Added,
506 before: None,
507 after: Some((*tgt_idx).clone()),
508 });
509 }
510 }
511 diffs
512}
513
514pub fn build_migration_statements(
516 source_schema: &str,
517 target_schema: &str,
518 diffs: &[TableDiff],
519) -> Vec<String> {
520 let mut stmts = Vec::new();
521 for table in diffs {
522 if table.status == DiffStatus::Unchanged {
523 continue;
524 }
525 if table.status == DiffStatus::Added {
526 let cols: Vec<String> = table
527 .column_diffs
528 .iter()
529 .filter(|c| c.status == DiffStatus::Added)
530 .filter_map(|c| c.after.as_ref())
531 .map(|c| {
532 let nn = if c.not_null { " NOT NULL" } else { "" };
533 let def = c
534 .default_value
535 .as_ref()
536 .map(|d| format!(" DEFAULT {d}"))
537 .unwrap_or_default();
538 format!(" \"{}\" {}{}{}", c.column_name, c.data_type, nn, def)
539 })
540 .collect();
541 stmts.push(format!(
542 "-- Table added in {target_schema}\nCREATE TABLE \"{source_schema}\".\"{}\" (\n{}\n);",
543 table.name,
544 cols.join(",\n")
545 ));
546 continue;
547 }
548 if table.status == DiffStatus::Removed {
549 stmts.push(format!(
550 "-- Table removed in {target_schema}\n-- DROP TABLE \"{source_schema}\".\"{}\"; -- Uncomment to drop",
551 table.name
552 ));
553 continue;
554 }
555
556 stmts.push(format!("-- Changes for table: {}", table.name));
557 for col in &table.column_diffs {
558 match col.status {
559 DiffStatus::Added => {
560 if let Some(after) = &col.after {
561 let nn = if after.not_null { " NOT NULL" } else { "" };
562 let def = after
563 .default_value
564 .as_ref()
565 .map(|d| format!(" DEFAULT {d}"))
566 .unwrap_or_default();
567 stmts.push(format!(
568 "ALTER TABLE \"{source_schema}\".\"{}\"\n ADD COLUMN \"{}\" {}{}{};",
569 table.name, col.name, after.data_type, nn, def
570 ));
571 }
572 }
573 DiffStatus::Removed => stmts.push(format!(
574 "-- ALTER TABLE \"{source_schema}\".\"{}\"\n-- DROP COLUMN \"{}\"; -- Uncomment to drop",
575 table.name, col.name
576 )),
577 DiffStatus::Changed => {
578 if let (Some(before), Some(after)) = (&col.before, &col.after) {
579 if before.data_type != after.data_type {
580 stmts.push(format!(
581 "ALTER TABLE \"{source_schema}\".\"{}\"\n ALTER COLUMN \"{}\" TYPE {};",
582 table.name, col.name, after.data_type
583 ));
584 }
585 if before.not_null != after.not_null {
586 let op = if after.not_null { "SET" } else { "DROP" };
587 stmts.push(format!(
588 "ALTER TABLE \"{source_schema}\".\"{}\"\n ALTER COLUMN \"{}\" {op} NOT NULL;",
589 table.name, col.name
590 ));
591 }
592 let before_def = before.default_value.as_deref().unwrap_or("");
593 let after_def = after.default_value.as_deref().unwrap_or("");
594 if before_def != after_def {
595 if after_def.is_empty() {
596 stmts.push(format!(
597 "ALTER TABLE \"{source_schema}\".\"{}\"\n ALTER COLUMN \"{}\" DROP DEFAULT;",
598 table.name, col.name
599 ));
600 } else {
601 stmts.push(format!(
602 "ALTER TABLE \"{source_schema}\".\"{}\"\n ALTER COLUMN \"{}\" SET DEFAULT {after_def};",
603 table.name, col.name
604 ));
605 }
606 }
607 }
608 }
609 DiffStatus::Unchanged => {}
610 }
611 }
612 for con in &table.constraint_diffs {
613 match con.status {
614 DiffStatus::Added => {
615 if let Some(after) = &con.after {
616 stmts.push(format!(
617 "ALTER TABLE \"{source_schema}\".\"{}\"\n ADD CONSTRAINT \"{}\" {};",
618 table.name, con.name, after.definition
619 ));
620 }
621 }
622 DiffStatus::Removed => stmts.push(format!(
623 "-- ALTER TABLE \"{source_schema}\".\"{}\"\n-- DROP CONSTRAINT \"{}\"; -- Uncomment to drop",
624 table.name, con.name
625 )),
626 _ => {}
627 }
628 }
629 for idx in &table.index_diffs {
630 match idx.status {
631 DiffStatus::Added => {
632 if let Some(after) = &idx.after {
633 let rewritten = after.definition.replace(
634 &format!("ON {target_schema}."),
635 &format!("ON {source_schema}."),
636 );
637 stmts.push(format!("{rewritten};"));
638 }
639 }
640 DiffStatus::Removed => {
641 stmts.push(format!(
642 "-- DROP INDEX \"{}\"; -- Uncomment to drop",
643 idx.name
644 ));
645 }
646 _ => {}
647 }
648 }
649 }
650 stmts
651}
652
653pub fn diffs_to_json(diffs: &[TableDiff]) -> Value {
654 json!(diffs)
655}
656
657#[cfg(test)]
658mod tests {
659 use super::*;
660
661 fn col(name: &str, ty: &str) -> ColumnSnapshot {
662 ColumnSnapshot {
663 column_name: name.into(),
664 data_type: ty.into(),
665 not_null: false,
666 default_value: None,
667 ordinal: 1,
668 }
669 }
670
671 #[test]
672 fn detects_added_table() {
673 let source = SchemaSnapshot { tables: vec![] };
674 let target = SchemaSnapshot {
675 tables: vec![TableSnapshot {
676 name: "orders".into(),
677 schema: "public".into(),
678 columns: vec![col("id", "integer")],
679 constraints: vec![],
680 indexes: vec![],
681 }],
682 };
683 let diffs = compute_schema_diff(&source, &target);
684 assert_eq!(diffs.len(), 1);
685 assert_eq!(diffs[0].status, DiffStatus::Added);
686 let stmts = build_migration_statements("public", "public", &diffs);
687 assert!(stmts[0].contains("CREATE TABLE"));
688 }
689
690 #[test]
691 fn detects_column_type_change() {
692 let source = SchemaSnapshot {
693 tables: vec![TableSnapshot {
694 name: "t".into(),
695 schema: "public".into(),
696 columns: vec![col("n", "integer")],
697 constraints: vec![],
698 indexes: vec![],
699 }],
700 };
701 let mut tgt_col = col("n", "bigint");
702 tgt_col.ordinal = 1;
703 let target = SchemaSnapshot {
704 tables: vec![TableSnapshot {
705 name: "t".into(),
706 schema: "public".into(),
707 columns: vec![tgt_col],
708 constraints: vec![],
709 indexes: vec![],
710 }],
711 };
712 let diffs = compute_schema_diff(&source, &target);
713 assert_eq!(diffs[0].status, DiffStatus::Changed);
714 let stmts = build_migration_statements("public", "public", &diffs);
715 assert!(stmts.iter().any(|s| s.contains("TYPE bigint")));
716 }
717}