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

systemprompt_database/lifecycle/migrations/
mod.rs
