diff --git a/README.md b/README.md index b2449e3aadb1fe..9c956382b853e1 100644 --- a/README.md +++ b/README.md @@ -17,7 +17,7 @@ primary_address,bid_amount_dollars ``` ```bash -solana-tokens distribute-tokens --from --dollars-per-sol --bids-csv --fee-payer +solana-tokens distribute-tokens --from --dollars-per-sol --input_csv --fee-payer ``` Example transaction log before: @@ -31,7 +31,7 @@ Send tokens to the recipients in `` if the distribution is not already recordered in the transaction log. ```bash -solana-tokens distribute-tokens --from --dollars-per-sol --bids-csv --fee-payer +solana-tokens distribute-tokens --from --dollars-per-sol --input_csv --fee-payer ``` Example output: @@ -60,7 +60,7 @@ List the differences between a list of expected distributions and the record of transactions have already been sent. ```bash -solana-tokens distribute-tokens --dollars-per-sol --dry-run --bids-csv +solana-tokens distribute-tokens --dollars-per-sol --dry-run --input_csv ``` Example bids.csv: diff --git a/src/arg_parser.rs b/src/arg_parser.rs index 96ebe02d1b6807..47f382b73fea19 100644 --- a/src/arg_parser.rs +++ b/src/arg_parser.rs @@ -42,12 +42,17 @@ where .help("Transactions database file"), ) .arg( - Arg::with_name("bids_csv") - .long("bids-csv") + Arg::with_name("from_bids") + .long("from-bids") + .help("Input CSV contains bids in dollars, not allocations in SOL"), + ) + .arg( + Arg::with_name("input_csv") + .long("input_csv") .required(true) .takes_value(true) .value_name("FILE") - .help("Bids CSV file"), + .help("Input CSV file"), ) .arg( Arg::with_name("dollars_per_sol") @@ -55,7 +60,7 @@ where .required(true) .takes_value(true) .value_name("NUMBER") - .help("Dollars per SOL"), + .help("Dollars per SOL, if input CSV contains bids") ) .arg( Arg::with_name("dry_run") @@ -149,8 +154,8 @@ where SubCommand::with_name("balances") .about("Balance of each account") .arg( - Arg::with_name("bids_csv") - .long("bids-csv") + Arg::with_name("input_csv") + .long("input_csv") .required(true) .takes_value(true) .value_name("FILE") @@ -170,9 +175,10 @@ where fn parse_distribute_tokens_args(matches: &ArgMatches<'_>) -> DistributeTokensArgs { DistributeTokensArgs { - bids_csv: value_t_or_exit!(matches, "bids_csv", String), + input_csv: value_t_or_exit!(matches, "input_csv", String), + from_bids: matches.is_present("from_bids"), transactions_db: value_t_or_exit!(matches, "transactions_db", String), - dollars_per_sol: value_t_or_exit!(matches, "dollars_per_sol", f64), + dollars_per_sol: value_t!(matches, "dollars_per_sol", f64).ok(), dry_run: matches.is_present("dry_run"), sender_keypair: value_t!(matches, "sender_keypair", String).ok(), fee_payer: value_t!(matches, "fee_payer", String).ok(), @@ -194,7 +200,7 @@ fn parse_distribute_stake_args(matches: &ArgMatches<'_>) -> DistributeStakeArgs< fn parse_balances_args(matches: &ArgMatches<'_>) -> BalancesArgs { BalancesArgs { - bids_csv: value_t_or_exit!(matches, "bids_csv", String), + input_csv: value_t_or_exit!(matches, "input_csv", String), dollars_per_sol: value_t_or_exit!(matches, "dollars_per_sol", f64), } } diff --git a/src/args.rs b/src/args.rs index f918535e0f519e..ea913681cd595c 100644 --- a/src/args.rs +++ b/src/args.rs @@ -5,9 +5,10 @@ use solana_sdk::{pubkey::Pubkey, signature::Signer}; use std::error::Error; pub struct DistributeTokensArgs { - pub bids_csv: String, + pub input_csv: String, + pub from_bids: bool, pub transactions_db: String, - pub dollars_per_sol: f64, + pub dollars_per_sol: Option, pub dry_run: bool, pub sender_keypair: Option, pub fee_payer: Option, @@ -25,7 +26,7 @@ pub struct DistributeStakeArgs { } pub struct BalancesArgs { - pub bids_csv: String, + pub input_csv: String, pub dollars_per_sol: f64, } @@ -49,7 +50,8 @@ pub fn resolve_command( let mut wallet_manager = maybe_wallet_manager()?; let matches = ArgMatches::default(); let resolved_args = DistributeTokensArgs { - bids_csv: args.bids_csv, + input_csv: args.input_csv, + from_bids: args.from_bids, transactions_db: args.transactions_db, dollars_per_sol: args.dollars_per_sol, dry_run: args.dry_run, diff --git a/src/tokens.rs b/src/tokens.rs index ab9462481408ba..7ad7a8c3336894 100644 --- a/src/tokens.rs +++ b/src/tokens.rs @@ -272,26 +272,42 @@ fn set_transaction_info( Ok(()) } +fn read_allocations( + args: &DistributeTokensArgs> +) -> Vec { + let rdr = ReaderBuilder::new() + .trim(Trim::All) + .from_path(&args.input_csv); + if args.from_bids { + let bids: Vec = rdr.unwrap().deserialize().map(|bid| bid.unwrap()).collect(); + bids + .into_iter() + .map(|bid| create_allocation(&bid, args.dollars_per_sol.unwrap())) + .collect() + } else { + rdr.unwrap().deserialize().map(|entry| entry.unwrap()).collect() + } +} + pub fn process_distribute_tokens( client: &ThinClient, args: &DistributeTokensArgs>, ) -> Result<(), Error> { - let mut rdr = ReaderBuilder::new() - .trim(Trim::All) - .from_path(&args.bids_csv)?; - let bids: Vec = rdr.deserialize().map(|bid| bid.unwrap()).collect(); - let mut allocations: Vec = bids - .into_iter() - .map(|bid| create_allocation(&bid, args.dollars_per_sol)) - .collect(); + let mut allocations: Vec = read_allocations(&args); let starting_total_tokens: f64 = allocations.iter().map(|x| x.amount).sum(); println!( - "{} ◎{} (${})", + "{} ◎{}", style(format!("{}", "Total in allocations_csv:")).bold(), starting_total_tokens, - starting_total_tokens * args.dollars_per_sol, ); + if let Some(dollars_per_sol) = args.dollars_per_sol { + println!( + "{} ${}", + style(format!("{}", "Total in allocations_csv:")).bold(), + starting_total_tokens * dollars_per_sol, + ); + } let mut db = open_db(&args.transactions_db)?; let transaction_infos = read_transaction_infos(&db); @@ -335,23 +351,41 @@ pub fn process_distribute_tokens( let distributed_tokens: f64 = transaction_infos.iter().map(|x| x.amount).sum(); let undistributed_tokens: f64 = allocations.iter().map(|x| x.amount).sum(); println!( - "{} ◎{} (${})", + "{} ◎{}", style(format!("{}", "Distributed:")).bold(), distributed_tokens, - distributed_tokens * args.dollars_per_sol, ); + if let Some(dollars_per_sol) = args.dollars_per_sol { + println!( + "{} ${}", + style(format!("{}", "Distributed:")).bold(), + distributed_tokens * dollars_per_sol, + ); + } println!( - "{} ◎{} (${})", + "{} ◎{}", style(format!("{}", "Undistributed:")).bold(), undistributed_tokens, - undistributed_tokens * args.dollars_per_sol, ); + if let Some(dollars_per_sol) = args.dollars_per_sol { + println!( + "{} ${}", + style(format!("{}", "Undistributed:")).bold(), + undistributed_tokens * dollars_per_sol, + ); + } println!( - "{} ◎{} (${})\n", + "{} ◎{}", style(format!("{}", "Total:")).bold(), distributed_tokens + undistributed_tokens, - (distributed_tokens + undistributed_tokens) * args.dollars_per_sol, ); + if let Some(dollars_per_sol) = args.dollars_per_sol { + println!( + "{} ${}", + style(format!("{}", "Total:")).bold(), + (distributed_tokens + undistributed_tokens) * dollars_per_sol, + ); + } distribute_tokens(client, &mut db, &allocations, args)?; @@ -391,7 +425,7 @@ pub fn process_balances( ) -> Result<(), csv::Error> { let mut rdr = ReaderBuilder::new() .trim(Trim::All) - .from_path(&args.bids_csv)?; + .from_path(&args.input_csv)?; let bids: Vec = rdr.deserialize().map(|bid| bid.unwrap()).collect(); let allocations: Vec = bids .into_iter() @@ -426,7 +460,7 @@ pub fn process_balances( use solana_sdk::{pubkey::Pubkey, signature::Keypair}; use tempfile::{tempdir, NamedTempFile}; -pub fn test_process_distribute_with_client(client: C, sender_keypair: Keypair) { +pub fn test_process_distribute_bids_with_client(client: C, sender_keypair: Keypair) { let thin_client = ThinClient(client); let fee_payer = Keypair::new(); thin_client @@ -439,7 +473,7 @@ pub fn test_process_distribute_with_client(client: C, sender_keypair: accepted_amount_dollars: 1000.0, }; let bids_file = NamedTempFile::new().unwrap(); - let bids_csv = bids_file.path().to_str().unwrap().to_string(); + let input_csv = bids_file.path().to_str().unwrap().to_string(); let mut wtr = csv::WriterBuilder::new().from_writer(bids_file); wtr.serialize(&bid).unwrap(); wtr.flush().unwrap(); @@ -456,15 +490,16 @@ pub fn test_process_distribute_with_client(client: C, sender_keypair: sender_keypair: Some(Box::new(sender_keypair)), fee_payer: Some(Box::new(fee_payer)), dry_run: false, - bids_csv, + from_bids: true, + input_csv, transactions_db: transactions_db.clone(), - dollars_per_sol: 0.22, + dollars_per_sol: Some(0.22), }; process_distribute_tokens(&thin_client, &args).unwrap(); let transaction_infos = read_transaction_infos(&open_db(&transactions_db).unwrap()); assert_eq!(transaction_infos.len(), 1); assert_eq!(transaction_infos[0].recipient, alice_pubkey.to_string()); - let expected_amount = bid.accepted_amount_dollars / args.dollars_per_sol; + let expected_amount = bid.accepted_amount_dollars / args.dollars_per_sol.unwrap(); assert_eq!(transaction_infos[0].amount, expected_amount); assert_eq!( @@ -477,7 +512,68 @@ pub fn test_process_distribute_with_client(client: C, sender_keypair: let transaction_infos = read_transaction_infos(&open_db(&transactions_db).unwrap()); assert_eq!(transaction_infos.len(), 1); assert_eq!(transaction_infos[0].recipient, alice_pubkey.to_string()); - let expected_amount = bid.accepted_amount_dollars / args.dollars_per_sol; + let expected_amount = bid.accepted_amount_dollars / args.dollars_per_sol.unwrap(); + assert_eq!(transaction_infos[0].amount, expected_amount); + + assert_eq!( + thin_client.get_balance(&alice_pubkey).unwrap(), + sol_to_lamports(expected_amount), + ); +} + +pub fn test_process_distribute_allocations_with_client(client: C, sender_keypair: Keypair) { + let thin_client = ThinClient(client); + let fee_payer = Keypair::new(); + thin_client + .transfer(sol_to_lamports(1.0), &sender_keypair, &fee_payer.pubkey()) + .unwrap(); + + let alice_pubkey = Pubkey::new_rand(); + let allocation = Allocation { + recipient: alice_pubkey.to_string(), + amount: 1000.0, + }; + let allocations_file = NamedTempFile::new().unwrap(); + let input_csv = allocations_file.path().to_str().unwrap().to_string(); + let mut wtr = csv::WriterBuilder::new().from_writer(allocations_file); + wtr.serialize(&allocation).unwrap(); + wtr.flush().unwrap(); + + let dir = tempdir().unwrap(); + let transactions_db = dir + .path() + .join("transactions.csv") + .to_str() + .unwrap() + .to_string(); + + let args: DistributeTokensArgs> = DistributeTokensArgs { + sender_keypair: Some(Box::new(sender_keypair)), + fee_payer: Some(Box::new(fee_payer)), + dry_run: false, + input_csv, + from_bids: false, + transactions_db: transactions_db.clone(), + dollars_per_sol: None, + }; + process_distribute_tokens(&thin_client, &args).unwrap(); + let transaction_infos = read_transaction_infos(&open_db(&transactions_db).unwrap()); + assert_eq!(transaction_infos.len(), 1); + assert_eq!(transaction_infos[0].recipient, alice_pubkey.to_string()); + let expected_amount = allocation.amount; + assert_eq!(transaction_infos[0].amount, expected_amount); + + assert_eq!( + thin_client.get_balance(&alice_pubkey).unwrap(), + sol_to_lamports(expected_amount), + ); + + // Now, run it again, and check there's no double-spend. + process_distribute_tokens(&thin_client, &args).unwrap(); + let transaction_infos = read_transaction_infos(&open_db(&transactions_db).unwrap()); + assert_eq!(transaction_infos.len(), 1); + assert_eq!(transaction_infos[0].recipient, alice_pubkey.to_string()); + let expected_amount = allocation.amount; assert_eq!(transaction_infos[0].amount, expected_amount); assert_eq!( @@ -589,11 +685,19 @@ mod tests { use solana_sdk::genesis_config::create_genesis_config; #[test] - fn test_process_distribute() { + fn test_process_distribute_bids() { + let (genesis_config, sender_keypair) = create_genesis_config(sol_to_lamports(9_000_000.0)); + let bank = Bank::new(&genesis_config); + let bank_client = BankClient::new(bank); + test_process_distribute_bids_with_client(bank_client, sender_keypair); + } + + #[test] + fn test_process_distribute_allocations() { let (genesis_config, sender_keypair) = create_genesis_config(sol_to_lamports(9_000_000.0)); let bank = Bank::new(&genesis_config); let bank_client = BankClient::new(bank); - test_process_distribute_with_client(bank_client, sender_keypair); + test_process_distribute_allocations_with_client(bank_client, sender_keypair); } #[test] diff --git a/tests/tokens.rs b/tests/tokens.rs index 7a669a3944f18d..c2c3fd255fbb65 100644 --- a/tests/tokens.rs +++ b/tests/tokens.rs @@ -1,7 +1,7 @@ use solana_client::rpc_client::RpcClient; use solana_core::validator::{TestValidator, TestValidatorOptions}; use solana_sdk::native_token::sol_to_lamports; -use solana_tokens::tokens::test_process_distribute_with_client; +use solana_tokens::tokens::test_process_distribute_bids_with_client; use std::fs::remove_dir_all; #[test] @@ -12,7 +12,7 @@ fn test_process_distribute_with_rpc_client() { ..TestValidatorOptions::default() }); let rpc_client = RpcClient::new_socket(validator.leader_data.rpc); - test_process_distribute_with_client(rpc_client, validator.alice); + test_process_distribute_bids_with_client(rpc_client, validator.alice); validator.server.close().unwrap(); remove_dir_all(validator.ledger_path).unwrap();