use metaldb_derive::{BinaryValue, FromAccess};
use rand::{seq::SliceRandom, thread_rng, Rng};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use metaldb::{
access::{Access, CopyAccessExt, FromAccess, Prefixed},
migration::{flush_migration, Migration},
Database, Entry, Group, ListIndex, MapIndex, Snapshot, TemporaryDB,
};
const USER_COUNT: usize = 10_000;
pub type PublicKey = u16;
pub type Hash = u32;
pub mod v1 {
use super::*;
#[derive(Debug, Serialize, Deserialize, BinaryValue)]
#[binary_value(codec = "bincode")]
pub struct Wallet {
pub public_key: PublicKey, pub username: String,
pub balance: u32,
}
#[derive(Debug, FromAccess)]
pub struct Schema<T: Access> {
pub ticker: Entry<T::Base, String>,
pub divisibility: Entry<T::Base, u8>,
pub wallets: MapIndex<T::Base, PublicKey, Wallet>,
pub histories: Group<T, PublicKey, ListIndex<T::Base, Hash>>,
}
impl<T: Access> Schema<T> {
pub fn new(access: T) -> Self {
Self::from_root(access).unwrap()
}
pub fn print_wallets(&self) {
for (public_key, wallet) in self.wallets.iter().take(10) {
println!("Wallet[{:?}] = {:?}", public_key, wallet);
println!(
"History = {:?}",
self.histories.get(&public_key).iter().collect::<Vec<_>>()
);
}
}
}
}
fn create_initial_data() -> TemporaryDB {
let db = TemporaryDB::new();
let fork = db.fork();
{
const NAMES: &[&str] = &["Alice", "Bob", "Carol", "Dave", "Eve"];
let mut schema = v1::Schema::new(Prefixed::new("test", &fork));
schema.ticker.set("XNM".to_owned());
schema.divisibility.set(8);
let mut rng = thread_rng();
for user_id in 0..USER_COUNT {
let public_key = user_id as u16;
let username = (*NAMES.choose(&mut rng).unwrap()).to_string();
let wallet = v1::Wallet {
public_key,
username,
balance: rng.gen_range(0..1_000),
};
schema.wallets.put(&public_key, wallet);
let history_len = rng.gen_range(0..10);
schema
.histories
.get(&public_key)
.extend((0..history_len).map(|idx| idx as u32));
}
}
fork.get_list("unrelated.list").extend(vec![1, 2, 3]);
db.merge(fork.into_patch()).unwrap();
db
}
pub mod v2 {
use super::*;
#[derive(Debug, Serialize, Deserialize, BinaryValue)]
#[binary_value(codec = "bincode")]
pub struct Wallet {
pub username: String,
pub balance: u32,
pub history_hash: Hash, }
#[derive(Debug, Serialize, Deserialize, BinaryValue)]
#[binary_value(codec = "bincode")]
pub struct Config {
pub ticker: String,
pub divisibility: u8,
}
#[derive(Debug, FromAccess)]
pub struct Schema<T: Access> {
pub config: Entry<T::Base, Config>,
pub wallets: MapIndex<T::Base, PublicKey, Wallet>,
pub histories: Group<T, PublicKey, ListIndex<T::Base, Hash>>,
}
impl<T: Access> Schema<T> {
pub fn new(access: T) -> Self {
Self::from_root(access).unwrap()
}
pub fn print_wallets(&self) {
for (public_key, wallet) in self.wallets.iter().take(10) {
println!("Wallet[{:?}] = {:?}", public_key, wallet);
println!(
"History = {:?}",
self.histories.get(&public_key).iter().collect::<Vec<_>>()
);
}
}
}
}
fn check_data_before_flush(snapshot: &dyn Snapshot) {
let old_schema = v1::Schema::new(Prefixed::new("test", snapshot));
assert_eq!(old_schema.ticker.get().unwrap(), "XNM");
let new_schema = v2::Schema::new(Migration::new("test", snapshot));
assert_eq!(new_schema.config.get().unwrap().ticker, "XNM");
}
fn check_data_after_flush(snapshot: &dyn Snapshot) {
let new_schema = v2::Schema::new(Prefixed::new("test", snapshot));
assert_eq!(new_schema.config.get().unwrap().divisibility, 8);
assert!(!snapshot.get_entry::<_, u8>("test.divisibility").exists());
}
pub fn perform_migration<F>(migrate: F)
where
F: FnOnce(Arc<dyn Database>),
{
let db: Arc<dyn Database> = Arc::new(create_initial_data());
let fork = db.fork();
{
let old_data = Prefixed::new("test", fork.readonly());
let old_schema = v1::Schema::new(old_data.clone());
println!("Before migration:");
old_schema.print_wallets();
}
migrate(Arc::clone(&db));
let snapshot = db.snapshot();
check_data_before_flush(&snapshot);
let mut fork = db.fork();
flush_migration(&mut fork, "test");
let patch = fork.into_patch();
check_data_after_flush(&patch);
db.merge(patch).unwrap();
let snapshot = db.snapshot();
check_data_after_flush(&snapshot);
let schema = v2::Schema::new(Prefixed::new("test", &snapshot));
println!("After migration:");
schema.print_wallets();
}