adk_memory/migration.rs
1//! Lightweight, embedded migration runner for SQL-backed memory services.
2//!
3//! This module provides shared types and free functions that track applied
4//! schema versions in a per-backend registry table and execute only unapplied
5//! forward-only migration steps.
6//!
7//! The types ([`MigrationStep`], [`AppliedMigration`], [`MigrationError`]) are
8//! always compiled. The SQL runner functions ([`run_sql_migrations`],
9//! [`sql_schema_version`]) require the `sqlite-memory` or `database-memory` feature.
10
11use chrono::{DateTime, Utc};
12
13/// A single forward-only migration step.
14///
15/// The struct intentionally does not contain the SQL itself — each backend
16/// defines its own step list as `&[(i64, &str, &str)]` tuples of
17/// `(version, description, sql)`.
18#[derive(Debug, Clone, Copy)]
19pub struct MigrationStep {
20 /// Monotonically increasing version number, starting at 1.
21 pub version: i64,
22 /// Human-readable description of what this step does.
23 pub description: &'static str,
24}
25
26/// Record of an applied migration stored in the registry table.
27#[derive(Debug, Clone)]
28pub struct AppliedMigration {
29 /// The applied version number.
30 pub version: i64,
31 /// Description recorded at apply time.
32 pub description: String,
33 /// UTC timestamp of application.
34 pub applied_at: DateTime<Utc>,
35}
36
37/// Error context for a failed migration step.
38#[derive(Debug)]
39pub struct MigrationError {
40 /// The version that failed.
41 pub version: i64,
42 /// Description of the failed step.
43 pub description: String,
44 /// Underlying cause (database error message, etc.).
45 pub cause: String,
46}
47
48impl std::fmt::Display for MigrationError {
49 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50 write!(f, "migration v{} ({}) failed: {}", self.version, self.description, self.cause)
51 }
52}
53
54impl std::error::Error for MigrationError {}
55
56// ---------------------------------------------------------------------------
57// SQL runner — macro generates concrete implementations per database backend
58// ---------------------------------------------------------------------------
59
60/// Generates `run_sql_migrations` and `sql_schema_version` for a concrete
61/// sqlx pool type. Each SQL backend (`sqlite-memory`, `database-memory`) gets
62/// its own monomorphised copy, avoiding complex generic trait bounds.
63#[cfg(any(feature = "sqlite-memory", feature = "database-memory"))]
64macro_rules! impl_sql_migration_runner {
65 ($mod_name:ident, $pool_ty:ty) => {
66 pub mod $mod_name {
67 use super::MigrationError;
68 use chrono::Utc;
69 use sqlx::Row;
70 use std::future::Future;
71
72 /// Run all pending migrations for a SQL backend.
73 ///
74 /// 1. Creates the registry table if it does not exist.
75 /// 2. Calls `detect_existing` to check for pre-existing schema
76 /// tables. If tables exist but the registry is empty, records
77 /// version 1 as already applied (baseline).
78 /// 3. Reads the maximum applied version from the registry.
79 /// 4. If the database version exceeds the compiled-in maximum,
80 /// returns a version-mismatch error.
81 /// 5. Executes each unapplied step inside a transaction and
82 /// records it in the registry.
83 pub async fn run_sql_migrations<F, Fut>(
84 pool: &$pool_ty,
85 registry_table: &str,
86 steps: &[(i64, &str, &str)],
87 detect_existing: F,
88 ) -> Result<(), adk_core::AdkError>
89 where
90 F: FnOnce() -> Fut,
91 Fut: Future<Output = Result<bool, adk_core::AdkError>>,
92 {
93 // Step 1: Create registry table if missing
94 let create_sql = format!(
95 "CREATE TABLE IF NOT EXISTS {registry_table} (\
96 version INTEGER PRIMARY KEY, \
97 description TEXT NOT NULL, \
98 applied_at TEXT NOT NULL\
99 )"
100 );
101 sqlx::query(&create_sql).execute(pool).await.map_err(|e| {
102 adk_core::AdkError::Memory(format!("migration registry creation failed: {e}"))
103 })?;
104
105 // Step 2: Read current max applied version
106 let max_sql =
107 format!("SELECT COALESCE(MAX(version), 0) AS max_v FROM {registry_table}");
108 let row = sqlx::query(&max_sql).fetch_one(pool).await.map_err(|e| {
109 adk_core::AdkError::Memory(format!("migration registry read failed: {e}"))
110 })?;
111 let mut max_applied: i64 = row.try_get("max_v").map_err(|e| {
112 adk_core::AdkError::Memory(format!("migration registry read failed: {e}"))
113 })?;
114
115 // Step 3: Baseline detection — if registry is empty but
116 // tables already exist, record v1 as applied.
117 if max_applied == 0 {
118 let existing = detect_existing().await?;
119 if existing {
120 if let Some(&(v, desc, _)) = steps.first() {
121 let now = Utc::now().to_rfc3339();
122 let ins = format!(
123 "INSERT INTO {registry_table} \
124 (version, description, applied_at) \
125 VALUES ({v}, '{desc}', '{now}')"
126 );
127 sqlx::query(&ins).execute(pool).await.map_err(|e| {
128 adk_core::AdkError::Memory(format!(
129 "{}",
130 MigrationError {
131 version: v,
132 description: desc.to_string(),
133 cause: e.to_string(),
134 }
135 ))
136 })?;
137 max_applied = v;
138 }
139 }
140 }
141
142 // Step 4: Compiled-in max version
143 let max_compiled = steps.last().map(|s| s.0).unwrap_or(0);
144
145 // Step 5: Version mismatch check
146 if max_applied > max_compiled {
147 return Err(adk_core::AdkError::Memory(format!(
148 "schema version mismatch: database is at v{max_applied} \
149 but code only knows up to v{max_compiled}. \
150 Upgrade your ADK version."
151 )));
152 }
153
154 // Step 6: Execute unapplied steps in transactions
155 for &(version, description, sql) in steps {
156 if version <= max_applied {
157 continue;
158 }
159
160 let mut tx = pool.begin().await.map_err(|e| {
161 adk_core::AdkError::Memory(format!(
162 "{}",
163 MigrationError {
164 version,
165 description: description.to_string(),
166 cause: format!("transaction begin failed: {e}"),
167 }
168 ))
169 })?;
170
171 // Execute the migration SQL (raw_sql supports multiple
172 // semicolon-separated statements in a single call).
173 sqlx::raw_sql(sql).execute(&mut *tx).await.map_err(|e| {
174 adk_core::AdkError::Memory(format!(
175 "{}",
176 MigrationError {
177 version,
178 description: description.to_string(),
179 cause: e.to_string(),
180 }
181 ))
182 })?;
183
184 // Record the step in the registry
185 let now = Utc::now().to_rfc3339();
186 let rec = format!(
187 "INSERT INTO {registry_table} \
188 (version, description, applied_at) \
189 VALUES ({version}, '{description}', '{now}')"
190 );
191 sqlx::query(&rec).execute(&mut *tx).await.map_err(|e| {
192 adk_core::AdkError::Memory(format!(
193 "{}",
194 MigrationError {
195 version,
196 description: description.to_string(),
197 cause: format!("registry record failed: {e}"),
198 }
199 ))
200 })?;
201
202 tx.commit().await.map_err(|e| {
203 adk_core::AdkError::Memory(format!(
204 "{}",
205 MigrationError {
206 version,
207 description: description.to_string(),
208 cause: format!("transaction commit failed: {e}"),
209 }
210 ))
211 })?;
212 }
213
214 Ok(())
215 }
216
217 /// Returns the highest applied migration version, or 0 if no
218 /// registry table exists or the registry is empty.
219 pub async fn sql_schema_version(
220 pool: &$pool_ty,
221 registry_table: &str,
222 ) -> Result<i64, adk_core::AdkError> {
223 let sql =
224 format!("SELECT COALESCE(MAX(version), 0) AS max_v FROM {registry_table}");
225 match sqlx::query(&sql).fetch_one(pool).await {
226 Ok(row) => {
227 let version: i64 = row.try_get("max_v").unwrap_or(0);
228 Ok(version)
229 }
230 Err(_) => Ok(0),
231 }
232 }
233 }
234 };
235}
236
237#[cfg(feature = "sqlite-memory")]
238impl_sql_migration_runner!(sqlite_runner, sqlx::SqlitePool);
239
240#[cfg(feature = "database-memory")]
241impl_sql_migration_runner!(pg_runner, sqlx::PgPool);