fix: add last_allocation_id to indexing_agreements, and improve get_by_id

Author: pcarranzav

.sqlx/query-47f848e049f3ff22e65bb53edc7ddc1646e68d6db58124b6e0d780c037f73513.json

@@ -310,11 +310,11 @@ pub async fn validate_and_cancel_agreement(
             decoded_request.request.agreement_id.into(),
         ))
         .await?;
-    let (agreement, cancelled) = result.ok_or(DipsError::AgreementNotFound)?;
-    if cancelled {
+    let stored_agreement = result.ok_or(DipsError::AgreementNotFound)?;
+    if stored_agreement.cancelled {
         return Err(DipsError::AgreementCancelled);
     }
-    let expected_signer = agreement.voucher.payer;
+    let expected_signer = stored_agreement.voucher.voucher.payer;
     let id = Uuid::from_bytes(decoded_request.request.agreement_id.into());
     decoded_request.validate(domain, &expected_signer)?;
 
@@ -391,11 +391,10 @@ mod test {
         .unwrap();
         assert_eq!(actual_id, id);
 
-        let actual = store.get_by_id(actual_id).await.unwrap();
+        let stored_agreement = store.get_by_id(actual_id).await.unwrap().unwrap();
 
-        let (actual_voucher, actual_cancelled) = actual.unwrap();
-        assert_eq!(voucher, actual_voucher);
-        assert!(!actual_cancelled);
+        assert_eq!(voucher, stored_agreement.voucher);
+        assert!(!stored_agreement.cancelled);
         Ok(())
     }
 
@@ -596,9 +595,8 @@ mod test {
         assert_eq!(agreement_id, cancelled_id);
 
         // Verify agreement is cancelled
-        let result = store.get_by_id(agreement_id).await?;
-        let (_, cancelled) = result.ok_or(DipsError::AgreementNotFound)?;
-        assert!(cancelled);
+        let stored_agreement = store.get_by_id(agreement_id).await?.unwrap();
+        assert!(stored_agreement.cancelled);
 
         Ok(())
     }

.sqlx/query-cb01d485323ed4b697d3c0c645f5174cc42ea3884e5b54ffb5c94605637d0baa.json renamed to .sqlx/query-5e4575f3b015895c5ba3f066c5d010a09c48290e2200fc1c6426234192ce0f82.json

@@ -11,12 +11,18 @@ use crate::{
     SubgraphIndexingVoucherMetadata,
 };
 
+#[derive(Debug, Clone)]
+pub struct StoredIndexingAgreement {
+    pub voucher: SignedIndexingAgreementVoucher,
+    pub metadata: SubgraphIndexingVoucherMetadata,
+    pub cancelled: bool,
+    pub current_allocation_id: Option<String>,
+    pub last_allocation_id: Option<String>,
+}
+
 #[async_trait]
 pub trait AgreementStore: Sync + Send + std::fmt::Debug {
-    async fn get_by_id(
-        &self,
-        id: Uuid,
-    ) -> Result<Option<(SignedIndexingAgreementVoucher, bool)>, DipsError>;
+    async fn get_by_id(&self, id: Uuid) -> Result<Option<StoredIndexingAgreement>, DipsError>;
     async fn create_agreement(
         &self,
         agreement: SignedIndexingAgreementVoucher,
@@ -30,15 +36,12 @@ pub trait AgreementStore: Sync + Send + std::fmt::Debug {
 
 #[derive(Default, Debug)]
 pub struct InMemoryAgreementStore {
-    pub data: tokio::sync::RwLock<HashMap<Uuid, (SignedIndexingAgreementVoucher, bool)>>,
+    pub data: tokio::sync::RwLock<HashMap<Uuid, StoredIndexingAgreement>>,
 }
 
 #[async_trait]
 impl AgreementStore for InMemoryAgreementStore {
-    async fn get_by_id(
-        &self,
-        id: Uuid,
-    ) -> Result<Option<(SignedIndexingAgreementVoucher, bool)>, DipsError> {
+    async fn get_by_id(&self, id: Uuid) -> Result<Option<StoredIndexingAgreement>, DipsError> {
         Ok(self
             .data
             .try_read()
@@ -49,15 +52,20 @@ impl AgreementStore for InMemoryAgreementStore {
     async fn create_agreement(
         &self,
         agreement: SignedIndexingAgreementVoucher,
-        _medatadata: SubgraphIndexingVoucherMetadata,
+        metadata: SubgraphIndexingVoucherMetadata,
     ) -> Result<(), DipsError> {
+        let id = Uuid::from_bytes(agreement.voucher.agreement_id.into());
+        let stored_agreement = StoredIndexingAgreement {
+            voucher: agreement,
+            metadata,
+            cancelled: false,
+            current_allocation_id: None,
+            last_allocation_id: None,
+        };
         self.data
             .try_write()
             .map_err(|e| DipsError::UnknownError(e.into()))?
-            .insert(
-                Uuid::from_bytes(agreement.voucher.agreement_id.into()),
-                (agreement.clone(), false),
-            );
+            .insert(id, stored_agreement);
 
         Ok(())
     }
@@ -67,7 +75,7 @@ impl AgreementStore for InMemoryAgreementStore {
     ) -> Result<Uuid, DipsError> {
         let id = Uuid::from_bytes(signed_cancellation.request.agreement_id.into());
 
-        let agreement = {
+        let mut agreement = {
             let read_lock = self
                 .data
                 .try_read()
@@ -78,11 +86,17 @@ impl AgreementStore for InMemoryAgreementStore {
                 .ok_or(DipsError::AgreementNotFound)?
         };
 
+        if agreement.cancelled {
+            return Err(DipsError::AgreementCancelled);
+        }
+
+        agreement.cancelled = true;
+
         let mut write_lock = self
             .data
             .try_write()
             .map_err(|e| DipsError::UnknownError(e.into()))?;
-        write_lock.insert(id, (agreement.0, true));
+        write_lock.insert(id, agreement);
 
         Ok(id)
     }

.sqlx/query-d0258d53c19174720f2d4d9cde481e457931b65eca471efed650b10d70d2282a.json

@@ -6,7 +6,8 @@ use std::str::FromStr;
 use axum::async_trait;
 use build_info::chrono::{DateTime, Utc};
 use indexer_dips::{
-    store::AgreementStore, DipsError, SignedCancellationRequest, SignedIndexingAgreementVoucher,
+    store::{AgreementStore, StoredIndexingAgreement},
+    DipsError, SignedCancellationRequest, SignedIndexingAgreementVoucher,
     SubgraphIndexingVoucherMetadata,
 };
 use sqlx::{types::BigDecimal, PgPool};
@@ -25,10 +26,7 @@ fn uint256_to_bigdecimal(value: &uint256, field: &str) -> Result<BigDecimal, Dip
 
 #[async_trait]
 impl AgreementStore for PsqlAgreementStore {
-    async fn get_by_id(
-        &self,
-        id: Uuid,
-    ) -> Result<Option<(SignedIndexingAgreementVoucher, bool)>, DipsError> {
+    async fn get_by_id(&self, id: Uuid) -> Result<Option<StoredIndexingAgreement>, DipsError> {
         let item = sqlx::query!("SELECT * FROM indexing_agreements WHERE id=$1", id,)
             .fetch_one(&self.pool)
             .await;
@@ -41,8 +39,17 @@ impl AgreementStore for PsqlAgreementStore {
 
         let signed = SignedIndexingAgreementVoucher::abi_decode(item.signed_payload.as_ref(), true)
             .map_err(|e| DipsError::AbiDecoding(e.to_string()))?;
+        let metadata =
+            SubgraphIndexingVoucherMetadata::abi_decode(signed.voucher.metadata.as_ref(), true)
+                .map_err(|e| DipsError::AbiDecoding(e.to_string()))?;
         let cancelled = item.cancelled_at.is_some();
-        Ok(Some((signed, cancelled)))
+        Ok(Some(StoredIndexingAgreement {
+            voucher: signed,
+            metadata,
+            cancelled,
+            current_allocation_id: item.current_allocation_id,
+            last_allocation_id: item.last_allocation_id,
+        }))
     }
     async fn create_agreement(
         &self,
@@ -72,7 +79,7 @@ impl AgreementStore for PsqlAgreementStore {
         let min_epochs_per_collection: i64 = agreement.voucher.minEpochsPerCollection.into();
         let max_epochs_per_collection: i64 = agreement.voucher.maxEpochsPerCollection.into();
         sqlx::query!(
-            "INSERT INTO indexing_agreements VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,null,null,null)",
+            "INSERT INTO indexing_agreements VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,null,null,null,null)",
             id,
             agreement.signature.as_ref(),
             bs,
@@ -130,7 +137,7 @@ pub(crate) mod test {
     use sqlx::PgPool;
     use thegraph_core::alloy::{
         primitives::{ruint::aliases::U256, Address},
-        sol_types::{SolType, SolValue},
+        sol_types::SolValue,
     };
     use uuid::Uuid;
 
@@ -226,19 +233,15 @@ pub(crate) mod test {
             .unwrap();
 
         // Retrieve agreement
-        let (retrieved_signed_voucher, cancelled) = store.get_by_id(id).await.unwrap().unwrap();
-
-        let retrieved_voucher = &retrieved_signed_voucher.voucher;
-        let retrieved_metadata =
-            <indexer_dips::SubgraphIndexingVoucherMetadata as SolType>::abi_decode(
-                retrieved_voucher.metadata.as_ref(),
-                true,
-            )
-            .unwrap();
+        let stored_agreement = store.get_by_id(id).await.unwrap().unwrap();
+
+        let retrieved_voucher = &stored_agreement.voucher;
+        let retrieved_metadata = stored_agreement.metadata;
+
         // Verify retrieved agreement matches original
-        assert_eq!(retrieved_signed_voucher.signature, agreement.signature);
+        assert_eq!(retrieved_voucher.signature, agreement.signature);
         assert_eq!(
-            retrieved_voucher.durationEpochs,
+            retrieved_voucher.voucher.durationEpochs,
             agreement.voucher.durationEpochs
         );
         assert_eq!(retrieved_metadata.protocolNetwork, metadata.protocolNetwork);
@@ -247,26 +250,29 @@ pub(crate) mod test {
             retrieved_metadata.subgraphDeploymentId,
             metadata.subgraphDeploymentId
         );
-        assert_eq!(retrieved_voucher.payer, agreement.voucher.payer);
-        assert_eq!(retrieved_voucher.recipient, agreement.voucher.recipient);
-        assert_eq!(retrieved_voucher.service, agreement.voucher.service);
+        assert_eq!(retrieved_voucher.voucher.payer, agreement.voucher.payer);
+        assert_eq!(
+            retrieved_voucher.voucher.recipient,
+            agreement.voucher.recipient
+        );
+        assert_eq!(retrieved_voucher.voucher.service, agreement.voucher.service);
         assert_eq!(
-            retrieved_voucher.maxInitialAmount,
+            retrieved_voucher.voucher.maxInitialAmount,
             agreement.voucher.maxInitialAmount
         );
         assert_eq!(
-            retrieved_voucher.maxOngoingAmountPerEpoch,
+            retrieved_voucher.voucher.maxOngoingAmountPerEpoch,
             agreement.voucher.maxOngoingAmountPerEpoch
         );
         assert_eq!(
-            retrieved_voucher.maxEpochsPerCollection,
+            retrieved_voucher.voucher.maxEpochsPerCollection,
             agreement.voucher.maxEpochsPerCollection
         );
         assert_eq!(
-            retrieved_voucher.minEpochsPerCollection,
+            retrieved_voucher.voucher.minEpochsPerCollection,
             agreement.voucher.minEpochsPerCollection
         );
-        assert!(!cancelled);
+        assert!(!stored_agreement.cancelled);
     }
 
     #[sqlx::test(migrations = "../../migrations")]

crates/dips/src/lib.rs

@@ -28,7 +28,8 @@ CREATE TABLE IF NOT EXISTS indexing_agreements (
     cancelled_at TIMESTAMP WITH TIME ZONE,
     signed_cancellation_payload BYTEA,
 
-    current_allocation_id CHAR(40)
+    current_allocation_id CHAR(40),
+    last_allocation_id CHAR(40)
 );
 
 CREATE UNIQUE INDEX IX_UNIQ_SIGNATURE_AGREEMENT on indexing_agreements(signature);