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