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