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