use crate::server::{ServerError, ServerResult};
use serde_json::Value;
use std::collections::{HashMap, HashSet, VecDeque};
pub struct MigrationContext {
doc: Value,
}
impl MigrationContext {
pub fn new(doc: Value) -> Self {
Self { doc }
}
pub fn get(&self, key: &str) -> Option<&Value> {
self.doc.get(key)
}
pub fn set(&mut self, key: &str, value: Value) {
if let Value::Object(map) = &mut self.doc {
map.insert(key.to_owned(), value);
}
}
pub fn remove(&mut self, key: &str) -> Option<Value> {
if let Value::Object(map) = &mut self.doc {
map.remove(key)
} else {
None
}
}
pub fn into_doc(self) -> Value {
self.doc
}
pub fn doc(&self) -> &Value {
&self.doc
}
}
pub trait Migration: Send + Sync {
#[allow(clippy::wrong_self_convention)]
fn from_version(&self) -> (u64, u64, u64);
fn to_version(&self) -> (u64, u64, u64);
fn description(&self) -> &str;
fn migrate(&self, ctx: &mut MigrationContext) -> ServerResult<()>;
}
pub struct MigrationPlan<'a> {
migrations: &'a [Box<dyn Migration>],
indices: Vec<usize>,
}
impl<'a> std::fmt::Debug for MigrationPlan<'a> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("MigrationPlan")
.field("steps", &self.indices.len())
.finish()
}
}
impl<'a> MigrationPlan<'a> {
fn new(migrations: &'a [Box<dyn Migration>], indices: Vec<usize>) -> Self {
Self {
migrations,
indices,
}
}
pub fn is_empty(&self) -> bool {
self.indices.is_empty()
}
pub fn len(&self) -> usize {
self.indices.len()
}
pub fn apply(&self, ctx: &mut MigrationContext) -> ServerResult<()> {
for &idx in &self.indices {
self.migrations[idx].migrate(ctx)?;
}
Ok(())
}
pub fn step_descriptions(&self) -> Vec<&str> {
self.indices
.iter()
.map(|&idx| self.migrations[idx].description())
.collect()
}
}
pub struct MigrationRegistry {
migrations: Vec<Box<dyn Migration>>,
}
impl MigrationRegistry {
pub fn new() -> Self {
Self {
migrations: Vec::new(),
}
}
pub fn register(&mut self, m: impl Migration + 'static) -> &mut Self {
self.migrations.push(Box::new(m));
self
}
pub fn plan(
&self,
from: (u64, u64, u64),
to: (u64, u64, u64),
) -> ServerResult<MigrationPlan<'_>> {
if from == to {
return Ok(MigrationPlan::new(&self.migrations, vec![]));
}
let mut adj: HashMap<(u64, u64, u64), Vec<usize>> = HashMap::new();
for (idx, m) in self.migrations.iter().enumerate() {
adj.entry(m.from_version()).or_default().push(idx);
}
let mut queue: VecDeque<((u64, u64, u64), Vec<usize>)> = VecDeque::new();
let mut visited: HashSet<(u64, u64, u64)> = HashSet::new();
queue.push_back((from, vec![]));
visited.insert(from);
while let Some((cur, path)) = queue.pop_front() {
if let Some(neighbors) = adj.get(&cur) {
for &idx in neighbors {
let next = self.migrations[idx].to_version();
let mut new_path = path.clone();
new_path.push(idx);
if next == to {
return Ok(MigrationPlan::new(&self.migrations, new_path));
}
if !visited.contains(&next) {
visited.insert(next);
queue.push_back((next, new_path));
}
}
}
}
Err(ServerError::Migration(format!(
"No migration path from {from:?} to {to:?}"
)))
}
}
impl Default for MigrationRegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
struct SetFieldMigration {
from: (u64, u64, u64),
to: (u64, u64, u64),
desc: String,
field: String,
value: Value,
}
impl SetFieldMigration {
fn new(from: (u64, u64, u64), to: (u64, u64, u64), field: &str, value: Value) -> Self {
let desc = format!("{from:?} -> {to:?}: set {field}");
Self {
from,
to,
desc,
field: field.to_owned(),
value,
}
}
}
impl Migration for SetFieldMigration {
fn from_version(&self) -> (u64, u64, u64) {
self.from
}
fn to_version(&self) -> (u64, u64, u64) {
self.to
}
fn description(&self) -> &str {
&self.desc
}
fn migrate(&self, ctx: &mut MigrationContext) -> ServerResult<()> {
ctx.set(&self.field, self.value.clone());
Ok(())
}
}
#[test]
fn test_linear_path() {
let mut registry = MigrationRegistry::new();
registry
.register(SetFieldMigration::new(
(0, 1, 0),
(0, 2, 0),
"step1",
json!("applied"),
))
.register(SetFieldMigration::new(
(0, 2, 0),
(0, 3, 0),
"step2",
json!("applied"),
));
let plan = registry
.plan((0, 1, 0), (0, 3, 0))
.expect("plan should exist");
assert_eq!(plan.len(), 2);
let mut ctx = MigrationContext::new(json!({}));
plan.apply(&mut ctx).expect("apply should succeed");
let doc = ctx.into_doc();
assert_eq!(doc.get("step1"), Some(&json!("applied")));
assert_eq!(doc.get("step2"), Some(&json!("applied")));
}
#[test]
fn test_already_at_target_returns_empty_plan() {
let registry = MigrationRegistry::new();
let plan = registry
.plan((0, 2, 0), (0, 2, 0))
.expect("same-version plan");
assert!(plan.is_empty());
assert_eq!(plan.len(), 0);
}
#[test]
fn test_no_path_returns_migration_error() {
let registry = MigrationRegistry::new(); let result = registry.plan((0, 1, 0), (0, 9, 0));
assert!(
matches!(result, Err(ServerError::Migration(_))),
"expected Migration error, got: {result:?}"
);
}
#[test]
fn test_branch_picks_shortest_path() {
let mut registry = MigrationRegistry::new();
registry
.register(SetFieldMigration::new(
(0, 1, 0),
(0, 2, 0),
"intermediate",
json!(true),
))
.register(SetFieldMigration::new(
(0, 2, 0),
(0, 3, 0),
"via_intermediate",
json!(true),
))
.register(SetFieldMigration::new(
(0, 1, 0),
(0, 3, 0),
"direct",
json!(true),
));
let plan = registry
.plan((0, 1, 0), (0, 3, 0))
.expect("plan should exist");
assert_eq!(plan.len(), 1, "BFS should pick the direct 1-hop path");
let mut ctx = MigrationContext::new(json!({}));
plan.apply(&mut ctx).expect("apply should succeed");
let doc = ctx.into_doc();
assert_eq!(doc.get("direct"), Some(&json!(true)));
assert_eq!(
doc.get("intermediate"),
None,
"2-hop path must not be taken"
);
}
#[test]
fn test_cycle_terminates() {
let mut registry = MigrationRegistry::new();
registry
.register(SetFieldMigration::new(
(0, 1, 0),
(0, 2, 0),
"fwd",
json!(1),
))
.register(SetFieldMigration::new(
(0, 2, 0),
(0, 1, 0),
"back",
json!(2),
));
let result = registry.plan((0, 1, 0), (0, 3, 0));
assert!(
matches!(result, Err(ServerError::Migration(_))),
"expected Migration error for unreachable target, got: {result:?}"
);
}
#[test]
fn test_step_descriptions() {
let mut registry = MigrationRegistry::new();
registry
.register(SetFieldMigration::new(
(0, 1, 0),
(0, 2, 0),
"f",
json!(null),
))
.register(SetFieldMigration::new(
(0, 2, 0),
(0, 3, 0),
"g",
json!(null),
));
let plan = registry.plan((0, 1, 0), (0, 3, 0)).expect("plan");
let descs = plan.step_descriptions();
assert_eq!(descs.len(), 2);
for d in descs {
assert!(!d.is_empty());
}
}
#[test]
fn test_migration_context_set_get_remove() {
let mut ctx = MigrationContext::new(json!({"x": 1}));
assert_eq!(ctx.get("x"), Some(&json!(1)));
ctx.set("y", json!(2));
assert_eq!(ctx.get("y"), Some(&json!(2)));
ctx.remove("x");
assert_eq!(ctx.get("x"), None);
let doc = ctx.into_doc();
assert_eq!(doc.get("y"), Some(&json!(2)));
}
#[test]
fn test_registry_default() {
let registry = MigrationRegistry::default();
assert!(registry.migrations.is_empty());
}
}