diff --git a/accountsdb-plugin-postgres/src/lib.rs b/accountsdb-plugin-postgres/src/lib.rs index 55eb28af00..a2d78c1eef 100644 --- a/accountsdb-plugin-postgres/src/lib.rs +++ b/accountsdb-plugin-postgres/src/lib.rs @@ -1,3 +1,4 @@ pub mod accounts_selector; pub mod accountsdb_plugin_postgres; pub mod postgres_client; +pub mod transaction_selector; diff --git a/accountsdb-plugin-postgres/src/transaction_selector.rs b/accountsdb-plugin-postgres/src/transaction_selector.rs new file mode 100644 index 0000000000..7c256f3c07 --- /dev/null +++ b/accountsdb-plugin-postgres/src/transaction_selector.rs @@ -0,0 +1,194 @@ +/// The transaction selector is responsible for filtering transactions +/// in the plugin framework. +use {log::*, solana_sdk::pubkey::Pubkey, std::collections::HashSet}; + +pub(crate) struct TransactionSelector { + pub mentioned_addresses: HashSet>, + pub select_all_transactions: bool, + pub select_all_vote_transactions: bool, +} + +#[allow(dead_code)] +impl TransactionSelector { + pub fn default() -> Self { + Self { + mentioned_addresses: HashSet::default(), + select_all_transactions: false, + select_all_vote_transactions: false, + } + } + + /// Create a selector based on the mentioned addresses + /// To select all transactions use ["*"] or ["all"] + /// To select all vote transactions, use ["all_votes"] + /// To select transactions mentioning specific addresses use ["", "", ...] + pub fn new(mentioned_addresses: &[String]) -> Self { + info!( + "Creating TransactionSelector from addresses: {:?}", + mentioned_addresses + ); + + let select_all_transactions = mentioned_addresses + .iter() + .any(|key| key == "*" || key == "all"); + if select_all_transactions { + return Self { + mentioned_addresses: HashSet::default(), + select_all_transactions, + select_all_vote_transactions: true, + }; + } + let select_all_vote_transactions = mentioned_addresses.iter().any(|key| key == "all_votes"); + if select_all_vote_transactions { + return Self { + mentioned_addresses: HashSet::default(), + select_all_transactions, + select_all_vote_transactions: true, + }; + } + + let mentioned_addresses = mentioned_addresses + .iter() + .map(|key| bs58::decode(key).into_vec().unwrap()) + .collect(); + + Self { + mentioned_addresses, + select_all_transactions: false, + select_all_vote_transactions: false, + } + } + + /// Check if a transaction is of interest. + pub fn is_transaction_selected( + &self, + is_vote: bool, + mentioned_addresses: Box + '_>, + ) -> bool { + if !self.is_enabled() { + return false; + } + + if self.select_all_transactions || (self.select_all_vote_transactions && is_vote) { + return true; + } + for address in mentioned_addresses { + if self.mentioned_addresses.contains(address.as_ref()) { + return true; + } + } + false + } + + /// Check if any transaction is of interest at all + pub fn is_enabled(&self) -> bool { + self.select_all_transactions + || self.select_all_vote_transactions + || !self.mentioned_addresses.is_empty() + } +} + +#[cfg(test)] +pub(crate) mod tests { + use super::*; + + #[test] + fn test_select_transaction() { + let pubkey1 = Pubkey::new_unique(); + let pubkey2 = Pubkey::new_unique(); + + let selector = TransactionSelector::new(&[pubkey1.to_string()]); + + assert!(selector.is_enabled()); + + let addresses = [pubkey1]; + + assert!(selector.is_transaction_selected(false, Box::new(addresses.iter()))); + + let addresses = [pubkey2]; + assert!(!selector.is_transaction_selected(false, Box::new(addresses.iter()))); + + let addresses = [pubkey1, pubkey2]; + assert!(selector.is_transaction_selected(false, Box::new(addresses.iter()))); + } + + #[test] + fn test_select_all_transaction_using_wildcard() { + let pubkey1 = Pubkey::new_unique(); + let pubkey2 = Pubkey::new_unique(); + + let selector = TransactionSelector::new(&["*".to_string()]); + + assert!(selector.is_enabled()); + + let addresses = [pubkey1]; + + assert!(selector.is_transaction_selected(false, Box::new(addresses.iter()))); + + let addresses = [pubkey2]; + assert!(selector.is_transaction_selected(false, Box::new(addresses.iter()))); + + let addresses = [pubkey1, pubkey2]; + assert!(selector.is_transaction_selected(false, Box::new(addresses.iter()))); + } + + #[test] + fn test_select_all_transaction_all() { + let pubkey1 = Pubkey::new_unique(); + let pubkey2 = Pubkey::new_unique(); + + let selector = TransactionSelector::new(&["all".to_string()]); + + assert!(selector.is_enabled()); + + let addresses = [pubkey1]; + + assert!(selector.is_transaction_selected(false, Box::new(addresses.iter()))); + + let addresses = [pubkey2]; + assert!(selector.is_transaction_selected(false, Box::new(addresses.iter()))); + + let addresses = [pubkey1, pubkey2]; + assert!(selector.is_transaction_selected(false, Box::new(addresses.iter()))); + } + + #[test] + fn test_select_all_vote_transaction() { + let pubkey1 = Pubkey::new_unique(); + let pubkey2 = Pubkey::new_unique(); + + let selector = TransactionSelector::new(&["all_votes".to_string()]); + + assert!(selector.is_enabled()); + + let addresses = [pubkey1]; + + assert!(!selector.is_transaction_selected(false, Box::new(addresses.iter()))); + + let addresses = [pubkey2]; + assert!(selector.is_transaction_selected(true, Box::new(addresses.iter()))); + + let addresses = [pubkey1, pubkey2]; + assert!(selector.is_transaction_selected(true, Box::new(addresses.iter()))); + } + + #[test] + fn test_select_no_transaction() { + let pubkey1 = Pubkey::new_unique(); + let pubkey2 = Pubkey::new_unique(); + + let selector = TransactionSelector::new(&[]); + + assert!(!selector.is_enabled()); + + let addresses = [pubkey1]; + + assert!(!selector.is_transaction_selected(false, Box::new(addresses.iter()))); + + let addresses = [pubkey2]; + assert!(!selector.is_transaction_selected(true, Box::new(addresses.iter()))); + + let addresses = [pubkey1, pubkey2]; + assert!(!selector.is_transaction_selected(true, Box::new(addresses.iter()))); + } +}