add possibility to filter by processed to get_batch_by_index

Author: jonastheis

crates/chain-orchestrator/src/lib.rs

@@ -722,12 +722,14 @@ impl<
 
                     // Perform a consistency check to ensure the previous commit batch exists in the
                     // database.
-                    if tx.get_batch_by_index(prev_batch_index).await?.is_none() {
+                    if tx.get_batch_by_index(prev_batch_index, None).await?.is_none() {
                         return Err(ChainOrchestratorError::BatchCommitGap(batch_clone.index));
                     }
 
                     // Check if batch already exists in DB.
-                    if let Some(existing_batch) = tx.get_batch_by_index(batch_clone.index).await? {
+                    if let Some(existing_batch) =
+                        tx.get_batch_by_index(batch_clone.index, Some(true)).await?
+                    {
                         if existing_batch.hash == batch_clone.hash {
                             // This means we have already processed this batch commit, we will skip
                             // it.

crates/database/db/src/db.rs

@@ -413,11 +413,12 @@ impl DatabaseReadOperations for Database {
     async fn get_batch_by_index(
         &self,
         batch_index: u64,
+        processed: Option<bool>,
     ) -> Result<Option<BatchCommitData>, DatabaseError> {
         metered!(
             DatabaseOperation::GetBatchByIndex,
             self,
-            tx(move |tx| async move { tx.get_batch_by_index(batch_index).await })
+            tx(move |tx| async move { tx.get_batch_by_index(batch_index, processed).await })
         )
     }
 
@@ -735,7 +736,7 @@ mod test {
         // Round trip the BatchCommitData through the database.
         db.insert_batch(batch_commit.clone()).await.unwrap();
         let batch_commit_from_db =
-            db.get_batch_by_index(batch_commit.index).await.unwrap().unwrap();
+            db.get_batch_by_index(batch_commit.index, None).await.unwrap().unwrap();
 
         assert_eq!(batch_commit, batch_commit_from_db);
     }
@@ -1249,7 +1250,7 @@ mod test {
 
         // Insert L2 blocks with different batch indices
         for i in 100..110 {
-            let batch_data = db.get_batch_by_index(i).await.unwrap().unwrap();
+            let batch_data = db.get_batch_by_index(i, None).await.unwrap().unwrap();
             let batch_info: BatchInfo = batch_data.into();
             let block_info = BlockInfo { number: 500 + i, hash: B256::arbitrary(&mut u).unwrap() };
 
@@ -1418,9 +1419,9 @@ mod test {
         db.set_finalized_l1_block_number(21).await.unwrap();
 
         // Verify the batches and blocks were inserted correctly
-        let retrieved_batch_1 = db.get_batch_by_index(1).await.unwrap().unwrap();
-        let retrieved_batch_2 = db.get_batch_by_index(2).await.unwrap().unwrap();
-        let retrieved_batch_3 = db.get_batch_by_index(3).await.unwrap().unwrap();
+        let retrieved_batch_1 = db.get_batch_by_index(1, None).await.unwrap().unwrap();
+        let retrieved_batch_2 = db.get_batch_by_index(2, None).await.unwrap().unwrap();
+        let retrieved_batch_3 = db.get_batch_by_index(3, None).await.unwrap().unwrap();
         let retried_block_1 = db.get_l2_block_info_by_number(1).await.unwrap().unwrap();
         let retried_block_2 = db.get_l2_block_info_by_number(2).await.unwrap().unwrap();
         let retried_block_3 = db.get_l2_block_info_by_number(3).await.unwrap().unwrap();
@@ -1441,9 +1442,9 @@ mod test {
         assert_eq!(result, (Some(block_2), Some(11)));
 
         // Verify that batches 2 and 3 are deleted
-        let batch_1 = db.get_batch_by_index(1).await.unwrap();
-        let batch_2 = db.get_batch_by_index(2).await.unwrap();
-        let batch_3 = db.get_batch_by_index(3).await.unwrap();
+        let batch_1 = db.get_batch_by_index(1, None).await.unwrap();
+        let batch_2 = db.get_batch_by_index(2, None).await.unwrap();
+        let batch_3 = db.get_batch_by_index(3, None).await.unwrap();
         assert!(batch_1.is_some());
         assert!(batch_2.is_none());
         assert!(batch_3.is_none());

crates/database/db/src/operations.rs

@@ -374,7 +374,8 @@ impl<T: WriteConnectionProvider + ?Sized + Sync> DatabaseWriteOperations for T {
             .map(|(_, batch_info)| batch_info)
             .filter(|b| b.index > 1)
         {
-            let batch = self.get_batch_by_index(batch_info.index).await?.expect("batch must exist");
+            let batch =
+                self.get_batch_by_index(batch_info.index, None).await?.expect("batch must exist");
             self.delete_batches_gt_block_number(batch.block_number.saturating_sub(1)).await?;
         };
 
@@ -383,7 +384,8 @@ impl<T: WriteConnectionProvider + ?Sized + Sync> DatabaseWriteOperations for T {
         else {
             return Ok((None, None));
         };
-        let batch = self.get_batch_by_index(batch_info.index).await?.expect("batch must exist");
+        let batch =
+            self.get_batch_by_index(batch_info.index, None).await?.expect("batch must exist");
         Ok((Some(block_info), Some(batch.block_number.saturating_add(1))))
     }
 
@@ -649,6 +651,7 @@ pub trait DatabaseReadOperations {
     async fn get_batch_by_index(
         &self,
         batch_index: u64,
+        processed: Option<bool>,
     ) -> Result<Option<BatchCommitData>, DatabaseError>;
 
     /// Get the latest L1 block number from the database.
@@ -727,13 +730,21 @@ impl<T: ReadConnectionProvider + Sync + ?Sized> DatabaseReadOperations for T {
     async fn get_batch_by_index(
         &self,
         batch_index: u64,
+        processed: Option<bool>,
     ) -> Result<Option<BatchCommitData>, DatabaseError> {
-        Ok(models::batch_commit::Entity::find_by_id(
-            TryInto::<i64>::try_into(batch_index).expect("index should fit in i64"),
-        )
-        .one(self.get_connection())
-        .await
-        .map(|x| x.map(Into::into))?)
+        let query = if let Some(p) = processed {
+            models::batch_commit::Entity::find().filter(
+                models::batch_commit::Column::Index
+                    .eq(TryInto::<i64>::try_into(batch_index).expect("index should fit in i64"))
+                    .and(models::batch_commit::Column::Processed.eq(p)),
+            )
+        } else {
+            models::batch_commit::Entity::find_by_id(
+                TryInto::<i64>::try_into(batch_index).expect("index should fit in i64"),
+            )
+        };
+
+        Ok(query.one(self.get_connection()).await.map(|x| x.map(Into::into))?)
     }
 
     async fn get_latest_l1_block_number(&self) -> Result<u64, DatabaseError> {

crates/derivation-pipeline/src/lib.rs

@@ -215,7 +215,7 @@ where
 
             // get the batch commit data.
             let batch = db
-                .get_batch_by_index(batch_info.index)
+                .get_batch_by_index(batch_info.index, None)
                 .await
                 .map_err(|err| (batch_info.clone(), err.into()))?
                 .ok_or((