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