0.7.24-pre.0

kb_lib/src/trade_aggregation.rs

@@ -74,7 +74,7 @@ impl KbTradeAggregationService {
         };
         let mut results = std::vec::Vec::new();
         for decoded_event in &decoded_events {
-            if !decoded_event.event_kind.ends_with(".swap") {
+            if !kb_is_trade_event_kind(decoded_event.event_kind.as_str()) {
                 continue;
             }
             let decoded_event_id = match decoded_event.id {
@@ -135,6 +135,16 @@ impl KbTradeAggregationService {
                     )));
                 }
             };
+            let pool_tokens_result =
+                crate::list_pool_tokens_by_pool_id(self.database.as_ref(), pool_id).await;
+            let pool_tokens = match pool_tokens_result {
+                Ok(pool_tokens) => pool_tokens,
+                Err(error) => return Err(error),
+            };
+            let base_vault_address =
+                kb_find_pool_token_vault_address_by_token_id(&pool_tokens, pair.base_token_id);
+            let quote_vault_address =
+                kb_find_pool_token_vault_address_by_token_id(&pool_tokens, pair.quote_token_id);
             let payload_result =
                 serde_json::from_str::<serde_json::Value>(decoded_event.payload_json.as_str());
             let payload = match payload_result {
@@ -146,8 +156,8 @@ impl KbTradeAggregationService {
                     )));
                 }
             };
-            let trade_side = kb_extract_trade_side(&payload);
-            let base_amount_raw = kb_extract_amount_string(
+            let trade_side = kb_extract_trade_side(decoded_event.event_kind.as_str(), &payload);
+            let mut base_amount_raw = kb_extract_amount_string(
                 &payload,
                 &[
                     "baseAmountRaw",
@@ -157,7 +167,7 @@ impl KbTradeAggregationService {
                     "amountInBase",
                 ],
             );
-            let quote_amount_raw = kb_extract_amount_string(
+            let mut quote_amount_raw = kb_extract_amount_string(
                 &payload,
                 &[
                     "quoteAmountRaw",
@@ -167,45 +177,85 @@ impl KbTradeAggregationService {
                     "amountOutQuote",
                 ],
             );
-            let price_quote_per_base =
-                kb_compute_price_quote_per_base(base_amount_raw.clone(), quote_amount_raw.clone());
-            let slot_i64 = kb_convert_slot_to_i64(transaction.slot);
-            let created_trade_event = existing_trade_option.is_none();
-            let trade_event_id = if let Some(existing_trade) = existing_trade_option {
-                match existing_trade.id {
-                    Some(trade_event_id) => trade_event_id,
-                    None => {
-                        return Err(crate::KbError::InvalidState(
-                            "trade event has no internal id".to_string(),
-                        ));
-                    }
-                }
-            } else {
-                let trade_event_dto = crate::KbTradeEventDto::new(
-                    pool.dex_id,
-                    pool_id,
-                    pair_id,
-                    transaction_id,
-                    decoded_event_id,
-                    transaction.signature.clone(),
-                    slot_i64,
-                    trade_side,
-                    pair.base_token_id,
-                    pair.quote_token_id,
-                    base_amount_raw.clone(),
-                    quote_amount_raw.clone(),
-                    price_quote_per_base,
-                    crate::KbObservationSourceKind::Dex,
-                    transaction.source_endpoint_name.clone(),
-                    decoded_event.payload_json.clone(),
+            let mut price_quote_per_base = None;
+            if decoded_event.event_kind.starts_with("pump_swap.")
+                && (base_amount_raw.is_none()
+                    || quote_amount_raw.is_none()
+                    || price_quote_per_base.is_none())
+            {
+                let inferred_result = kb_extract_pump_swap_amounts_from_transaction(
+                    transaction.transaction_json.as_str(),
+                    transaction.meta_json.as_deref(),
+                    base_vault_address.as_deref(),
+                    quote_vault_address.as_deref(),
                 );
-                let upsert_result =
-                    crate::upsert_trade_event(self.database.as_ref(), &trade_event_dto).await;
-                match upsert_result {
-                    Ok(trade_event_id) => trade_event_id,
+                let inferred = match inferred_result {
+                    Ok(inferred) => inferred,
                     Err(error) => return Err(error),
+                };
+                if base_amount_raw.is_none() {
+                    base_amount_raw = inferred.0;
                 }
+                if quote_amount_raw.is_none() {
+                    quote_amount_raw = inferred.1;
+                }
+                if price_quote_per_base.is_none() {
+                    price_quote_per_base = inferred.2;
+                }
+            }
+            if price_quote_per_base.is_none() {
+                price_quote_per_base = kb_compute_price_quote_per_base_with_decimals(
+                    transaction.meta_json.as_deref(),
+                    transaction.transaction_json.as_str(),
+                    base_vault_address.as_deref(),
+                    quote_vault_address.as_deref(),
+                );
+            }
+            let slot_i64 = kb_convert_slot_to_i64(transaction.slot);
+            let existing_trade_was_empty = match &existing_trade_option {
+                Some(existing_trade) => {
+                    existing_trade.base_amount_raw.is_none()
+                        && existing_trade.quote_amount_raw.is_none()
+                        && existing_trade.price_quote_per_base.is_none()
+                }
+                None => false,
             };
+            let trade_event_dto = crate::KbTradeEventDto::new(
+                pool.dex_id,
+                pool_id,
+                pair_id,
+                transaction_id,
+                decoded_event_id,
+                transaction.signature.clone(),
+                slot_i64,
+                trade_side,
+                pair.base_token_id,
+                pair.quote_token_id,
+                base_amount_raw.clone(),
+                quote_amount_raw.clone(),
+                price_quote_per_base,
+                crate::KbObservationSourceKind::Dex,
+                transaction.source_endpoint_name.clone(),
+                decoded_event.payload_json.clone(),
+            );
+            tracing::debug!(
+                event_kind = %decoded_event.event_kind,
+                pool_account = ?decoded_event.pool_account,
+                decoded_event_id = ?decoded_event.id,
+                "trade aggregation candidate"
+            );
+            let upsert_result =
+                crate::upsert_trade_event(self.database.as_ref(), &trade_event_dto).await;
+            let trade_event_id = match upsert_result {
+                Ok(trade_event_id) => trade_event_id,
+                Err(error) => return Err(error),
+            };
+            let created_trade_event = existing_trade_option.is_none();
+            let repaired_trade_event = !created_trade_event
+                && existing_trade_was_empty
+                && (base_amount_raw.is_some()
+                    || quote_amount_raw.is_some()
+                    || price_quote_per_base.is_some());
             let pair_metric_result =
                 crate::get_pair_metric_by_pair_id(self.database.as_ref(), pair_id).await;
             let pair_metric_option = match pair_metric_result {
@@ -221,7 +271,7 @@ impl KbTradeAggregationService {
                         ));
                     }
                 };
-                if created_trade_event {
+                if created_trade_event || repaired_trade_event {
                     let mut updated_metric = existing_metric.clone();
                     kb_apply_trade_to_pair_metric(
                         &mut updated_metric,
@@ -310,6 +360,19 @@ impl KbTradeAggregationService {
     }
 }
 
+fn kb_is_trade_event_kind(event_kind: &str) -> bool {
+    if event_kind.ends_with(".swap") {
+        return true;
+    }
+    if event_kind.ends_with(".buy") {
+        return true;
+    }
+    if event_kind.ends_with(".sell") {
+        return true;
+    }
+    false
+}
+
 fn kb_convert_slot_to_i64(slot: std::option::Option<u64>) -> std::option::Option<i64> {
     match slot {
         Some(slot) => match i64::try_from(slot) {
@@ -320,13 +383,20 @@ fn kb_convert_slot_to_i64(slot: std::option::Option<u64>) -> std::option::Option
     }
 }
 
-fn kb_extract_trade_side(payload: &serde_json::Value) -> crate::KbSwapTradeSide {
+fn kb_extract_trade_side(event_kind: &str, payload: &serde_json::Value) -> crate::KbSwapTradeSide {
     let trade_side_option = kb_extract_string_by_candidate_keys(payload, &["tradeSide"]);
     match trade_side_option.as_deref() {
-        Some("BuyBase") => crate::KbSwapTradeSide::BuyBase,
-        Some("SellBase") => crate::KbSwapTradeSide::SellBase,
-        _ => crate::KbSwapTradeSide::Unknown,
+        Some("BuyBase") => return crate::KbSwapTradeSide::BuyBase,
+        Some("SellBase") => return crate::KbSwapTradeSide::SellBase,
+        _ => {}
     }
+    if event_kind.ends_with(".buy") {
+        return crate::KbSwapTradeSide::BuyBase;
+    }
+    if event_kind.ends_with(".sell") {
+        return crate::KbSwapTradeSide::SellBase;
+    }
+    crate::KbSwapTradeSide::Unknown
 }
 
 fn kb_extract_amount_string(
@@ -336,34 +406,6 @@ fn kb_extract_amount_string(
     kb_extract_scalar_as_string_by_candidate_keys(payload, candidate_keys)
 }
 
-fn kb_compute_price_quote_per_base(
-    base_amount_raw: std::option::Option<std::string::String>,
-    quote_amount_raw: std::option::Option<std::string::String>,
-) -> std::option::Option<f64> {
-    let base_amount_text = match base_amount_raw {
-        Some(base_amount_text) => base_amount_text,
-        None => return None,
-    };
-    let quote_amount_text = match quote_amount_raw {
-        Some(quote_amount_text) => quote_amount_text,
-        None => return None,
-    };
-    let base_amount_result = base_amount_text.parse::<f64>();
-    let base_amount = match base_amount_result {
-        Ok(base_amount) => base_amount,
-        Err(_) => return None,
-    };
-    if base_amount <= 0.0 {
-        return None;
-    }
-    let quote_amount_result = quote_amount_text.parse::<f64>();
-    let quote_amount = match quote_amount_result {
-        Ok(quote_amount) => quote_amount,
-        Err(_) => return None,
-    };
-    Some(quote_amount / base_amount)
-}
-
 fn kb_apply_trade_to_pair_metric(
     metric: &mut crate::KbPairMetricDto,
     slot: std::option::Option<i64>,
@@ -495,10 +537,292 @@ fn kb_extract_scalar_as_string_by_candidate_keys(
             }
         }
     }
-
     None
 }
 
+fn kb_find_pool_token_vault_address_by_token_id(
+    pool_tokens: &[crate::KbPoolTokenDto],
+    token_id: i64,
+) -> std::option::Option<std::string::String> {
+    for pool_token in pool_tokens {
+        if pool_token.token_id != token_id {
+            continue;
+        }
+        let vault_address_option = pool_token.vault_address.clone();
+        let vault_address = match vault_address_option {
+            Some(vault_address) => vault_address.trim().to_string(),
+            None => continue,
+        };
+        if vault_address.is_empty() {
+            continue;
+        }
+        return Some(vault_address);
+    }
+    None
+}
+
+fn kb_extract_pump_swap_amounts_from_transaction(
+    transaction_json: &str,
+    meta_json: std::option::Option<&str>,
+    base_vault_address: std::option::Option<&str>,
+    quote_vault_address: std::option::Option<&str>,
+) -> Result<
+    (
+        std::option::Option<std::string::String>,
+        std::option::Option<std::string::String>,
+        std::option::Option<f64>,
+    ),
+    crate::KbError,
+> {
+    let meta_json = match meta_json {
+        Some(meta_json) => meta_json,
+        None => return Ok((None, None, None)),
+    };
+    let transaction_value_result = serde_json::from_str::<serde_json::Value>(transaction_json);
+    let transaction_value = match transaction_value_result {
+        Ok(transaction_value) => transaction_value,
+        Err(error) => {
+            return Err(crate::KbError::Json(format!(
+                "cannot parse transaction_json for pump_swap amount extraction: {}",
+                error
+            )));
+        }
+    };
+    let meta_value_result = serde_json::from_str::<serde_json::Value>(meta_json);
+    let meta_value = match meta_value_result {
+        Ok(meta_value) => meta_value,
+        Err(error) => {
+            return Err(crate::KbError::Json(format!(
+                "cannot parse meta_json for pump_swap amount extraction: {}",
+                error
+            )));
+        }
+    };
+    let account_keys_result = kb_extract_transaction_account_keys(&transaction_value);
+    let account_keys = match account_keys_result {
+        Ok(account_keys) => account_keys,
+        Err(error) => return Err(error),
+    };
+    let pre_balances_result =
+        kb_extract_token_balance_map(&meta_value, &account_keys, "preTokenBalances");
+    let pre_balances = match pre_balances_result {
+        Ok(pre_balances) => pre_balances,
+        Err(error) => return Err(error),
+    };
+    let post_balances_result =
+        kb_extract_token_balance_map(&meta_value, &account_keys, "postTokenBalances");
+    let post_balances = match post_balances_result {
+        Ok(post_balances) => post_balances,
+        Err(error) => return Err(error),
+    };
+    let mut base_amount_raw = None;
+    let mut quote_amount_raw = None;
+    let mut price_quote_per_base = None;
+    if let Some(base_vault_address) = base_vault_address {
+        let base_pre = pre_balances.get(base_vault_address);
+        let base_post = post_balances.get(base_vault_address);
+        let base_pre_raw = base_pre.map(|value| value.0.clone());
+        let base_post_raw = base_post.map(|value| value.0.clone());
+        base_amount_raw = kb_compute_amount_delta_abs(base_pre_raw, base_post_raw);
+        let base_pre_ui = base_pre.and_then(|value| value.1);
+        let base_post_ui = base_post.and_then(|value| value.1);
+        let base_delta_ui = kb_compute_ui_delta_abs(base_pre_ui, base_post_ui);
+        if let Some(quote_vault_address) = quote_vault_address {
+            let quote_pre = pre_balances.get(quote_vault_address);
+            let quote_post = post_balances.get(quote_vault_address);
+            let quote_pre_raw = quote_pre.map(|value| value.0.clone());
+            let quote_post_raw = quote_post.map(|value| value.0.clone());
+            quote_amount_raw = kb_compute_amount_delta_abs(quote_pre_raw, quote_post_raw);
+            let quote_pre_ui = quote_pre.and_then(|value| value.1);
+            let quote_post_ui = quote_post.and_then(|value| value.1);
+            let quote_delta_ui = kb_compute_ui_delta_abs(quote_pre_ui, quote_post_ui);
+            match (base_delta_ui, quote_delta_ui) {
+                (Some(base_delta_ui), Some(quote_delta_ui)) => {
+                    if base_delta_ui > 0.0 {
+                        price_quote_per_base = Some(quote_delta_ui / base_delta_ui);
+                    }
+                }
+                _ => {}
+            }
+        }
+    }
+    Ok((base_amount_raw, quote_amount_raw, price_quote_per_base))
+}
+
+fn kb_extract_transaction_account_keys(
+    transaction_value: &serde_json::Value,
+) -> Result<std::vec::Vec<std::string::String>, crate::KbError> {
+    let candidate_arrays = [
+        transaction_value
+            .get("message")
+            .and_then(|value| value.get("accountKeys")),
+        transaction_value
+            .get("transaction")
+            .and_then(|value| value.get("message"))
+            .and_then(|value| value.get("accountKeys")),
+        transaction_value.get("accountKeys"),
+    ];
+    for candidate_array_option in candidate_arrays {
+        let candidate_array = match candidate_array_option {
+            Some(candidate_array) => candidate_array,
+            None => continue,
+        };
+        let array = match candidate_array.as_array() {
+            Some(array) => array,
+            None => continue,
+        };
+        let mut account_keys = std::vec::Vec::new();
+        for item in array {
+            if let Some(value) = item.as_str() {
+                account_keys.push(value.to_string());
+                continue;
+            }
+            let pubkey_option = item.get("pubkey").and_then(|value| value.as_str());
+            if let Some(pubkey) = pubkey_option {
+                account_keys.push(pubkey.to_string());
+                continue;
+            }
+        }
+        if !account_keys.is_empty() {
+            return Ok(account_keys);
+        }
+    }
+    Err(crate::KbError::Json(
+        "cannot extract accountKeys from transaction_json".to_string(),
+    ))
+}
+
+fn kb_extract_token_balance_map(
+    meta_value: &serde_json::Value,
+    account_keys: &[std::string::String],
+    field_name: &str,
+) -> Result<
+    std::collections::BTreeMap<
+        std::string::String,
+        (std::string::String, std::option::Option<f64>),
+    >,
+    crate::KbError,
+> {
+    let mut result = std::collections::BTreeMap::<
+        std::string::String,
+        (std::string::String, std::option::Option<f64>),
+    >::new();
+    let balances_option = meta_value
+        .get(field_name)
+        .and_then(|value| value.as_array());
+    let balances = match balances_option {
+        Some(balances) => balances,
+        None => return Ok(result),
+    };
+    for balance in balances {
+        let account_index_option = balance.get("accountIndex").and_then(|value| value.as_u64());
+        let account_index = match account_index_option {
+            Some(account_index) => account_index as usize,
+            None => continue,
+        };
+        if account_index >= account_keys.len() {
+            continue;
+        }
+        let account_address = account_keys[account_index].clone();
+        let ui_token_amount_option = balance.get("uiTokenAmount");
+        let ui_token_amount = match ui_token_amount_option {
+            Some(ui_token_amount) => ui_token_amount,
+            None => continue,
+        };
+        let raw_amount_option = ui_token_amount
+            .get("amount")
+            .and_then(|value| value.as_str());
+        let raw_amount = match raw_amount_option {
+            Some(raw_amount) => raw_amount.to_string(),
+            None => continue,
+        };
+        let ui_amount_string_option = ui_token_amount
+            .get("uiAmountString")
+            .and_then(|value| value.as_str());
+        let ui_amount = match ui_amount_string_option {
+            Some(ui_amount_string) => {
+                let parse_result = ui_amount_string.parse::<f64>();
+                match parse_result {
+                    Ok(ui_amount) => Some(ui_amount),
+                    Err(_) => None,
+                }
+            }
+            None => None,
+        };
+        result.insert(account_address, (raw_amount, ui_amount));
+    }
+    Ok(result)
+}
+
+fn kb_compute_amount_delta_abs(
+    pre_amount: std::option::Option<std::string::String>,
+    post_amount: std::option::Option<std::string::String>,
+) -> std::option::Option<std::string::String> {
+    let pre_amount = match pre_amount {
+        Some(pre_amount) => pre_amount,
+        None => "0".to_string(),
+    };
+    let post_amount = match post_amount {
+        Some(post_amount) => post_amount,
+        None => "0".to_string(),
+    };
+    let pre_value_result = pre_amount.parse::<i128>();
+    let pre_value = match pre_value_result {
+        Ok(pre_value) => pre_value,
+        Err(_) => return None,
+    };
+    let post_value_result = post_amount.parse::<i128>();
+    let post_value = match post_value_result {
+        Ok(post_value) => post_value,
+        Err(_) => return None,
+    };
+    let delta = if post_value >= pre_value {
+        post_value - pre_value
+    } else {
+        pre_value - post_value
+    };
+    Some(delta.to_string())
+}
+
+fn kb_compute_ui_delta_abs(
+    pre_amount: std::option::Option<f64>,
+    post_amount: std::option::Option<f64>,
+) -> std::option::Option<f64> {
+    let pre_amount = match pre_amount {
+        Some(pre_amount) => pre_amount,
+        None => 0.0,
+    };
+    let post_amount = match post_amount {
+        Some(post_amount) => post_amount,
+        None => 0.0,
+    };
+    let delta = if post_amount >= pre_amount {
+        post_amount - pre_amount
+    } else {
+        pre_amount - post_amount
+    };
+    Some(delta)
+}
+
+fn kb_compute_price_quote_per_base_with_decimals(
+    meta_json: std::option::Option<&str>,
+    transaction_json: &str,
+    base_vault_address: std::option::Option<&str>,
+    quote_vault_address: std::option::Option<&str>,
+) -> std::option::Option<f64> {
+    let inferred_result = kb_extract_pump_swap_amounts_from_transaction(
+        transaction_json,
+        meta_json,
+        base_vault_address,
+        quote_vault_address,
+    );
+    let inferred = match inferred_result {
+        Ok(inferred) => inferred,
+        Err(_) => return None,
+    };
+    inferred.2
+}
+
 #[cfg(test)]
 mod tests {
     async fn make_database() -> std::sync::Arc<crate::KbDatabase> {