use alloc::boxed::Box;
use alloc::string::{String, ToString};
use alloc::vec::Vec;
use core::fmt;
pub trait MigrationStep: Send + Sync {
fn source_version(&self) -> u32;
fn target_version(&self) -> u32;
fn migrate(&self, data: &[u8]) -> Result<Vec<u8>, MigrationError>;
}
#[derive(Debug, Clone, PartialEq)]
pub enum MigrationError {
NoPath { from: u32, to: u32 },
StepFailed { from: u32, to: u32, reason: String },
GapInChain { missing: u32 },
FutureVersion { found: u32, current: u32 },
Deserialization(String),
Serialization(String),
}
impl fmt::Display for MigrationError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::NoPath { from, to } => {
write!(f, "no migration path from v{from} to v{to}")
}
Self::StepFailed { from, to, reason } => {
write!(f, "migration v{from}→v{to} failed: {reason}")
}
Self::GapInChain { missing } => {
write!(f, "missing migration step for v{missing}")
}
Self::FutureVersion { found, current } => {
write!(f, "data version v{found} is newer than current v{current}")
}
Self::Deserialization(msg) => write!(f, "deserialization error: {msg}"),
Self::Serialization(msg) => write!(f, "serialization error: {msg}"),
}
}
}
#[derive(Debug, Clone)]
pub struct MigrationConfig {
pub write_back_on_read: bool,
pub eager_migration: bool,
}
impl Default for MigrationConfig {
fn default() -> Self {
Self {
write_back_on_read: true,
eager_migration: false,
}
}
}
pub struct MigrationEngine {
current_version: u32,
steps: Vec<Box<dyn MigrationStep>>,
}
impl MigrationEngine {
pub fn new(current_version: u32) -> Self {
Self {
current_version,
steps: Vec::new(),
}
}
pub fn register(&mut self, step: Box<dyn MigrationStep>) {
self.steps.push(step);
self.steps.sort_by_key(|s| s.source_version());
}
pub fn current_version(&self) -> u32 {
self.current_version
}
pub fn needs_migration(&self, data_version: u32) -> bool {
data_version != self.current_version
}
pub fn migrate_to_current(
&self,
data: &[u8],
from_version: u32,
) -> Result<Vec<u8>, MigrationError> {
if from_version == self.current_version {
return Ok(data.to_vec());
}
if from_version > self.current_version {
return Err(MigrationError::FutureVersion {
found: from_version,
current: self.current_version,
});
}
let mut current_data = data.to_vec();
let mut version = from_version;
while version < self.current_version {
let step = self
.steps
.iter()
.find(|s| s.source_version() == version)
.ok_or(MigrationError::GapInChain { missing: version })?;
current_data = step
.migrate(¤t_data)
.map_err(|e| MigrationError::StepFailed {
from: version,
to: step.target_version(),
reason: e.to_string(),
})?;
version = step.target_version();
}
Ok(current_data)
}
pub fn validate_chain(&self, min_version: u32) -> Result<(), MigrationError> {
let mut version = min_version;
while version < self.current_version {
if !self.steps.iter().any(|s| s.source_version() == version) {
return Err(MigrationError::GapInChain { missing: version });
}
version += 1;
}
Ok(())
}
pub fn registered_steps(&self) -> Vec<(u32, u32)> {
self.steps
.iter()
.map(|s| (s.source_version(), s.target_version()))
.collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
struct SimpleStep {
from: u32,
to: u32,
suffix: &'static [u8],
}
impl MigrationStep for SimpleStep {
fn source_version(&self) -> u32 {
self.from
}
fn target_version(&self) -> u32 {
self.to
}
fn migrate(&self, data: &[u8]) -> Result<Vec<u8>, MigrationError> {
let mut result = data.to_vec();
result.extend_from_slice(self.suffix);
Ok(result)
}
}
#[test]
fn no_migration_needed() {
let engine = MigrationEngine::new(1);
let data = b"hello";
let result = engine.migrate_to_current(data, 1).unwrap();
assert_eq!(result, b"hello");
}
#[test]
fn single_step_migration() {
let mut engine = MigrationEngine::new(2);
engine.register(Box::new(SimpleStep {
from: 1,
to: 2,
suffix: b"+v2",
}));
let result = engine.migrate_to_current(b"data", 1).unwrap();
assert_eq!(result, b"data+v2");
}
#[test]
fn multi_step_chain() {
let mut engine = MigrationEngine::new(4);
engine.register(Box::new(SimpleStep {
from: 1,
to: 2,
suffix: b"+v2",
}));
engine.register(Box::new(SimpleStep {
from: 2,
to: 3,
suffix: b"+v3",
}));
engine.register(Box::new(SimpleStep {
from: 3,
to: 4,
suffix: b"+v4",
}));
let result = engine.migrate_to_current(b"v1", 1).unwrap();
assert_eq!(result, b"v1+v2+v3+v4");
let result = engine.migrate_to_current(b"v2", 2).unwrap();
assert_eq!(result, b"v2+v3+v4");
let result = engine.migrate_to_current(b"v3", 3).unwrap();
assert_eq!(result, b"v3+v4");
}
#[test]
fn future_version_error() {
let engine = MigrationEngine::new(2);
let err = engine.migrate_to_current(b"data", 5).unwrap_err();
assert_eq!(
err,
MigrationError::FutureVersion {
found: 5,
current: 2
}
);
}
#[test]
fn gap_in_chain_error() {
let mut engine = MigrationEngine::new(3);
engine.register(Box::new(SimpleStep {
from: 1,
to: 2,
suffix: b"+v2",
}));
let err = engine.migrate_to_current(b"data", 1).unwrap_err();
assert_eq!(err, MigrationError::GapInChain { missing: 2 });
}
#[test]
fn validate_chain_ok() {
let mut engine = MigrationEngine::new(3);
engine.register(Box::new(SimpleStep {
from: 1,
to: 2,
suffix: b"",
}));
engine.register(Box::new(SimpleStep {
from: 2,
to: 3,
suffix: b"",
}));
assert!(engine.validate_chain(1).is_ok());
}
#[test]
fn validate_chain_gap() {
let mut engine = MigrationEngine::new(3);
engine.register(Box::new(SimpleStep {
from: 1,
to: 2,
suffix: b"",
}));
let err = engine.validate_chain(1).unwrap_err();
assert_eq!(err, MigrationError::GapInChain { missing: 2 });
}
#[test]
fn needs_migration() {
let engine = MigrationEngine::new(3);
assert!(engine.needs_migration(1));
assert!(engine.needs_migration(2));
assert!(!engine.needs_migration(3));
}
#[test]
fn registered_steps_list() {
let mut engine = MigrationEngine::new(3);
engine.register(Box::new(SimpleStep {
from: 2,
to: 3,
suffix: b"",
}));
engine.register(Box::new(SimpleStep {
from: 1,
to: 2,
suffix: b"",
}));
let steps = engine.registered_steps();
assert_eq!(steps, vec![(1, 2), (2, 3)]); }
}