adk_session/migration.rs
1//! Lightweight, embedded migration runner for SQL-backed session 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` or `postgres` 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`, `postgres`) gets its own
62/// monomorphised copy, avoiding complex generic trait bounds.
63#[cfg(any(feature = "sqlite", feature = "postgres"))]
64macro_rules! impl_sql_migration_runner {
65 ($mod_name:ident, $pool_ty:ty, $int_type:expr) => {
66 /// SQL migration runner for this database backend.
67 pub mod $mod_name {
68 use super::MigrationError;
69 use chrono::Utc;
70 use sqlx::Row;
71 use std::future::Future;
72
73 /// Run all pending migrations for a SQL backend.
74 ///
75 /// 1. Creates the registry table if it does not exist.
76 /// 2. Calls `detect_existing` to check for pre-existing schema
77 /// tables. If tables exist but the registry is empty, records
78 /// version 1 as already applied (baseline).
79 /// 3. Reads the maximum applied version from the registry.
80 /// 4. If the database version exceeds the compiled-in maximum,
81 /// returns a version-mismatch error.
82 /// 5. Executes each unapplied step inside a transaction and
83 /// records it in the registry.
84 pub async fn run_sql_migrations<F, Fut>(
85 pool: &$pool_ty,
86 registry_table: &str,
87 steps: &[(i64, &str, &str)],
88 detect_existing: F,
89 ) -> Result<(), adk_core::AdkError>
90 where
91 F: FnOnce() -> Fut,
92 Fut: Future<Output = Result<bool, adk_core::AdkError>>,
93 {
94 // Step 1: Create registry table if missing
95 let create_sql = format!(
96 "CREATE TABLE IF NOT EXISTS {registry_table} (\
97 version {} PRIMARY KEY, \
98 description TEXT NOT NULL, \
99 applied_at TEXT NOT NULL\
100 )",
101 $int_type
102 );
103 sqlx::query(&create_sql).execute(pool).await.map_err(|e| {
104 adk_core::AdkError::session(format!("migration registry creation failed: {e}"))
105 })?;
106
107 // Step 2: Read current max applied version
108 let max_sql =
109 format!("SELECT COALESCE(MAX(version), 0) AS max_v FROM {registry_table}");
110 let row = sqlx::query(&max_sql).fetch_one(pool).await.map_err(|e| {
111 adk_core::AdkError::session(format!("migration registry read failed: {e}"))
112 })?;
113 let mut max_applied: i64 = row.try_get("max_v").map_err(|e| {
114 adk_core::AdkError::session(format!("migration registry read failed: {e}"))
115 })?;
116
117 // Step 3: Baseline detection — if registry is empty but
118 // tables already exist, record v1 as applied.
119 if max_applied == 0 {
120 let existing = detect_existing().await?;
121 if existing {
122 if let Some(&(v, desc, _)) = steps.first() {
123 let now = Utc::now().to_rfc3339();
124 let ins = format!(
125 "INSERT INTO {registry_table} \
126 (version, description, applied_at) \
127 VALUES ({v}, '{desc}', '{now}')"
128 );
129 sqlx::query(&ins).execute(pool).await.map_err(|e| {
130 adk_core::AdkError::session(format!(
131 "{}",
132 MigrationError {
133 version: v,
134 description: desc.to_string(),
135 cause: e.to_string(),
136 }
137 ))
138 })?;
139 max_applied = v;
140 }
141 }
142 }
143
144 // Step 4: Compiled-in max version
145 let max_compiled = steps.last().map(|s| s.0).unwrap_or(0);
146
147 // Step 5: Version mismatch check
148 if max_applied > max_compiled {
149 return Err(adk_core::AdkError::session(format!(
150 "schema version mismatch: database is at v{max_applied} \
151 but code only knows up to v{max_compiled}. \
152 Upgrade your ADK version."
153 )));
154 }
155
156 // Step 6: Execute unapplied steps in transactions
157 for &(version, description, sql) in steps {
158 if version <= max_applied {
159 continue;
160 }
161
162 let mut tx = pool.begin().await.map_err(|e| {
163 adk_core::AdkError::session(format!(
164 "{}",
165 MigrationError {
166 version,
167 description: description.to_string(),
168 cause: format!("transaction begin failed: {e}"),
169 }
170 ))
171 })?;
172
173 // Execute the migration SQL (raw_sql supports multiple
174 // semicolon-separated statements in a single call).
175 sqlx::raw_sql(sql).execute(&mut *tx).await.map_err(|e| {
176 adk_core::AdkError::session(format!(
177 "{}",
178 MigrationError {
179 version,
180 description: description.to_string(),
181 cause: e.to_string(),
182 }
183 ))
184 })?;
185
186 // Record the step in the registry
187 let now = Utc::now().to_rfc3339();
188 let rec = format!(
189 "INSERT INTO {registry_table} \
190 (version, description, applied_at) \
191 VALUES ({version}, '{description}', '{now}')"
192 );
193 sqlx::query(&rec).execute(&mut *tx).await.map_err(|e| {
194 adk_core::AdkError::session(format!(
195 "{}",
196 MigrationError {
197 version,
198 description: description.to_string(),
199 cause: format!("registry record failed: {e}"),
200 }
201 ))
202 })?;
203
204 tx.commit().await.map_err(|e| {
205 adk_core::AdkError::session(format!(
206 "{}",
207 MigrationError {
208 version,
209 description: description.to_string(),
210 cause: format!("transaction commit failed: {e}"),
211 }
212 ))
213 })?;
214 }
215
216 Ok(())
217 }
218
219 /// Returns the highest applied migration version, or 0 if no
220 /// registry table exists or the registry is empty.
221 pub async fn sql_schema_version(
222 pool: &$pool_ty,
223 registry_table: &str,
224 ) -> Result<i64, adk_core::AdkError> {
225 let sql =
226 format!("SELECT COALESCE(MAX(version), 0) AS max_v FROM {registry_table}");
227 match sqlx::query(&sql).fetch_one(pool).await {
228 Ok(row) => {
229 let version: i64 = row.try_get("max_v").unwrap_or(0);
230 Ok(version)
231 }
232 Err(_) => Ok(0),
233 }
234 }
235 }
236 };
237}
238
239#[cfg(feature = "sqlite")]
240impl_sql_migration_runner!(sqlite_runner, sqlx::SqlitePool, "INTEGER");
241
242#[cfg(feature = "postgres")]
243impl_sql_migration_runner!(pg_runner, sqlx::PgPool, "BIGINT");