Bläddra i källkod

solana-tokens: Add capability to perform the same transfer to a batch of recipients (#12259)

* Add transfer-amount argument, use simplified input-csv

* Add transfer-amount to readme
Tyera Eulberg 5 år sedan
förälder
incheckning
a48cc073cf
5 ändrade filer med 130 tillägg och 9 borttagningar
  1. 26 0
      tokens/README.md
  2. 12 1
      tokens/src/arg_parser.rs
  3. 1 0
      tokens/src/args.rs
  4. 90 7
      tokens/src/commands.rs
  5. 1 1
      tokens/tests/commands.rs

+ 26 - 0
tokens/README.md

@@ -82,6 +82,32 @@ Recipient                                     Expected Balance (◎)
 7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr  42
 ```
 
+## Distribute tokens: transfer-amount
+
+This tool also makes it straightforward to transfer the same amount of tokens to a simple list of recipients. Just add the `--transfer-amount` arg to specify the amount:
+
+Example recipients.csv:
+
+```text
+recipient
+6Vo87BaDhp4v4GHwVDhw5huhxVF8CyxSXYtkUwVHbbPv
+7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr
+CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT
+```
+
+```bash
+solana-tokens distribute-tokens --transfer-amount 10 --from <KEYPAIR> --input-csv <RECIPIENTS_CSV> --fee-payer <KEYPAIR>
+```
+
+Example output:
+
+```text
+Recipient                                     Expected Balance (◎)
+6Vo87BaDhp4v4GHwVDhw5huhxVF8CyxSXYtkUwVHbbPv  10
+7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr  10
+CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT  10
+```
+
 ## Distribute stake accounts
 
 Distributing tokens via stake accounts works similarly to how tokens are distributed. The

+ 12 - 1
tokens/src/arg_parser.rs

@@ -3,7 +3,8 @@ use crate::args::{
 };
 use clap::{value_t, value_t_or_exit, App, Arg, ArgMatches, SubCommand};
 use solana_clap_utils::{
-    input_validators::{is_valid_pubkey, is_valid_signer},
+    input_parsers::value_of,
+    input_validators::{is_amount, is_valid_pubkey, is_valid_signer},
     keypair::{pubkey_from_path, signer_from_path},
 };
 use solana_cli_config::CONFIG_FILE;
@@ -60,6 +61,14 @@ where
                         .value_name("FILE")
                         .help("Input CSV file"),
                 )
+                .arg(
+                    Arg::with_name("transfer_amount")
+                        .long("transfer-amount")
+                        .takes_value(true)
+                        .value_name("AMOUNT")
+                        .validator(is_amount)
+                        .help("The amount to send to each recipient, in SOL"),
+                )
                 .arg(
                     Arg::with_name("dry_run")
                         .long("dry-run")
@@ -255,6 +264,7 @@ fn parse_distribute_tokens_args(
         sender_keypair,
         fee_payer,
         stake_args: None,
+        transfer_amount: value_of(matches, "transfer_amount"),
     })
 }
 
@@ -330,6 +340,7 @@ fn parse_distribute_stake_args(
         sender_keypair,
         fee_payer,
         stake_args: Some(stake_args),
+        transfer_amount: None,
     })
 }
 

+ 1 - 0
tokens/src/args.rs

@@ -8,6 +8,7 @@ pub struct DistributeTokensArgs {
     pub sender_keypair: Box<dyn Signer>,
     pub fee_payer: Box<dyn Signer>,
     pub stake_args: Option<StakeArgs>,
+    pub transfer_amount: Option<f64>,
 }
 
 pub struct StakeArgs {

+ 90 - 7
tokens/src/commands.rs

@@ -274,9 +274,25 @@ async fn distribute_allocations(
     Ok(())
 }
 
-fn read_allocations(input_csv: &str) -> io::Result<Vec<Allocation>> {
+fn read_allocations(input_csv: &str, transfer_amount: Option<f64>) -> io::Result<Vec<Allocation>> {
     let mut rdr = ReaderBuilder::new().trim(Trim::All).from_path(input_csv)?;
-    Ok(rdr.deserialize().map(|entry| entry.unwrap()).collect())
+    let allocations = if let Some(amount) = transfer_amount {
+        let recipients: Vec<String> = rdr
+            .deserialize()
+            .map(|recipient| recipient.unwrap())
+            .collect();
+        recipients
+            .into_iter()
+            .map(|recipient| Allocation {
+                recipient,
+                amount,
+                lockup_date: "".to_string(),
+            })
+            .collect()
+    } else {
+        rdr.deserialize().map(|entry| entry.unwrap()).collect()
+    };
+    Ok(allocations)
 }
 
 fn new_spinner_progress_bar() -> ProgressBar {
@@ -291,7 +307,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)?;
+    let mut allocations: Vec<Allocation> = read_allocations(&args.input_csv, args.transfer_amount)?;
 
     let starting_total_tokens: f64 = allocations.iter().map(|x| x.amount).sum();
     println!(
@@ -467,7 +483,7 @@ pub async fn process_balances(
     client: &mut BanksClient,
     args: &BalancesArgs,
 ) -> Result<(), csv::Error> {
-    let allocations: Vec<Allocation> = read_allocations(&args.input_csv)?;
+    let allocations: Vec<Allocation> = read_allocations(&args.input_csv, None)?;
     let allocations = merge_allocations(&allocations);
 
     println!(
@@ -507,6 +523,7 @@ use tempfile::{tempdir, NamedTempFile};
 pub async fn test_process_distribute_tokens_with_client(
     client: &mut BanksClient,
     sender_keypair: Keypair,
+    transfer_amount: Option<f64>,
 ) {
     let fee_payer = Keypair::new();
     let transaction = transfer(
@@ -532,7 +549,11 @@ pub async fn test_process_distribute_tokens_with_client(
     let alice_pubkey = Pubkey::new_rand();
     let allocation = Allocation {
         recipient: alice_pubkey.to_string(),
-        amount: 1000.0,
+        amount: if let Some(amount) = transfer_amount {
+            amount
+        } else {
+            1000.0
+        },
         lockup_date: "".to_string(),
     };
     let allocations_file = NamedTempFile::new().unwrap();
@@ -560,6 +581,7 @@ pub async fn test_process_distribute_tokens_with_client(
         transaction_db: transaction_db.clone(),
         output_path: Some(output_path.clone()),
         stake_args: None,
+        transfer_amount,
     };
     let confirmations = process_allocations(client, &args).await.unwrap();
     assert_eq!(confirmations, None);
@@ -683,6 +705,7 @@ pub async fn test_process_distribute_stake_with_client(
         output_path: Some(output_path.clone()),
         stake_args: Some(stake_args),
         sender_keypair: Box::new(sender_keypair),
+        transfer_amount: None,
     };
     let confirmations = process_allocations(client, &args).await.unwrap();
     assert_eq!(confirmations, None);
@@ -751,7 +774,24 @@ mod tests {
         Runtime::new().unwrap().block_on(async {
             let transport = start_local_server(&bank_forks).await;
             let mut banks_client = start_client(transport).await.unwrap();
-            test_process_distribute_tokens_with_client(&mut banks_client, sender_keypair).await;
+            test_process_distribute_tokens_with_client(&mut banks_client, sender_keypair, None)
+                .await;
+        });
+    }
+
+    #[test]
+    fn test_process_transfer_amount_allocations() {
+        let (genesis_config, sender_keypair) = create_genesis_config(sol_to_lamports(9_000_000.0));
+        let bank_forks = Arc::new(RwLock::new(BankForks::new(Bank::new(&genesis_config))));
+        Runtime::new().unwrap().block_on(async {
+            let transport = start_local_server(&bank_forks).await;
+            let mut banks_client = start_client(transport).await.unwrap();
+            test_process_distribute_tokens_with_client(
+                &mut banks_client,
+                sender_keypair,
+                Some(1.5),
+            )
+            .await;
         });
     }
 
@@ -780,7 +820,49 @@ mod tests {
         wtr.serialize(&allocation).unwrap();
         wtr.flush().unwrap();
 
-        assert_eq!(read_allocations(&input_csv).unwrap(), vec![allocation]);
+        assert_eq!(
+            read_allocations(&input_csv, None).unwrap(),
+            vec![allocation]
+        );
+    }
+
+    #[test]
+    fn test_read_allocations_transfer_amount() {
+        let pubkey0 = Pubkey::new_rand();
+        let pubkey1 = Pubkey::new_rand();
+        let pubkey2 = Pubkey::new_rand();
+        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("recipient".to_string()).unwrap();
+        wtr.serialize(&pubkey0.to_string()).unwrap();
+        wtr.serialize(&pubkey1.to_string()).unwrap();
+        wtr.serialize(&pubkey2.to_string()).unwrap();
+        wtr.flush().unwrap();
+
+        let amount = 1.5;
+
+        let expected_allocations = vec![
+            Allocation {
+                recipient: pubkey0.to_string(),
+                amount,
+                lockup_date: "".to_string(),
+            },
+            Allocation {
+                recipient: pubkey1.to_string(),
+                amount,
+                lockup_date: "".to_string(),
+            },
+            Allocation {
+                recipient: pubkey2.to_string(),
+                amount,
+                lockup_date: "".to_string(),
+            },
+        ];
+        assert_eq!(
+            read_allocations(&input_csv, Some(amount)).unwrap(),
+            expected_allocations
+        );
     }
 
     #[test]
@@ -889,6 +971,7 @@ mod tests {
             output_path: None,
             stake_args: Some(stake_args),
             sender_keypair: Box::new(Keypair::new()),
+            transfer_amount: None,
         };
         let lockup_date = lockup_date_str.parse().unwrap();
         let instructions = distribution_instructions(

+ 1 - 1
tokens/tests/commands.rs

@@ -20,7 +20,7 @@ fn test_process_distribute_with_rpc_client() {
 
     Runtime::new().unwrap().block_on(async {
         let mut banks_client = start_tcp_client(leader_data.rpc_banks).await.unwrap();
-        test_process_distribute_tokens_with_client(&mut banks_client, alice).await
+        test_process_distribute_tokens_with_client(&mut banks_client, alice, None).await
     });
 
     // Explicit cleanup, otherwise "pure virtual method called" crash in Docker