systemprompt_database/lifecycle/migrations/
mod.rs

1//! Extension migration runner backed by the `extension_migrations`
2//! bookkeeping table. [`MigrationService`] applies, reverts, inspects, and
3//! squashes per-extension migration history; reverts live in [`down`],
4//! status/plan queries in [`status`], squash in [`squash`].
5//!
6//! Copyright (c) systemprompt.io — Business Source License 1.1.
7//! See <https://systemprompt.io> for licensing details.
8
9mod down;
10mod exec;
11mod mark_applied;
12mod repair;
13mod squash;
14mod status;
15
16pub use mark_applied::MarkAppliedOutcome;
17pub use repair::RepairResult;
18pub use squash::SquashPlan;
19pub use status::{
20    AppliedMigration, ChecksumDrift, ExtensionMigrationStatus, MigrationResult, MigrationStatus,
21    PendingMigration,
22};
23
24use crate::services::{DatabaseProvider, SqlExecutor};
25use exec::{check_cross_extension_alters, execute_statements_transactional};
26use std::collections::HashSet;
27use systemprompt_extension::{Extension, LoaderError, Migration};
28use tracing::{debug, info, warn};
29
30#[derive(Debug, Default, Clone, Copy)]
31pub struct MigrationConfig {
32    pub allow_checksum_drift: bool,
33}
34
35pub struct MigrationService<'a> {
36    db: &'a dyn DatabaseProvider,
37    config: MigrationConfig,
38}
39
40impl std::fmt::Debug for MigrationService<'_> {
41    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42        f.debug_struct("MigrationService")
43            .field("config", &self.config)
44            .finish_non_exhaustive()
45    }
46}
47
48impl<'a> MigrationService<'a> {
49    pub fn new(db: &'a dyn DatabaseProvider) -> Self {
50        Self {
51            db,
52            config: MigrationConfig::default(),
53        }
54    }
55
56    #[must_use]
57    pub const fn with_config(mut self, config: MigrationConfig) -> Self {
58        self.config = config;
59        self
60    }
61
62    async fn ensure_migrations_table_exists(&self) -> Result<(), LoaderError> {
63        let sql = include_str!("../../../schema/extension_migrations.sql");
64        SqlExecutor::execute_statements_parsed(self.db, sql)
65            .await
66            .map_err(|e| LoaderError::MigrationFailed {
67                extension: "database".to_owned(),
68                message: format!("Failed to ensure migrations table exists: {e}"),
69            })
70    }
71
72    pub async fn get_applied_migrations(
73        &self,
74        extension_id: &str,
75    ) -> Result<Vec<AppliedMigration>, LoaderError> {
76        let result = self
77            .db
78            .query_raw_with(
79                &"SELECT extension_id, version, name, checksum, applied_at FROM \
80                  extension_migrations WHERE extension_id = $1 ORDER BY version",
81                &[&extension_id],
82            )
83            .await
84            .map_err(|e| LoaderError::MigrationFailed {
85                extension: extension_id.to_owned(),
86                message: format!("Failed to query applied migrations: {e}"),
87            })?;
88
89        let migrations = result
90            .rows
91            .iter()
92            .filter_map(|row| {
93                Some(AppliedMigration {
94                    extension_id: row.get("extension_id")?.as_str()?.to_owned(),
95                    version: row.get("version")?.as_i64()? as u32,
96                    name: row.get("name")?.as_str()?.to_owned(),
97                    checksum: row.get("checksum")?.as_str()?.to_owned(),
98                    applied_at: row
99                        .get("applied_at")
100                        .and_then(|v| v.as_str().map(String::from)),
101                })
102            })
103            .collect();
104
105        Ok(migrations)
106    }
107
108    pub async fn run_pending_migrations(
109        &self,
110        extension: &dyn Extension,
111    ) -> Result<MigrationResult, LoaderError> {
112        let ext_id = extension.metadata().id;
113        let migrations = extension.migrations();
114
115        if migrations.is_empty() {
116            return Ok(MigrationResult::default());
117        }
118
119        self.ensure_migrations_table_exists().await?;
120
121        let applied = self.get_applied_migrations(ext_id).await?;
122        let applied_versions: HashSet<u32> = applied.iter().map(|m| m.version).collect();
123        let applied_checksums: std::collections::HashMap<u32, &str> = applied
124            .iter()
125            .map(|m| (m.version, m.checksum.as_str()))
126            .collect();
127
128        let mut migrations_run = 0;
129        let mut migrations_skipped = 0;
130
131        for migration in &migrations {
132            if applied_versions.contains(&migration.version) {
133                self.verify_checksum(
134                    ext_id,
135                    migration,
136                    applied_checksums.get(&migration.version).copied(),
137                )?;
138                migrations_skipped += 1;
139                debug!(
140                    extension = %ext_id,
141                    version = migration.version,
142                    "Migration already applied, skipping"
143                );
144                continue;
145            }
146
147            self.execute_migration(extension, migration).await?;
148            migrations_run += 1;
149        }
150
151        if migrations_run > 0 {
152            info!(
153                extension = %ext_id,
154                migrations_run,
155                migrations_skipped,
156                "Migrations completed"
157            );
158        }
159
160        Ok(MigrationResult {
161            migrations_run,
162            migrations_skipped,
163        })
164    }
165
166    fn verify_checksum(
167        &self,
168        ext_id: &str,
169        migration: &Migration,
170        stored: Option<&str>,
171    ) -> Result<(), LoaderError> {
172        let Some(stored_checksum) = stored else {
173            return Ok(());
174        };
175        let current_checksum = migration.checksum();
176        if stored_checksum == current_checksum {
177            return Ok(());
178        }
179        if self.config.allow_checksum_drift {
180            warn!(
181                extension = %ext_id,
182                version = migration.version,
183                name = %migration.name,
184                stored_checksum = %stored_checksum,
185                current_checksum = %current_checksum,
186                "Migration checksum mismatch tolerated by --allow-checksum-drift"
187            );
188            return Ok(());
189        }
190        Err(LoaderError::MigrationFailed {
191            extension: ext_id.to_owned(),
192            message: format!(
193                "Migration {ver} ('{name}') has been edited since it was applied (stored checksum \
194                 {stored_checksum}, current {current_checksum}). Refusing to proceed. Run \
195                 `systemprompt infra db migrate-repair --apply` to reconcile the tracking table, \
196                 or pass --allow-checksum-drift to bypass the check without fixing it.",
197                ver = migration.version,
198                name = migration.name,
199            ),
200        })
201    }
202
203    async fn execute_migration(
204        &self,
205        extension: &dyn Extension,
206        migration: &Migration,
207    ) -> Result<(), LoaderError> {
208        let ext_id = extension.metadata().id;
209
210        check_cross_extension_alters(extension, migration)?;
211
212        info!(
213            extension = %ext_id,
214            version = migration.version,
215            name = %migration.name,
216            no_transaction = migration.no_transaction,
217            "Running migration"
218        );
219
220        if migration.no_transaction {
221            SqlExecutor::execute_statements_parsed(self.db, migration.sql)
222                .await
223                .map_err(|e| LoaderError::MigrationFailed {
224                    extension: ext_id.to_owned(),
225                    message: format!(
226                        "Failed to execute migration {} ({}): {e}",
227                        migration.version, migration.name
228                    ),
229                })?;
230        } else {
231            let statements = SqlExecutor::parse_sql_statements(migration.sql).map_err(|e| {
232                LoaderError::MigrationFailed {
233                    extension: ext_id.to_owned(),
234                    message: format!(
235                        "Failed to parse migration {} ({}): {e}",
236                        migration.version, migration.name
237                    ),
238                }
239            })?;
240            execute_statements_transactional(self.db, &statements, ext_id, migration).await?;
241        }
242
243        self.record_migration(ext_id, migration).await?;
244
245        Ok(())
246    }
247
248    async fn record_migration(
249        &self,
250        ext_id: &str,
251        migration: &Migration,
252    ) -> Result<(), LoaderError> {
253        let id = format!("{}_{:03}", ext_id, migration.version);
254        let checksum = migration.checksum();
255
256        self.db
257            .execute(
258                &"INSERT INTO extension_migrations (id, extension_id, version, name, checksum) \
259                  VALUES ($1, $2, $3, $4, $5)",
260                &[&id, &ext_id, &migration.version, &migration.name, &checksum],
261            )
262            .await
263            .map_err(|e| LoaderError::MigrationFailed {
264                extension: ext_id.to_owned(),
265                message: format!("Failed to record migration: {e}"),
266            })?;
267
268        Ok(())
269    }
270}
systemprompt_database/lifecycle/migrations/mod.rs

systemprompt_database/lifecycle/migrations/
mod.rs
