Skip to main content

lean_ctx_patterns/
alembic.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Alembic (SQLAlchemy migrations) output compression.
4//!
5//! `alembic upgrade`/`downgrade` log each step as
6//! `INFO  [alembic.runtime.migration] Running upgrade <from> -> <to>, <msg>`
7//! plus boilerplate (`Context impl ...`, `Will assume transactional DDL`).
8//! We keep one compact line per applied revision and any error, dropping the
9//! logger prefix and boilerplate.
10
11use crate::ansi::strip_ansi;
12
13pub fn compress(_cmd: &str, output: &str) -> Option<String> {
14    let trimmed = output.trim();
15    if trimmed.is_empty() {
16        return Some("alembic: ok".to_string());
17    }
18
19    let mut upgrades: Vec<String> = Vec::new();
20    let mut downgrades: Vec<String> = Vec::new();
21    let mut errors: Vec<String> = Vec::new();
22
23    for raw in trimmed.lines() {
24        let stripped = strip_ansi(raw);
25        let line = stripped.trim();
26        if line.is_empty() {
27            continue;
28        }
29        let body = strip_log_prefix(line);
30
31        if let Some(rest) = body.strip_prefix("Running upgrade ") {
32            upgrades.push(parse_target(rest));
33        } else if let Some(rest) = body.strip_prefix("Running downgrade ") {
34            downgrades.push(parse_target(rest));
35        } else if is_error(line) {
36            errors.push(body.to_string());
37        }
38    }
39
40    if upgrades.is_empty() && downgrades.is_empty() && errors.is_empty() {
41        return Some(fallback(trimmed));
42    }
43
44    let mut parts: Vec<String> = Vec::new();
45    let mut header = String::from("alembic:");
46    if !upgrades.is_empty() {
47        header.push_str(&format!(" {} upgrade(s)", upgrades.len()));
48    }
49    if !downgrades.is_empty() {
50        header.push_str(&format!(" {} downgrade(s)", downgrades.len()));
51    }
52    if upgrades.is_empty() && downgrades.is_empty() {
53        header.push_str(" FAILED");
54    }
55    parts.push(header);
56    for u in &upgrades {
57        parts.push(format!("  ↑ {u}"));
58    }
59    for d in &downgrades {
60        parts.push(format!("  ↓ {d}"));
61    }
62    for e in errors.iter().take(5) {
63        parts.push(format!("  {e}"));
64    }
65    Some(parts.join("\n"))
66}
67
68/// Drop a leading `LEVEL  [alembic...] ` python-logging prefix.
69fn strip_log_prefix(line: &str) -> &str {
70    let is_log = line.starts_with("INFO")
71        || line.starts_with("WARNING")
72        || line.starts_with("ERROR")
73        || line.starts_with("DEBUG");
74    if is_log
75        && line.contains("[alembic")
76        && let Some(idx) = line.find("] ")
77    {
78        return line[idx + 2..].trim_start();
79    }
80    line
81}
82
83/// Parse `<from> -> <to>, <msg>` into `<to> <msg>` (msg optional).
84fn parse_target(rest: &str) -> String {
85    let after = rest.split("-> ").nth(1).unwrap_or(rest).trim();
86    match after.split_once(", ") {
87        Some((rev, msg)) => format!("{} {}", rev.trim(), msg.trim()),
88        None => after.to_string(),
89    }
90}
91
92fn is_error(line: &str) -> bool {
93    line.starts_with("ERROR")
94        || line.starts_with("FAILED")
95        || line.contains("Error:")
96        || line.contains("alembic.util.exc")
97        || line.contains("Can't locate revision")
98        || line.contains("Target database is not up to date")
99}
100
101fn fallback(text: &str) -> String {
102    let lines: Vec<&str> = text.lines().filter(|l| !l.trim().is_empty()).collect();
103    let n = lines.len().min(8);
104    let mut s = lines[..n].join("\n");
105    if lines.len() > n {
106        s.push_str(&format!("\n... (+{} lines)", lines.len() - n));
107    }
108    s
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    const UPGRADE: &str = "INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.\nINFO  [alembic.runtime.migration] Will assume transactional DDL.\nINFO  [alembic.runtime.migration] Running upgrade  -> a1b2c3, create users table\nINFO  [alembic.runtime.migration] Running upgrade a1b2c3 -> d4e5f6, add email index\n";
116
117    #[test]
118    fn keeps_revisions_drops_boilerplate() {
119        let r = compress("alembic upgrade head", UPGRADE).unwrap();
120        assert!(r.contains("2 upgrade(s)"), "counts upgrades: {r}");
121        assert!(r.contains("a1b2c3 create users table"), "{r}");
122        assert!(r.contains("d4e5f6 add email index"), "{r}");
123        assert!(!r.contains("transactional DDL"), "drops boilerplate: {r}");
124        assert!(!r.contains("Context impl"), "drops boilerplate: {r}");
125    }
126
127    #[test]
128    fn shorter_than_input() {
129        let r = compress("alembic upgrade head", UPGRADE).unwrap();
130        assert!(r.len() < UPGRADE.len());
131    }
132
133    #[test]
134    fn empty_is_ok() {
135        assert_eq!(compress("alembic upgrade head", "").unwrap(), "alembic: ok");
136    }
137
138    #[test]
139    fn surfaces_errors() {
140        let out = "INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.\nFAILED: Can't locate revision identified by 'deadbeef'\n";
141        let r = compress("alembic upgrade head", out).unwrap();
142        assert!(r.contains("FAILED"), "{r}");
143        assert!(r.contains("deadbeef"), "{r}");
144    }
145}