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
10mod checksum_transition;
11mod down;
12mod exec;
13mod mark_applied;
14mod repair;
15mod stamp;
16mod status;
17mod verify;
18
19pub use mark_applied::MarkAppliedOutcome;
20pub use repair::RepairResult;
21pub use stamp::{BaselineStamp, FreshnessCheck};
22pub use status::{
23    AppliedMigration, ChecksumDrift, ExtensionMigrationStatus, MigrationResult, MigrationStatus,
24    OrphanedMigration, PendingMigration, SlotCollision, TombstonedSlot,
25};
26
27use crate::services::{DatabaseProvider, SqlExecutor};
28use exec::{TrackingWrite, check_cross_extension_alters, execute_statements_transactional};
29use std::collections::HashSet;
30use systemprompt_extension::{Extension, LoaderError, Migration};
31use systemprompt_identifiers::ToDbValue;
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    async fn execute_migration(
173        &self,
174        extension: &dyn Extension,
175        migration: &Migration,
176    ) -> Result<(), LoaderError> {
177        let ext_id = extension.metadata().id;
178
179        check_cross_extension_alters(extension, migration)?;
180
181        info!(
182            extension = %ext_id,
183            version = migration.version,
184            name = %migration.name,
185            no_transaction = migration.no_transaction,
186            "Running migration"
187        );
188
189        let id = format!("{}_{:03}", ext_id, migration.version);
190        let checksum = migration.checksum();
191        let record_params: [&dyn ToDbValue; 5] =
192            [&id, &ext_id, &migration.version, &migration.name, &checksum];
193
194        if migration.no_transaction {
195            SqlExecutor::execute_statements_parsed(self.db, migration.sql)
196                .await
197                .map_err(|e| LoaderError::MigrationFailed {
198                    extension: ext_id.to_owned(),
199                    message: format!(
200                        "Failed to execute migration {} ({}): {e}",
201                        migration.version, migration.name
202                    ),
203                })?;
204            self.db
205                .execute(&RECORD_MIGRATION_SQL, &record_params)
206                .await
207                .map_err(|e| LoaderError::MigrationFailed {
208                    extension: ext_id.to_owned(),
209                    message: format!("Failed to record migration: {e}"),
210                })?;
211        } else {
212            let statements = SqlExecutor::parse_sql_statements(migration.sql).map_err(|e| {
213                LoaderError::MigrationFailed {
214                    extension: ext_id.to_owned(),
215                    message: format!(
216                        "Failed to parse migration {} ({}): {e}",
217                        migration.version, migration.name
218                    ),
219                }
220            })?;
221            execute_statements_transactional(
222                self.db,
223                &statements,
224                ext_id,
225                migration,
226                Some(TrackingWrite {
227                    sql: RECORD_MIGRATION_SQL,
228                    params: &record_params,
229                }),
230            )
231            .await?;
232        }
233
234        Ok(())
235    }
236}
237
238pub(crate) fn orphaned_versions(applied: &[AppliedMigration], defined: &[Migration]) -> Vec<u32> {
239    let declared: HashSet<u32> = defined.iter().map(|m| m.version).collect();
240    applied
241        .iter()
242        .map(|m| m.version)
243        .filter(|version| !declared.contains(version))
244        .collect()
245}
246
247fn warn_orphaned_versions(ext_id: &str, applied: &[AppliedMigration], defined: &[Migration]) {
248    let orphaned = orphaned_versions(applied, defined);
249    if orphaned.is_empty() {
250        return;
251    }
252    warn!(
253        extension = %ext_id,
254        versions = ?orphaned,
255        "Applied migrations are no longer declared by the extension; their files were deleted \
256         without leaving a tombstone, so the numbers look free but are spent"
257    );
258}
259
260fn decode_applied_row(
261    extension_id: &str,
262    row: &crate::models::JsonRow,
263) -> Result<AppliedMigration, LoaderError> {
264    let malformed = |column: &str| LoaderError::MigrationFailed {
265        extension: extension_id.to_owned(),
266        message: format!("extension_migrations row has a malformed `{column}` column"),
267    };
268    let text = |column: &str| -> Result<String, LoaderError> {
269        row.get(column)
270            .and_then(serde_json::Value::as_str)
271            .map(str::to_owned)
272            .ok_or_else(|| malformed(column))
273    };
274    let version = row
275        .get("version")
276        .and_then(serde_json::Value::as_i64)
277        .and_then(|v| u32::try_from(v).ok())
278        .ok_or_else(|| malformed("version"))?;
279    let checksum = text("checksum")?;
280    Ok(AppliedMigration {
281        extension_id: text("extension_id")?,
282        version,
283        name: text("name")?,
284        checksum,
285        applied_at: row
286            .get("applied_at")
287            .and_then(serde_json::Value::as_str)
288            .map(str::to_owned),
289    })
290}
systemprompt_database/lifecycle/migrations/mod.rs

systemprompt_database/lifecycle/migrations/
mod.rs
