refactor(all): use graphql_client crate

Signed-off-by: Gustavo Inacio <[email protected]>

Author: gusinacio

Cargo.lock

@@ -36,6 +36,8 @@ tower-http = { version = "0.5.2", features = [
     "trace",
 ] }
 tokio-util = "0.7.10"
+serde-inline-default = "0.2.0"
+graphql_client = { version = "0.14.0", features = ["reqwest-rustls"] }
 
 [dev-dependencies]
 env_logger = { version = "0.11.0", default-features = false }

common/Cargo.toml

@@ -3,16 +3,58 @@
 
 use std::{
     collections::HashMap,
+    str::FromStr,
     time::{Duration, SystemTime, UNIX_EPOCH},
 };
 
 use super::Allocation;
 use crate::prelude::SubgraphClient;
+use alloy::primitives::U256;
 use eventuals::{timer, Eventual, EventualExt};
-use thegraph_core::Address;
+use graphql_client::GraphQLQuery;
+use thegraph_core::{Address, DeploymentId};
 use tokio::time::sleep;
 use tracing::warn;
 
+type BigInt = String;
+type Bytes = String;
+
+#[derive(GraphQLQuery)]
+#[graphql(
+    schema_path = "../graphql/network.schema.graphql",
+    query_path = "../graphql/allocations.query.graphql",
+    response_derives = "Debug",
+    variables_derives = "Clone"
+)]
+pub struct AllocationsQuery;
+
+impl TryFrom<allocations_query::AllocationFragment> for Allocation {
+    type Error = anyhow::Error;
+
+    fn try_from(
+        value: allocations_query::AllocationsQueryAllocations,
+    ) -> Result<Self, Self::Error> {
+        Ok(Self {
+            id: Address::from_str(&value.id)?,
+            status: super::AllocationStatus::Null,
+            subgraph_deployment: super::SubgraphDeployment {
+                id: DeploymentId::from_str(&value.id)?,
+                denied_at: Some(value.subgraph_deployment.denied_at as u64),
+            },
+            indexer: Address::from_str(&value.indexer.id)?,
+            allocated_tokens: U256::from_str(&value.allocated_tokens)?,
+            created_at_epoch: value.created_at_epoch as u64,
+            created_at_block_hash: value.created_at_block_hash,
+            closed_at_epoch: value.closed_at_epoch.map(|v| v as u64),
+            closed_at_epoch_start_block_hash: None,
+            previous_epoch_start_block_hash: None,
+            poi: None,
+            query_fee_rebates: None,
+            query_fees_collected: None,
+        })
+    }
+}
+
 /// An always up-to-date list of an indexer's active and recently closed allocations.
 pub fn indexer_allocations(
     network_subgraph: &'static SubgraphClient,
@@ -56,49 +98,23 @@ pub async fn get_allocations(
         .expect("Time went backwards");
     let closed_at_threshold = since_the_epoch - recently_closed_allocation_buffer;
 
-    let query = format!(
-        r#"
-            allocations(
-                block: $block
-                orderBy: id
-                orderDirection: asc
-                first: $first
-                where: {{
-                and: [
-                    {{ id_gt: $last }}
-                    {{ indexer_: {{ id: "{}" }} }}
-                    {{
-                    or: [
-                        {{ status: Active }}
-                        {{ and: [{{ status: Closed, closedAt_gte: {} }}] }}
-                    ]
-                    }}
-                ]
-                }}
-            ) {{
-                id
-                indexer {{
-                    id
-                }}
-                allocatedTokens
-                createdAtBlockHash
-                createdAtEpoch
-                closedAtEpoch
-                subgraphDeployment {{
-                    id
-                    deniedAt
-                }}
-            }}
-        "#,
-        indexer_address.to_string().to_ascii_lowercase(),
-        closed_at_threshold.as_secs(),
-    );
     let responses = network_subgraph
-        .paginated_query::<Allocation>(query, 200)
+        .paginated_query::<AllocationsQuery, _>(
+            allocations_query::Variables {
+                indexer: indexer_address.to_string().to_ascii_lowercase(),
+                closed_at_threshold: closed_at_threshold.as_secs() as i64,
+            },
+            200,
+        )
         .await
         .map_err(|e| anyhow::anyhow!(e.to_string()))?;
 
-    Ok(HashMap::from_iter(responses.into_iter().map(|a| (a.id, a))))
+    Ok(HashMap::from_iter(
+        responses
+            .into_iter()
+            .flat_map(|a| a.allocations)
+            .map(|a| (Address::from_str(&a.id).unwrap(), a.try_into().unwrap())),
+    ))
 }
 
 #[cfg(test)]

common/src/allocations/monitor.rs

@@ -4,41 +4,32 @@
 use std::time::Duration;
 
 use eventuals::{timer, Eventual, EventualExt};
-use serde::Deserialize;
+use graphql_client::GraphQLQuery;
 use thegraph_core::Address;
 use tokio::time::sleep;
 use tracing::warn;
 
-use crate::subgraph_client::{Query, SubgraphClient};
+use crate::subgraph_client::SubgraphClient;
+
+type Bytes = Address;
+
+#[derive(GraphQLQuery)]
+#[graphql(
+    schema_path = "../graphql/network.schema.graphql",
+    query_path = "../graphql/dispute.query.graphql",
+    response_derives = "Debug",
+    variables_derives = "Clone"
+)]
+struct DisputeManager;
 
 pub fn dispute_manager(
     network_subgraph: &'static SubgraphClient,
     interval: Duration,
 ) -> Eventual<Address> {
-    #[derive(Deserialize)]
-    #[serde(rename_all = "camelCase")]
-    struct DisputeManagerResponse {
-        graph_network: Option<GraphNetwork>,
-    }
-
-    #[derive(Deserialize)]
-    #[serde(rename_all = "camelCase")]
-    struct GraphNetwork {
-        dispute_manager: Address,
-    }
-
     timer(interval).map_with_retry(
         move |_| async move {
             let response = network_subgraph
-                .query::<DisputeManagerResponse>(Query::new(
-                    r#"
-                        query network {
-                            graphNetwork(id: 1) {
-                                disputeManager
-                            }
-                        }
-                    "#,
-                ))
+                .query::<DisputeManager, _>(dispute_manager::Variables {})
                 .await
                 .map_err(|e| e.to_string())?;
 
@@ -49,7 +40,7 @@ pub fn dispute_manager(
             })
         },
         move |err: String| {
-            warn!("Failed to query dispute manager for network: {}", err,);
+            warn!("Failed to query dispute manager: {}", err);
 
             // Sleep for a bit before we retry
             sleep(interval.div_f32(2.0))

common/src/attestations/dispute_manager.rs

@@ -10,13 +10,13 @@ use std::{
 use alloy::primitives::U256;
 use anyhow::Result;
 use eventuals::{timer, Eventual, EventualExt};
-use serde::Deserialize;
+use graphql_client::GraphQLQuery;
 use thegraph_core::Address;
 use thiserror::Error;
 use tokio::time::sleep;
 use tracing::{error, warn};
 
-use crate::prelude::{Query, SubgraphClient};
+use crate::prelude::SubgraphClient;
 
 #[derive(Error, Debug)]
 pub enum EscrowAccountsError {
@@ -89,110 +89,59 @@ impl EscrowAccounts {
     }
 }
 
+type BigInt = U256;
+
+#[derive(GraphQLQuery)]
+#[graphql(
+    schema_path = "../graphql/tap.schema.graphql",
+    query_path = "../graphql/escrow_account.query.graphql",
+    response_derives = "Debug",
+    variables_derives = "Clone"
+)]
+pub struct EscrowAccountQuery;
+
 pub fn escrow_accounts(
     escrow_subgraph: &'static SubgraphClient,
     indexer_address: Address,
     interval: Duration,
     reject_thawing_signers: bool,
 ) -> Eventual<EscrowAccounts> {
-    // Types for deserializing the network subgraph response
-    #[derive(Deserialize)]
-    #[serde(rename_all = "camelCase")]
-    struct EscrowAccountsResponse {
-        escrow_accounts: Vec<EscrowAccount>,
-    }
-    // Note that U256's serde implementation is based on serializing the internal bytes, not the string decimal
-    // representation. This is why we deserialize them as strings below.
-    #[derive(Deserialize)]
-    #[serde(rename_all = "camelCase")]
-    struct EscrowAccount {
-        balance: String,
-        total_amount_thawing: String,
-        sender: Sender,
-    }
-    #[derive(Deserialize)]
-    #[serde(rename_all = "camelCase")]
-    struct Sender {
-        id: Address,
-        signers: Vec<Signer>,
-    }
-    #[derive(Deserialize)]
-    #[serde(rename_all = "camelCase")]
-    struct Signer {
-        id: Address,
-    }
-
-    // thawEndTimestamp == 0 means that the signer is not thawing. This also means
-    // that we don't wait for the thawing period to end before stopping serving
-    // queries for this signer.
-    // isAuthorized == true means that the signer is still authorized to sign
-    // payments in the name of the sender.
-    let query = if reject_thawing_signers {
-        r#"
-        query ($indexer: ID!) {
-            escrowAccounts(where: {receiver_: {id: $indexer}}) {
-                balance
-                totalAmountThawing
-                sender {
-                    id
-                    signers(
-                        where: {thawEndTimestamp: "0", isAuthorized: true}
-                    ) {
-                        id
-                    }
-                }
-            }
-        }
-    "#
-    } else {
-        r#"
-        query ($indexer: ID!) {
-            escrowAccounts(where: {receiver_: {id: $indexer}}) {
-                balance
-                totalAmountThawing
-                sender {
-                    id
-                    signers(
-                        where: {isAuthorized: true}
-                    ) {
-                        id
-                    }
-                }
-            }
-        }
-    "#
-    };
-
     timer(interval).map_with_retry(
         move |_| async move {
+            // thawEndTimestamp == 0 means that the signer is not thawing. This also means
+            // that we don't wait for the thawing period to end before stopping serving
+            // queries for this signer.
+            // isAuthorized == true means that the signer is still authorized to sign
+            // payments in the name of the sender.
             let response = escrow_subgraph
-                .query::<EscrowAccountsResponse>(Query::new_with_variables(
-                    query,
-                    [("indexer", format!("{:x?}", indexer_address).into())],
-                ))
+                .query::<EscrowAccountQuery, _>(escrow_account_query::Variables {
+                    indexer: format!("{:x?}", indexer_address),
+                    thaw_end_timestamp: if reject_thawing_signers {
+                        Some(U256::ZERO)
+                    } else {
+                        None
+                    },
+                })
                 .await
                 .map_err(|e| e.to_string())?;
 
-            let response = response.map_err(|e| e.to_string())?;
+            let response = response.unwrap();
 
-            let senders_balances = response
+            let senders_balances: HashMap<Address, U256> = response
                 .escrow_accounts
                 .iter()
                 .map(|account| {
-                    let balance = U256::checked_sub(
-                        U256::from_str(&account.balance)?,
-                        U256::from_str(&account.total_amount_thawing)?,
-                    )
-                    .unwrap_or_else(|| {
-                        warn!(
-                            "Balance minus total amount thawing underflowed for account {}. \
+                    let balance = U256::checked_sub(account.balance, account.total_amount_thawing)
+                        .unwrap_or_else(|| {
+                            warn!(
+                                "Balance minus total amount thawing underflowed for account {}. \
                                  Setting balance to 0, no queries will be served for this sender.",
-                            account.sender.id
-                        );
-                        U256::from(0)
-                    });
+                                account.sender.id
+                            );
+                            U256::from(0)
+                        });
 
-                    Ok((account.sender.id, balance))
+                    Ok((Address::from_str(&account.sender.id).unwrap(), balance))
                 })
                 .collect::<Result<HashMap<_, _>, anyhow::Error>>()
                 .map_err(|e| format!("{}", e))?;
@@ -201,12 +150,14 @@ pub fn escrow_accounts(
                 .escrow_accounts
                 .iter()
                 .map(|account| {
-                    let sender = account.sender.id;
+                    let sender = Address::from_str(&account.sender.id).unwrap();
                     let signers = account
                         .sender
                         .signers
+                        .as_ref()
+                        .unwrap()
                         .iter()
-                        .map(|signer| signer.id)
+                        .map(|signer| Address::from_str(&signer.id).unwrap())
                         .collect();
                     (sender, signers)
                 })