Skip to main content

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`] (whose
5//! rows the installer commits with the structural DDL they describe).
6//!
7//! Copyright (c) systemprompt.io — Business Source License 1.1.
8//! See <https://systemprompt.io> for licensing details.
9
10pub(crate) mod budget;
11mod checksum_transition;
12mod down;
13mod exec;
14mod mark_applied;
15mod repair;
16mod run;
17mod stamp;
18mod status;
19mod step_error;
20mod triggers;
21mod verify;
22
23pub use mark_applied::MarkAppliedOutcome;
24pub use repair::RepairResult;
25pub use stamp::{BaselineStamp, FreshnessCheck, is_retirement};
26pub use status::{
27    AppliedMigration, ChecksumDrift, ExtensionMigrationStatus, MigrationResult, MigrationStatus,
28    OrphanedMigration, PendingMigration, SlotCollision, TombstonedSlot,
29};
30
31use crate::services::{DatabaseProvider, SqlExecutor};
32use std::collections::HashSet;
33use systemprompt_extension::{Extension, LoaderError, Migration};
34use systemprompt_identifiers::ExtensionId;
35use tracing::{debug, info, warn};
36
37pub(crate) const RECORD_MIGRATION_SQL: &str = "INSERT INTO extension_migrations (id, extension_id, version, \
38                                    name, checksum) VALUES ($1, $2, $3, $4, $5)";
39
40#[derive(Debug, Default, Clone, Copy)]
41pub struct MigrationConfig {
42    pub allow_checksum_drift: bool,
43}
44
45pub struct MigrationService<'a> {
46    db: &'a dyn DatabaseProvider,
47    config: MigrationConfig,
48}
49
50impl std::fmt::Debug for MigrationService<'_> {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        f.debug_struct("MigrationService")
53            .field("config", &self.config)
54            .finish_non_exhaustive()
55    }
56}
57
58impl<'a> MigrationService<'a> {
59    pub fn new(db: &'a dyn DatabaseProvider) -> Self {
60        Self {
61            db,
62            config: MigrationConfig::default(),
63        }
64    }
65
66    #[must_use]
67    pub const fn with_config(mut self, config: MigrationConfig) -> Self {
68        self.config = config;
69        self
70    }
71
72    async fn ensure_migrations_table_exists(&self) -> Result<(), LoaderError> {
73        let sql = include_str!("../../../schema/extension_migrations.sql");
74        SqlExecutor::execute_statements_parsed(self.db, sql)
75            .await
76            .map_err(|e| LoaderError::MigrationStepFailed {
77                extension: ExtensionId::new("database"),
78                context: "Failed to ensure migrations table exists".to_owned(),
79                source: Box::new(e),
80            })
81    }
82
83    pub async fn get_applied_migrations(
84        &self,
85        extension_id: &ExtensionId,
86    ) -> Result<Vec<AppliedMigration>, LoaderError> {
87        let result = self
88            .db
89            .query_raw_with(
90                &"SELECT extension_id, version, name, checksum, applied_at FROM \
91                  extension_migrations WHERE extension_id = $1 ORDER BY version",
92                &[&extension_id],
93            )
94            .await
95            .map_err(|e| LoaderError::MigrationStepFailed {
96                extension: extension_id.clone(),
97                context: "Failed to query applied migrations".to_owned(),
98                source: Box::new(e),
99            })?;
100
101        result
102            .rows
103            .iter()
104            .map(|row| decode_applied_row(extension_id, row))
105            .collect()
106    }
107
108    pub async fn run_pending_migrations(
109        &self,
110        extension: &dyn Extension,
111    ) -> Result<MigrationResult, LoaderError> {
112        let ext_id = &ExtensionId::new(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        self.transition_checksums(ext_id, &migrations, &applied)
123            .await?;
124        let applied_rows: std::collections::HashMap<u32, &AppliedMigration> =
125            applied.iter().map(|m| (m.version, m)).collect();
126
127        warn_orphaned_versions(ext_id, &applied, &migrations);
128
129        let mut migrations_run = 0;
130        let mut migrations_skipped = 0;
131
132        for migration in &migrations {
133            let row = applied_rows.get(&migration.version).copied();
134
135            if migration.tombstone {
136                debug!(
137                    extension = %ext_id,
138                    version = migration.version,
139                    name = %migration.name,
140                    tracked = row.is_some(),
141                    "Migration slot is tombstoned, nothing to run"
142                );
143                continue;
144            }
145
146            if let Some(row) = row {
147                self.verify_slot_identity(ext_id, migration, Some(row))?;
148                self.verify_checksum(ext_id, migration, &row.checksum)?;
149                migrations_skipped += 1;
150                debug!(
151                    extension = %ext_id,
152                    version = migration.version,
153                    "Migration already applied, skipping"
154                );
155                continue;
156            }
157
158            self.execute_migration(extension, migration).await?;
159            migrations_run += 1;
160        }
161
162        if migrations_run > 0 {
163            info!(
164                extension = %ext_id,
165                migrations_run,
166                migrations_skipped,
167                "Migrations completed"
168            );
169        }
170
171        Ok(MigrationResult {
172            migrations_run,
173            migrations_skipped,
174        })
175    }
176}
177
178pub(crate) fn orphaned_versions(applied: &[AppliedMigration], defined: &[Migration]) -> Vec<u32> {
179    let declared: HashSet<u32> = defined.iter().map(|m| m.version).collect();
180    applied
181        .iter()
182        .map(|m| m.version)
183        .filter(|version| !declared.contains(version))
184        .collect()
185}
186
187fn warn_orphaned_versions(
188    ext_id: &ExtensionId,
189    applied: &[AppliedMigration],
190    defined: &[Migration],
191) {
192    let orphaned = orphaned_versions(applied, defined);
193    if orphaned.is_empty() {
194        return;
195    }
196    warn!(
197        extension = %ext_id,
198        versions = ?orphaned,
199        "Applied migrations are no longer declared by the extension; their files were deleted \
200         without leaving a tombstone, so the numbers look free but are spent"
201    );
202}
203
204fn decode_applied_row(
205    extension_id: &ExtensionId,
206    row: &crate::models::JsonRow,
207) -> Result<AppliedMigration, LoaderError> {
208    let malformed = |column: &str| LoaderError::MigrationFailed {
209        extension: extension_id.clone(),
210        message: format!("extension_migrations row has a malformed `{column}` column"),
211    };
212    let text = |column: &str| -> Result<String, LoaderError> {
213        row.get(column)
214            .and_then(serde_json::Value::as_str)
215            .map(str::to_owned)
216            .ok_or_else(|| malformed(column))
217    };
218    let version = row
219        .get("version")
220        .and_then(serde_json::Value::as_i64)
221        .and_then(|v| u32::try_from(v).ok())
222        .ok_or_else(|| malformed("version"))?;
223    let checksum = text("checksum")?;
224    Ok(AppliedMigration {
225        extension_id: extension_id.clone(),
226        version,
227        name: text("name")?,
228        checksum,
229        applied_at: row
230            .get("applied_at")
231            .and_then(serde_json::Value::as_str)
232            .map(str::to_owned),
233    })
234}