systemprompt_database/lifecycle/migrations/
mod.rs1pub(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}