systemprompt_database/lifecycle/migrations/
squash.rs

1//! Collapsing a contiguous range of applied migrations into a single
2//! version-0 baseline row.
3//!
4//! Copyright (c) systemprompt.io — Business Source License 1.1.
5//! See <https://systemprompt.io> for licensing details.
6
7use super::MigrationService;
8use std::collections::HashSet;
9use systemprompt_extension::{Extension, LoaderError, Migration};
10use tracing::info;
11
12#[derive(Debug, Clone)]
13pub struct SquashPlan {
14    pub extension_id: String,
15    pub through: u32,
16    pub baseline_name: String,
17    pub baseline_sql: String,
18    pub baseline_checksum: String,
19    pub source_versions: Vec<u32>,
20    pub already_applied_versions: Vec<u32>,
21    pub applied: bool,
22}
23
24fn baseline_checksum(sql: &str) -> String {
25    use std::hash::{Hash, Hasher};
26    let mut hasher = std::collections::hash_map::DefaultHasher::new();
27    sql.hash(&mut hasher);
28    format!("{:x}", hasher.finish())
29}
30
31fn collect_squash_range<'m>(
32    ext_id: &str,
33    migrations: &'m [Migration],
34    through: u32,
35) -> Result<Vec<&'m Migration>, LoaderError> {
36    if through == 0 {
37        return Err(LoaderError::MigrationFailed {
38            extension: ext_id.to_owned(),
39            message: "--through must be >= 1; version 0 is reserved for the squash baseline"
40                .to_owned(),
41        });
42    }
43
44    let to_squash: Vec<&Migration> = migrations
45        .iter()
46        .filter(|m| m.version >= 1 && m.version <= through)
47        .collect();
48
49    if to_squash.is_empty() {
50        return Err(LoaderError::MigrationFailed {
51            extension: ext_id.to_owned(),
52            message: format!(
53                "No migrations in range 1..={through} are defined for extension '{ext_id}'"
54            ),
55        });
56    }
57
58    let mut covered: Vec<u32> = to_squash.iter().map(|m| m.version).collect();
59    covered.sort_unstable();
60    if covered != (1..=through).collect::<Vec<u32>>() {
61        return Err(LoaderError::MigrationFailed {
62            extension: ext_id.to_owned(),
63            message: format!(
64                "Migrations 1..={through} are not contiguous for extension '{ext_id}': have \
65                 {covered:?}"
66            ),
67        });
68    }
69
70    Ok(to_squash)
71}
72
73fn build_baseline_sql(to_squash: &[&Migration]) -> String {
74    let mut baseline_sql = String::new();
75    for m in to_squash {
76        baseline_sql.push_str(&format!(
77            "-- migration {ver:03}: {name}\n",
78            ver = m.version,
79            name = m.name
80        ));
81        baseline_sql.push_str(m.sql);
82        if !m.sql.ends_with('\n') {
83            baseline_sql.push('\n');
84        }
85        baseline_sql.push('\n');
86    }
87    baseline_sql
88}
89
90impl MigrationService<'_> {
91    pub async fn squash_through(
92        &self,
93        extension: &dyn Extension,
94        through: u32,
95        apply: bool,
96    ) -> Result<SquashPlan, LoaderError> {
97        let ext_id = extension.metadata().id;
98
99        let mut migrations = extension.migrations();
100        migrations.sort_by_key(|m| m.version);
101        let to_squash = collect_squash_range(ext_id, &migrations, through)?;
102
103        let baseline_sql = build_baseline_sql(&to_squash);
104        let checksum = baseline_checksum(&baseline_sql);
105        let baseline_name = format!("baseline_v{through}");
106        let covered: Vec<u32> = to_squash.iter().map(|m| m.version).collect();
107
108        self.verify_range_applied(ext_id, through).await?;
109
110        let plan = SquashPlan {
111            extension_id: ext_id.to_owned(),
112            through,
113            baseline_name: baseline_name.clone(),
114            baseline_sql,
115            baseline_checksum: checksum.clone(),
116            source_versions: covered,
117            already_applied_versions: (1..=through).collect(),
118            applied: false,
119        };
120
121        if !apply {
122            return Ok(plan);
123        }
124
125        self.apply_squash_rows(ext_id, through, &baseline_name, &checksum)
126            .await?;
127
128        Ok(SquashPlan {
129            applied: true,
130            ..plan
131        })
132    }
133
134    async fn verify_range_applied(&self, ext_id: &str, through: u32) -> Result<(), LoaderError> {
135        self.ensure_migrations_table_exists().await?;
136        let applied = self.get_applied_migrations(ext_id).await?;
137        let applied_versions: HashSet<u32> = applied.iter().map(|m| m.version).collect();
138        let not_applied: Vec<u32> = (1..=through)
139            .filter(|v| !applied_versions.contains(v))
140            .collect();
141        if not_applied.is_empty() {
142            return Ok(());
143        }
144        Err(LoaderError::MigrationFailed {
145            extension: ext_id.to_owned(),
146            message: format!(
147                "Refusing to squash through {through}: extension '{ext_id}' has not applied \
148                 migrations {not_applied:?}. Squashing would retire them behind the baseline \
149                 without ever running them. Apply migrations 1..={through} first."
150            ),
151        })
152    }
153
154    async fn apply_squash_rows(
155        &self,
156        ext_id: &str,
157        through: u32,
158        baseline_name: &str,
159        checksum: &str,
160    ) -> Result<(), LoaderError> {
161        let baseline_id = format!("{ext_id}_000");
162
163        self.db
164            .execute(
165                &"INSERT INTO extension_migrations (id, extension_id, version, name, checksum) \
166                  VALUES ($1, $2, 0, $3, $4) ON CONFLICT (extension_id, version) DO UPDATE SET \
167                  name = EXCLUDED.name, checksum = EXCLUDED.checksum",
168                &[&baseline_id, &ext_id, &baseline_name, &checksum],
169            )
170            .await
171            .map_err(|e| LoaderError::MigrationFailed {
172                extension: ext_id.to_owned(),
173                message: format!("Failed to record baseline migration row: {e}"),
174            })?;
175
176        self.db
177            .execute(
178                &"DELETE FROM extension_migrations WHERE extension_id = $1 AND version BETWEEN 1 \
179                  AND $2",
180                &[&ext_id, &through],
181            )
182            .await
183            .map_err(|e| LoaderError::MigrationFailed {
184                extension: ext_id.to_owned(),
185                message: format!("Failed to retire squashed migration rows: {e}"),
186            })?;
187
188        info!(
189            extension = %ext_id,
190            through,
191            baseline_name = %baseline_name,
192            "Squash applied: baseline row inserted, source rows retired"
193        );
194
195        Ok(())
196    }
197}
systemprompt_database/lifecycle/migrations/squash.rs

systemprompt_database/lifecycle/migrations/
squash.rs
