systemprompt_database/lifecycle/migrations/
mod.rs

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

systemprompt_database/lifecycle/migrations/
mod.rs
