gaspersic/solana
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Improve solana-tokens UX (#12253)

Browse Source 
* Fix computed banks port

* Readme incorrect

* Return error if csv cannot be read

* Move column headers over columns

* Add dry-run check for sender/fee-payer balances

* Use clap requires method for paired args

* Write transaction-log anytime outfile is specified

* Replace campaign-name with required db-path

* Remove bids

* Exclude new_stake_account_address from logs for non-stake distributions

* Fix readme
This commit is contained in:
Tyera Eulberg
2020-09-15 19:38:22 -06:00
committed by GitHub
parent 3930cb865a
commit 90a591da0e
6 changed files with 199 additions and 200 deletions
Download Patch File Download Diff File Expand all files Collapse all files

105
tokens/src/arg_parser.rs
View File

@ -41,16 +41,16 @@ where
SubCommand::with_name("distribute-tokens")
.about("Distribute tokens")
.arg(
Arg::with_name("campaign_name")
.long("campaign-name")
Arg::with_name("db_path")
.long("db-path")
.required(true)
.takes_value(true)
.value_name("NAME")
.help("Campaign name for storing transaction data"),
)
.arg(
Arg::with_name("from_bids")
.long("from-bids")
.help("Input CSV contains bids in dollars, not allocations in SOL"),
.value_name("FILE")
.help(
"Location for storing distribution database. \
The database is used for tracking transactions as they are finalized \
and preventing double spends.",
),
)
.arg(
Arg::with_name("input_csv")
@ -60,18 +60,19 @@ where
.value_name("FILE")
.help("Input CSV file"),
)
.arg(
Arg::with_name("dollars_per_sol")
.long("dollars-per-sol")
.takes_value(true)
.value_name("NUMBER")
.help("Dollars per SOL, if input CSV contains bids"),
)
.arg(
Arg::with_name("dry_run")
.long("dry-run")
.help("Do not execute any transfers"),
)
.arg(
Arg::with_name("output_path")
.long("output-path")
.short("o")
.value_name("FILE")
.takes_value(true)
.help("Write the transaction log to this file"),
)
.arg(
Arg::with_name("sender_keypair")
.long("from")
@ -95,11 +96,16 @@ where
SubCommand::with_name("distribute-stake")
.about("Distribute stake accounts")
.arg(
Arg::with_name("campaign_name")
.long("campaign-name")
Arg::with_name("db_path")
.long("db-path")
.required(true)
.takes_value(true)
.value_name("NAME")
.help("Campaign name for storing transaction data"),
.value_name("FILE")
.help(
"Location for storing distribution database. \
The database is used for tracking transactions as they are finalized \
and preventing double spends.",
),
)
.arg(
Arg::with_name("input_csv")
@ -114,6 +120,14 @@ where
.long("dry-run")
.help("Do not execute any transfers"),
)
.arg(
Arg::with_name("output_path")
.long("output-path")
.short("o")
.value_name("FILE")
.takes_value(true)
.help("Write the transaction log to this file"),
)
.arg(
Arg::with_name("sender_keypair")
.long("from")
@ -186,29 +200,18 @@ where
.takes_value(true)
.value_name("FILE")
.help("Bids CSV file"),
)
.arg(
Arg::with_name("from_bids")
.long("from-bids")
.help("Input CSV contains bids in dollars, not allocations in SOL"),
)
.arg(
Arg::with_name("dollars_per_sol")
.long("dollars-per-sol")
.takes_value(true)
.value_name("NUMBER")
.help("Dollars per SOL"),
),
)
.subcommand(
SubCommand::with_name("transaction-log")
.about("Print the database to a CSV file")
.arg(
Arg::with_name("campaign_name")
.long("campaign-name")
Arg::with_name("db_path")
.long("db-path")
.required(true)
.takes_value(true)
.value_name("NAME")
.help("Campaign name for storing transaction data"),
.value_name("FILE")
.help("Location of database to query"),
)
.arg(
Arg::with_name("output_path")
@ -222,22 +225,6 @@ where
.get_matches_from(args)
}
fn create_db_path(campaign_name: Option<String>) -> String {
let (prefix, hyphen) = if let Some(name) = campaign_name {
(name, "-")
} else {
("".to_string(), "")
};
let path = dirs::home_dir().unwrap();
let filename = format!("{}{}transactions.db", prefix, hyphen);
path.join(".config")
.join("solana-tokens")
.join(filename)
.to_str()
.unwrap()
.to_string()
}
fn parse_distribute_tokens_args(
matches: &ArgMatches<'_>,
) -> Result<DistributeTokensArgs, Box<dyn Error>> {
@ -262,9 +249,8 @@ fn parse_distribute_tokens_args(
Ok(DistributeTokensArgs {
input_csv: value_t_or_exit!(matches, "input_csv", String),
from_bids: matches.is_present("from_bids"),
transaction_db: create_db_path(value_t!(matches, "campaign_name", String).ok()),
dollars_per_sol: value_t!(matches, "dollars_per_sol", f64).ok(),
transaction_db: value_t_or_exit!(matches, "db_path", String),
output_path: matches.value_of("output_path").map(|path| path.to_string()),
dry_run: matches.is_present("dry_run"),
sender_keypair,
fee_payer,
@ -338,9 +324,8 @@ fn parse_distribute_stake_args(
};
Ok(DistributeTokensArgs {
input_csv: value_t_or_exit!(matches, "input_csv", String),
from_bids: false,
transaction_db: create_db_path(value_t!(matches, "campaign_name", String).ok()),
dollars_per_sol: None,
transaction_db: value_t_or_exit!(matches, "db_path", String),
output_path: matches.value_of("output_path").map(|path| path.to_string()),
dry_run: matches.is_present("dry_run"),
sender_keypair,
fee_payer,
@ -351,14 +336,12 @@ fn parse_distribute_stake_args(
fn parse_balances_args(matches: &ArgMatches<'_>) -> BalancesArgs {
BalancesArgs {
input_csv: value_t_or_exit!(matches, "input_csv", String),
from_bids: matches.is_present("from_bids"),
dollars_per_sol: value_t!(matches, "dollars_per_sol", f64).ok(),
}
}
fn parse_transaction_log_args(matches: &ArgMatches<'_>) -> TransactionLogArgs {
TransactionLogArgs {
transaction_db: create_db_path(value_t!(matches, "campaign_name", String).ok()),
transaction_db: value_t_or_exit!(matches, "db_path", String),
output_path: value_t_or_exit!(matches, "output_path", String),
}
}

5
tokens/src/args.rs
View File

@ -2,9 +2,8 @@ use solana_sdk::{pubkey::Pubkey, signature::Signer};
pub struct DistributeTokensArgs {
pub input_csv: String,
pub from_bids: bool,
pub transaction_db: String,
pub dollars_per_sol: Option<f64>,
pub output_path: Option<String>,
pub dry_run: bool,
pub sender_keypair: Box<dyn Signer>,
pub fee_payer: Box<dyn Signer>,
@ -21,8 +20,6 @@ pub struct StakeArgs {
pub struct BalancesArgs {
pub input_csv: String,
pub from_bids: bool,
pub dollars_per_sol: Option<f64>,
}
pub struct TransactionLogArgs {

201
tokens/src/commands.rs
View File

@ -54,6 +54,12 @@ pub enum Error {
TransportError(#[from] TransportError),
#[error("Missing lockup authority")]
MissingLockupAuthority,
#[error("insufficient funds for fee ({0} SOL)")]
InsufficientFundsForFees(f64),
#[error("insufficient funds for distribution ({0} SOL)")]
InsufficientFundsForDistribution(f64),
#[error("insufficient funds for distribution ({0} SOL) and fee ({1} SOL)")]
InsufficientFundsForDistributionAndFees(f64, f64),
}
fn merge_allocations(allocations: &[Allocation]) -> Vec<Allocation> {
@ -99,14 +105,6 @@ fn apply_previous_transactions(
allocations.retain(|x| x.amount > 0.5);
}
fn create_allocation(bid: &Bid, dollars_per_sol: f64) -> Allocation {
Allocation {
recipient: bid.primary_address.clone(),
amount: bid.accepted_amount_dollars / dollars_per_sol,
lockup_date: "".to_string(),
}
}
async fn transfer<S: Signer>(
client: &mut BanksClient,
lamports: u64,
@ -203,6 +201,7 @@ async fn distribute_allocations(
allocations: &[Allocation],
args: &DistributeTokensArgs,
) -> Result<(), Error> {
let mut num_signatures = 0;
for allocation in allocations {
let new_stake_account_keypair = Keypair::new();
let new_stake_account_address = new_stake_account_keypair.pubkey();
@ -221,6 +220,7 @@ async fn distribute_allocations(
}
}
let signers = unique_signers(signers);
num_signatures += signers.len();
let lockup_date = if allocation.lockup_date == "" {
None
@ -250,7 +250,7 @@ async fn distribute_allocations(
&allocation.recipient.parse().unwrap(),
allocation.amount,
&transaction,
Some(&new_stake_account_address),
args.stake_args.as_ref().map(|_| &new_stake_account_address),
false,
last_valid_slot,
lockup_date,
@ -261,26 +261,22 @@ async fn distribute_allocations(
}
};
}
if args.dry_run {
let undistributed_tokens: f64 = allocations.iter().map(|x| x.amount).sum();
check_payer_balances(
num_signatures,
sol_to_lamports(undistributed_tokens),
client,
args,
)
.await?;
}
Ok(())
}
fn read_allocations(
input_csv: &str,
from_bids: bool,
dollars_per_sol: Option<f64>,
) -> Vec<Allocation> {
let rdr = ReaderBuilder::new().trim(Trim::All).from_path(input_csv);
if from_bids {
let bids: Vec<Bid> = rdr.unwrap().deserialize().map(|bid| bid.unwrap()).collect();
bids.into_iter()
.map(|bid| create_allocation(&bid, dollars_per_sol.unwrap()))
.collect()
} else {
rdr.unwrap()
.deserialize()
.map(|entry| entry.unwrap())
.collect()
}
fn read_allocations(input_csv: &str) -> io::Result<Vec<Allocation>> {
let mut rdr = ReaderBuilder::new().trim(Trim::All).from_path(input_csv)?;
Ok(rdr.deserialize().map(|entry| entry.unwrap()).collect())
}
fn new_spinner_progress_bar() -> ProgressBar {
@ -295,8 +291,7 @@ pub async fn process_allocations(
client: &mut BanksClient,
args: &DistributeTokensArgs,
) -> Result<Option<usize>, Error> {
let mut allocations: Vec<Allocation> =
read_allocations(&args.input_csv, args.from_bids, args.dollars_per_sol);
let mut allocations: Vec<Allocation> = read_allocations(&args.input_csv)?;
let starting_total_tokens: f64 = allocations.iter().map(|x| x.amount).sum();
println!(
@ -304,13 +299,6 @@ pub async fn process_allocations(
style("Total in input_csv:").bold(),
starting_total_tokens,
);
if let Some(dollars_per_sol) = args.dollars_per_sol {
println!(
"{} ${}",
style("Total in input_csv:").bold(),
starting_total_tokens * dollars_per_sol,
);
}
let mut db = db::open_db(&args.transaction_db, args.dry_run)?;
@ -325,6 +313,20 @@ pub async fn process_allocations(
return Ok(confirmations);
}
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("Distributed:").bold(), distributed_tokens,);
println!(
"{} ◎{}",
style("Undistributed:").bold(),
undistributed_tokens,
);
println!(
"{} ◎{}",
style("Total:").bold(),
distributed_tokens + undistributed_tokens,
);
println!(
"{}",
style(format!(
@ -334,44 +336,16 @@ pub async fn process_allocations(
.bold()
);
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("Distributed:").bold(), distributed_tokens,);
if let Some(dollars_per_sol) = args.dollars_per_sol {
println!(
"{} ${}",
style("Distributed:").bold(),
distributed_tokens * dollars_per_sol,
);
}
println!(
"{} ◎{}",
style("Undistributed:").bold(),
undistributed_tokens,
);
if let Some(dollars_per_sol) = args.dollars_per_sol {
println!(
"{} ${}",
style("Undistributed:").bold(),
undistributed_tokens * dollars_per_sol,
);
}
println!(
"{} ◎{}",
style("Total:").bold(),
distributed_tokens + undistributed_tokens,
);
if let Some(dollars_per_sol) = args.dollars_per_sol {
println!(
"{} ${}",
style("Total:").bold(),
(distributed_tokens + undistributed_tokens) * dollars_per_sol,
);
}
distribute_allocations(client, &mut db, &allocations, args).await?;
let opt_confirmations = finalize_transactions(client, &mut db, args.dry_run).await?;
if !args.dry_run {
if let Some(output_path) = &args.output_path {
db::write_transaction_log(&db, &output_path)?;
}
}
Ok(opt_confirmations)
}
@ -455,12 +429,45 @@ async fn update_finalized_transactions(
Ok(confirmations)
}
async fn check_payer_balances(
num_signatures: usize,
allocation_lamports: u64,
client: &mut BanksClient,
args: &DistributeTokensArgs,
) -> Result<(), Error> {
let (fee_calculator, _blockhash, _last_valid_slot) = client.get_fees().await?;
let fees = fee_calculator
.lamports_per_signature
.checked_mul(num_signatures as u64)
.unwrap();
if args.fee_payer.pubkey() == args.sender_keypair.pubkey() {
let balance = client.get_balance(args.fee_payer.pubkey()).await?;
if balance < fees + allocation_lamports {
return Err(Error::InsufficientFundsForDistributionAndFees(
lamports_to_sol(allocation_lamports),
lamports_to_sol(fees),
));
}
} else {
let fee_payer_balance = client.get_balance(args.fee_payer.pubkey()).await?;
if fee_payer_balance < fees {
return Err(Error::InsufficientFundsForFees(lamports_to_sol(fees)));
}
let sender_balance = client.get_balance(args.sender_keypair.pubkey()).await?;
if sender_balance < allocation_lamports {
return Err(Error::InsufficientFundsForDistribution(lamports_to_sol(
allocation_lamports,
)));
}
}
Ok(())
}
pub async fn process_balances(
client: &mut BanksClient,
args: &BalancesArgs,
) -> Result<(), csv::Error> {
let allocations: Vec<Allocation> =
read_allocations(&args.input_csv, args.from_bids, args.dollars_per_sol);
let allocations: Vec<Allocation> = read_allocations(&args.input_csv)?;
let allocations = merge_allocations(&allocations);
println!(
@ -494,6 +501,7 @@ pub fn process_transaction_log(args: &TransactionLogArgs) -> Result<(), Error> {
Ok(())
}
use crate::db::check_output_file;
use solana_sdk::{pubkey::Pubkey, signature::Keypair};
use tempfile::{tempdir, NamedTempFile};
pub async fn test_process_distribute_tokens_with_client(
@ -541,14 +549,16 @@ pub async fn test_process_distribute_tokens_with_client(
.unwrap()
.to_string();
let output_file = NamedTempFile::new().unwrap();
let output_path = output_file.path().to_str().unwrap().to_string();
let args = DistributeTokensArgs {
sender_keypair: Box::new(sender_keypair),
fee_payer: Box::new(fee_payer),
dry_run: false,
input_csv,
from_bids: false,
transaction_db: transaction_db.clone(),
dollars_per_sol: None,
output_path: Some(output_path.clone()),
stake_args: None,
};
let confirmations = process_allocations(client, &args).await.unwrap();
@ -569,6 +579,8 @@ pub async fn test_process_distribute_tokens_with_client(
expected_amount,
);
check_output_file(&output_path, &db::open_db(&transaction_db, true).unwrap());
// Now, run it again, and check there's no double-spend.
process_allocations(client, &args).await.unwrap();
let transaction_infos =
@ -585,6 +597,8 @@ pub async fn test_process_distribute_tokens_with_client(
client.get_balance(alice_pubkey).await.unwrap(),
expected_amount,
);
check_output_file(&output_path, &db::open_db(&transaction_db, true).unwrap());
}
pub async fn test_process_distribute_stake_with_client(
@ -651,6 +665,9 @@ pub async fn test_process_distribute_stake_with_client(
.unwrap()
.to_string();
let output_file = NamedTempFile::new().unwrap();
let output_path = output_file.path().to_str().unwrap().to_string();
let stake_args = StakeArgs {
stake_account_address,
stake_authority: Box::new(stake_authority),
@ -663,10 +680,9 @@ pub async fn test_process_distribute_stake_with_client(
dry_run: false,
input_csv,
transaction_db: transaction_db.clone(),
output_path: Some(output_path.clone()),
stake_args: Some(stake_args),
from_bids: false,
sender_keypair: Box::new(sender_keypair),
dollars_per_sol: None,
};
let confirmations = process_allocations(client, &args).await.unwrap();
assert_eq!(confirmations, None);
@ -691,6 +707,8 @@ pub async fn test_process_distribute_stake_with_client(
expected_amount - sol_to_lamports(1.0),
);
check_output_file(&output_path, &db::open_db(&transaction_db, true).unwrap());
// Now, run it again, and check there's no double-spend.
process_allocations(client, &args).await.unwrap();
let transaction_infos =
@ -711,6 +729,8 @@ pub async fn test_process_distribute_stake_with_client(
client.get_balance(new_stake_account_address).await.unwrap(),
expected_amount - sol_to_lamports(1.0),
);
check_output_file(&output_path, &db::open_db(&transaction_db, true).unwrap());
}
#[cfg(test)]
@ -760,31 +780,7 @@ mod tests {
wtr.serialize(&allocation).unwrap();
wtr.flush().unwrap();
assert_eq!(read_allocations(&input_csv, false, None), vec![allocation]);
}
#[test]
fn test_read_allocations_from_bids() {
let alice_pubkey = Pubkey::new_rand();
let bid = Bid {
primary_address: alice_pubkey.to_string(),
accepted_amount_dollars: 42.0,
};
let file = NamedTempFile::new().unwrap();
let input_csv = file.path().to_str().unwrap().to_string();
let mut wtr = csv::WriterBuilder::new().from_writer(file);
wtr.serialize(&bid).unwrap();
wtr.flush().unwrap();
let allocation = Allocation {
recipient: bid.primary_address,
amount: 84.0,
lockup_date: "".to_string(),
};
assert_eq!(
read_allocations(&input_csv, true, Some(0.5)),
vec![allocation]
);
assert_eq!(read_allocations(&input_csv).unwrap(), vec![allocation]);
}
#[test]
@ -890,10 +886,9 @@ mod tests {
dry_run: false,
input_csv: "".to_string(),
transaction_db: "".to_string(),
output_path: None,
stake_args: Some(stake_args),
from_bids: false,
sender_keypair: Box::new(Keypair::new()),
dollars_per_sol: None,
};
let lockup_date = lockup_date_str.parse().unwrap();
let instructions = distribution_instructions(

28
tokens/src/db.rs
View File

@ -20,6 +20,7 @@ pub struct TransactionInfo {
struct SignedTransactionInfo {
recipient: String,
amount: f64,
#[serde(skip_serializing_if = "String::is_empty", default)]
new_stake_account_address: String,
finalized_date: Option<DateTime<Utc>>,
signature: String,
@ -178,6 +179,33 @@ pub fn update_finalized_transaction(
Ok(None)
}
use csv::{ReaderBuilder, Trim};
pub(crate) fn check_output_file(path: &str, db: &PickleDb) {
let mut rdr = ReaderBuilder::new()
.trim(Trim::All)
.from_path(path)
.unwrap();
let logged_infos: Vec<SignedTransactionInfo> =
rdr.deserialize().map(|entry| entry.unwrap()).collect();
let mut transaction_infos = read_transaction_infos(db);
transaction_infos.sort_by(compare_transaction_infos);
let transaction_infos: Vec<SignedTransactionInfo> = transaction_infos
.iter()
.map(|info| SignedTransactionInfo {
recipient: info.recipient.to_string(),
amount: info.amount,
new_stake_account_address: info
.new_stake_account_address
.map(|x| x.to_string())
.unwrap_or_else(|| "".to_string()),
finalized_date: info.finalized_date,
signature: info.transaction.signatures[0].to_string(),
})
.collect();
assert_eq!(logged_infos, transaction_infos);
}
#[cfg(test)]
mod tests {
use super::*;