spark_sdk/wallet/internal_handlers/implementations/
leaves.rs

1use spark_protos::spark::{query_nodes_request::Source, QueryNodesRequest};
2use tonic::{async_trait, Request};
3
4use crate::{
5    error::SparkSdkError,
6    signer::traits::SparkSigner,
7    wallet::{
8        internal_handlers::traits::leaves::{LeafSelectionResponse, LeavesInternalHandlers},
9        leaf_manager::{LeafNode, LeafNodeStatus, PublicLeafNode},
10    },
11    SparkSdk,
12};
13
14#[async_trait]
15impl<S: SparkSigner + Send + Sync + Clone + 'static> LeavesInternalHandlers<S> for SparkSdk<S> {
16    fn get_btc_value(
17        &self,
18        filter_cb: Option<Box<dyn Fn(&PublicLeafNode) -> bool>>,
19    ) -> Result<u64, SparkSdkError> {
20        let value = self.leaf_manager.get_btc_value(filter_cb)?;
21        Ok(value)
22    }
23
24    fn query_single_node(
25        &self,
26        cb: Option<Box<dyn Fn(&LeafNode) -> bool>>,
27        new_status: Option<LeafNodeStatus>,
28    ) -> Result<PublicLeafNode, SparkSdkError> {
29        self.leaf_manager.query_single_node(cb, new_status)
30    }
31
32    fn get_node_with_id(&self, leaf_id: &String) -> Result<PublicLeafNode, SparkSdkError> {
33        let node = self.leaf_manager.get_node(leaf_id)?;
34        Ok(node)
35    }
36
37    async fn prepare_leaves_for_amount(
38        &self,
39        amount: u64,
40        token_pubkey: Option<&Vec<u8>>,
41        new_status: LeafNodeStatus,
42    ) -> Result<LeafSelectionResponse, SparkSdkError> {
43        // get the leaves that are
44        println!("Getting ready to select leaves");
45        let leaf_selection_response =
46            self.leaf_manager
47                .select_leaves(amount, token_pubkey, new_status)?;
48
49        println!("Leaf selection response: {:?}", leaf_selection_response);
50
51        let total_value = leaf_selection_response.total_value.clone();
52        let leaves = leaf_selection_response.leaves.clone();
53
54        // if the total leaves are over the amount, swap with the SSP
55        if total_value > amount {
56            let leaf_ids = leaves.iter().map(|leaf| leaf.id.clone()).collect();
57            let leaves = self.request_leaves_swap(total_value, leaf_ids).await?;
58
59            todo!("Handle SSP swap for splits in leaf selection");
60        }
61
62        Ok(leaf_selection_response)
63    }
64
65    async fn verify_and_get_leaves(
66        &self,
67        leaf_ids: Vec<String>,
68        new_status: LeafNodeStatus,
69    ) -> Result<LeafSelectionResponse, SparkSdkError> {
70        let (leaves, unlocking_id) = self
71            .leaf_manager
72            .lock_leaf_ids_for_operation(&leaf_ids, new_status)?;
73
74        let total_value = leaves.iter().map(|leaf| leaf.value).sum();
75        let leaf_selection_response = LeafSelectionResponse {
76            leaves,
77            total_value,
78            unlocking_id: Some(unlocking_id),
79        };
80
81        Ok(leaf_selection_response)
82    }
83
84    async fn select_leaf_to_split(
85        &self,
86        target_value: u64,
87    ) -> Result<PublicLeafNode, SparkSdkError> {
88        let leaf = self.leaf_manager.select_leaf_to_split(target_value)?;
89        Ok(leaf)
90    }
91
92    async fn fetch_owned_leaves_from_spark(&self) -> Result<(), SparkSdkError> {
93        let mut spark_client = self.config.spark_config.get_spark_connection(None).await?;
94
95        let identity_pubkey = self.get_identity_public_key().to_vec();
96        let mut request = Request::new(QueryNodesRequest {
97            include_parents: true,
98            source: Some(Source::OwnerIdentityPubkey(identity_pubkey)),
99        });
100
101        self.add_authorization_header_to_request(&mut request, None);
102
103        let response = spark_client.query_nodes(request).await?;
104        let queried_nodes = response.into_inner().nodes;
105
106        // select the available leaves only
107        let available_leaves_with_ids = queried_nodes
108            .into_iter()
109            .filter(|node| node.1.status.to_uppercase() == "AVAILABLE")
110            .collect::<Vec<_>>();
111
112        let available_leaves = available_leaves_with_ids
113            .into_iter()
114            .map(|(_, node)| node.into())
115            .collect();
116
117        // refresh the leaf manager with the new leaves
118        self.leaf_manager.refresh_leaves(available_leaves)?;
119
120        Ok(())
121    }
122}