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                // Why: a tombstone's name labels the retirement ("retired_chain"),
141                // it is not the name of the migration that once held the slot, so
142                // comparing it to a tracked row is meaningless — and it failed on
143                // exactly the population tombstones exist for. Every established
144                // database carries the real names in a retired range, so slot
145                // identity was checked against a label and refused the boot.
146                debug!(
147                    extension = %ext_id,
148                    version = migration.version,
149                    name = %migration.name,
150                    tracked = row.is_some(),
151                    "Migration slot is tombstoned, nothing to run"
152                );
153                continue;
154            }
155
156            if let Some(row) = row {
157                self.verify_slot_identity(ext_id, migration, Some(row))?;
158                self.verify_checksum(ext_id, migration, Some(row.checksum.as_str()))?;
159                migrations_skipped += 1;
160                debug!(
161                    extension = %ext_id,
162                    version = migration.version,
163                    "Migration already applied, skipping"
164                );
165                continue;
166            }
167
168            self.execute_migration(extension, migration).await?;
169            migrations_run += 1;
170        }
171
172        if migrations_run > 0 {
173            info!(
174                extension = %ext_id,
175                migrations_run,
176                migrations_skipped,
177                "Migrations completed"
178            );
179        }
180
181        Ok(MigrationResult {
182            migrations_run,
183            migrations_skipped,
184        })
185    }
186
187    // Why: the recorded name is the only thing that distinguishes a migration
188    // edited in place from a slot whose file was deleted and its number reused.
189    // The checksum cannot tell them apart — it hashes the SQL alone.
190    async fn execute_migration(
191        &self,
192        extension: &dyn Extension,
193        migration: &Migration,
194    ) -> Result<(), LoaderError> {
195        let ext_id = extension.metadata().id;
196
197        check_cross_extension_alters(extension, migration)?;
198
199        info!(
200            extension = %ext_id,
201            version = migration.version,
202            name = %migration.name,
203            no_transaction = migration.no_transaction,
204            "Running migration"
205        );
206
207        let id = format!("{}_{:03}", ext_id, migration.version);
208        let checksum = migration.checksum();
209        let record_params: [&dyn ToDbValue; 5] =
210            [&id, &ext_id, &migration.version, &migration.name, &checksum];
211
212        if migration.no_transaction {
213            SqlExecutor::execute_statements_parsed(self.db, migration.sql)
214                .await
215                .map_err(|e| LoaderError::MigrationFailed {
216                    extension: ext_id.to_owned(),
217                    message: format!(
218                        "Failed to execute migration {} ({}): {e}",
219                        migration.version, migration.name
220                    ),
221                })?;
222            self.db
223                .execute(&RECORD_MIGRATION_SQL, &record_params)
224                .await
225                .map_err(|e| LoaderError::MigrationFailed {
226                    extension: ext_id.to_owned(),
227                    message: format!("Failed to record migration: {e}"),
228                })?;
229        } else {
230            let statements = SqlExecutor::parse_sql_statements(migration.sql).map_err(|e| {
231                LoaderError::MigrationFailed {
232                    extension: ext_id.to_owned(),
233                    message: format!(
234                        "Failed to parse migration {} ({}): {e}",
235                        migration.version, migration.name
236                    ),
237                }
238            })?;
239            execute_statements_transactional(
240                self.db,
241                &statements,
242                ext_id,
243                migration,
244                Some(TrackingWrite {
245                    sql: RECORD_MIGRATION_SQL,
246                    params: &record_params,
247                }),
248            )
249            .await?;
250        }
251
252        Ok(())
253    }
254}
255
256// Why: reported, never fatal. Databases predating tombstones carry rows for
257// every migration since deleted, and refusing to boot on those would strand
258// every established install. Adding the matching `.tombstone` file clears the
259// warning; `infra db migrate-status` lists the rows.
260pub(crate) fn orphaned_versions(applied: &[AppliedMigration], defined: &[Migration]) -> Vec<u32> {
261    let declared: HashSet<u32> = defined.iter().map(|m| m.version).collect();
262    applied
263        .iter()
264        .map(|m| m.version)
265        .filter(|version| !declared.contains(version))
266        .collect()
267}
268
269fn warn_orphaned_versions(ext_id: &str, applied: &[AppliedMigration], defined: &[Migration]) {
270    let orphaned = orphaned_versions(applied, defined);
271    if orphaned.is_empty() {
272        return;
273    }
274    warn!(
275        extension = %ext_id,
276        versions = ?orphaned,
277        "Applied migrations are no longer declared by the extension; their files were deleted \
278         without leaving a tombstone, so the numbers look free but are spent"
279    );
280}
systemprompt_database/lifecycle/migrations/mod.rs

systemprompt_database/lifecycle/migrations/
mod.rs
