Binary safe database import.

This commit is contained in:
chriseth 2022-12-20 20:40:43 +01:00
parent 60d294c6c8
commit 9fe536a853
3 changed files with 102 additions and 43 deletions

View file

@ -1,61 +1,60 @@
use std::env; use std::env;
use pathfinder2::{ use pathfinder2::io::*;
io::{read_edges_binary, read_edges_csv, write_edges_binary, write_edges_csv}, use pathfinder2::safe_db::safes_json::import_from_safes_json;
safe_db::safes_json::import_from_safes_json,
};
fn main() { fn main() {
let operation = env::args().nth(1).and_then(|op| { let input_format = env::args().nth(1).and_then(|op| {
if matches!( if matches!(
op.as_str(), op.as_str(),
"--safes-json-to-edges-bin" "--safes-json" | "--safes-bin" | "--edges-csv" | "--edges-bin"
| "--safes-json-to-edges-csv"
| "--edges-csv-to-edges-bin"
| "--edges-bin-to-edges-csv"
) { ) {
Some(op) Some(op)
} else { } else {
None None
} }
}); });
if env::args().len() != 4 || operation.is_none() { let output_format = env::args().nth(3).and_then(|op| {
println!("Usage: convert --safes-json-to-edges-bin <safes.json> <edges.dat>"); if matches!(op.as_str(), "--edges-csv" | "--edges-bin") {
println!("Usage: convert --safes-json-to-edges-csv <safes.json> <edges.csv>"); Some(op)
println!("Usage: convert --edges-csv-to-edges-bin <edges.csv> <edges.dat>"); } else {
println!("Usage: convert --edges-bin-to-edges-csv <edges.dat> <edges.csv>"); None
}
});
if env::args().len() != 5 || input_format.is_none() || output_format.is_none() {
println!("Usage: convert <input> <input_file> <output> <output_file>");
println!(" Where <input> is one of:");
println!(" --safes-json");
println!(" --safes-bin");
println!(" --edges-csv");
println!(" --edges-bin");
println!(" and <output>is one of:");
println!(" --edges-csv");
println!(" --edges-bin");
return; return;
} }
let input = env::args().nth(2).unwrap(); let input_file = env::args().nth(2).unwrap();
let output = env::args().nth(3).unwrap(); let edges = match input_format.unwrap().as_str() {
match operation.unwrap().as_str() { "--safes-json" => {
"--safes-json-to-edges-bin" => { let safes = import_from_safes_json(&input_file);
let safes = import_from_safes_json(&input); safes.edges().clone()
let edges = safes.edges();
println!("Imported {} edges.", edges.edge_count());
write_edges_binary(edges, &output).unwrap();
println!("Export done.");
} }
"--safes-json-to-edges-csv" => { "--safes-bin" => {
let safes = import_from_safes_json(&input); let safes = import_from_safes_binary(&input_file).unwrap();
let edges = safes.edges(); safes.edges().clone()
println!("Imported {} edges.", edges.edge_count());
write_edges_csv(edges, &output).unwrap();
println!("Export done.");
}
"--edges-csv-to-edges-bin" => {
let edges = &read_edges_csv(&input).unwrap();
println!("Imported {} edges.", edges.edge_count());
write_edges_binary(edges, &output).unwrap();
println!("Export done.");
}
"--edges-bin-to-edges-csv" => {
let edges = &read_edges_binary(&input).unwrap();
println!("Imported {} edges.", edges.edge_count());
write_edges_csv(edges, &output).unwrap();
println!("Export done.");
} }
"--edges-csv" => read_edges_csv(&input_file).unwrap(),
"--edges-bin" => read_edges_binary(&input_file).unwrap(),
_ => unreachable!(),
};
println!("Imported {} edges.", edges.edge_count());
let output_file = env::args().nth(4).unwrap();
match output_format.unwrap().as_str() {
"--edges-csv" => write_edges_csv(&edges, &output_file).unwrap(),
"--edges-bin" => write_edges_binary(&edges, &output_file).unwrap(),
_ => unreachable!(), _ => unreachable!(),
} }
println!("Export done.");
} }

View file

@ -1,11 +1,12 @@
use std::collections::BTreeSet; use std::collections::{BTreeMap, BTreeSet};
use std::fs::File; use std::fs::File;
use std::io::{self, BufRead}; use std::io::{self, BufRead};
use std::io::{Read, Write}; use std::io::{Read, Write};
use std::{collections::HashMap, io::BufReader}; use std::{collections::HashMap, io::BufReader};
use crate::safe_db::db::DB;
use crate::types::edge::EdgeDB; use crate::types::edge::EdgeDB;
use crate::types::{Address, Edge, U256}; use crate::types::{Address, Edge, Safe, U256};
pub fn read_edges_binary(path: &String) -> Result<EdgeDB, io::Error> { pub fn read_edges_binary(path: &String) -> Result<EdgeDB, io::Error> {
let mut f = File::open(path)?; let mut f = File::open(path)?;
@ -65,6 +66,63 @@ pub fn write_edges_csv(edges: &EdgeDB, path: &String) -> Result<(), io::Error> {
Ok(()) Ok(())
} }
pub fn import_from_safes_binary(path: &str) -> Result<DB, io::Error> {
let mut f = File::open(path)?;
let mut safes: BTreeMap<Address, Safe> = Default::default();
let address_index = read_address_index(&mut f)?;
// organizations
for _ in 0..read_u32(&mut f)? {
let org_address = read_address(&mut f, &address_index)?;
safes.entry(org_address).or_default().organization = true;
}
// trust edges
for _ in 0..read_u32(&mut f)? {
let user = read_address(&mut f, &address_index)?;
assert!(user != Address::default());
let send_to = read_address(&mut f, &address_index)?;
assert!(send_to != Address::default());
let limit_percentage = read_u8(&mut f)?;
assert!(limit_percentage <= 100);
if send_to != user && limit_percentage > 0 {
safes
.entry(user)
.or_default()
.limit_percentage
.insert(send_to, limit_percentage);
}
}
// balances
for _ in 0..read_u32(&mut f)? {
let user = read_address(&mut f, &address_index)?;
assert!(user != Address::default());
let token_owner = read_address(&mut f, &address_index)?;
assert!(token_owner != Address::default());
let balance = read_u256(&mut f)?;
if balance != U256::from(0) {
safes
.entry(user)
.or_default()
.balances
.insert(token_owner, balance);
}
}
// we use the safe address as token address
let mut token_owner = BTreeMap::default();
for (addr, safe) in &mut safes {
safe.token_address = *addr;
token_owner.insert(*addr, *addr);
}
Ok(DB::new(safes, token_owner))
}
fn read_address_index(file: &mut File) -> Result<HashMap<u32, Address>, io::Error> { fn read_address_index(file: &mut File) -> Result<HashMap<u32, Address>, io::Error> {
let address_count = read_u32(file)?; let address_count = read_u32(file)?;
let mut addresses = HashMap::new(); let mut addresses = HashMap::new();

View file

@ -4,6 +4,8 @@ use super::{Address, U256};
#[derive(Default, Debug)] #[derive(Default, Debug)]
pub struct Safe { pub struct Safe {
/// The address of the token, or the address of the safe if
/// the database does not use the distinction.
pub token_address: Address, pub token_address: Address,
pub balances: BTreeMap<Address, U256>, pub balances: BTreeMap<Address, U256>,
/// Limit percentage in "send to" direction /// Limit percentage in "send to" direction