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