refactor: DRY block-lookup logic in atomic-data-publisher

Extract shared helpers block_boundary_for_table and next_block_for_table
to consolidate duplicated SQL queries across correction job functions.
Simplify test helpers with run_block_range_query.

Author: bryzettler

utils/atomic-data-publisher/src/database.rs

@@ -1009,15 +1009,6 @@ impl DatabaseClient {
     &self,
     query_name: &str,
     parameters: &serde_json::Value,
-  ) -> Result<Option<u64>> {
-    self.get_block_for_query(query_name, parameters, true).await
-  }
-
-  async fn get_block_for_query(
-    &self,
-    query_name: &str,
-    parameters: &serde_json::Value,
-    use_max: bool,
   ) -> Result<Option<u64>> {
     let pool = self.get_pool().await;
     let block = match query_name {
@@ -1031,26 +1022,12 @@ impl DatabaseClient {
             ))
           })?;
 
-        let (block, has_data) = match hotspot_type {
+        let query = match hotspot_type {
           "mobile" => {
-            let query = if use_max {
-              r#"SELECT COALESCE((SELECT MAX(last_block) FROM mobile_hotspot_infos), -1)::bigint as block"#
-            } else {
-              r#"SELECT COALESCE((SELECT MIN(last_block) FROM mobile_hotspot_infos), -1)::bigint as block"#
-            };
-            let row = sqlx::query(query).fetch_one(&*pool).await?;
-            let block_value = row.get::<i64, _>("block");
-            (block_value, block_value >= 0)
+            r#"SELECT COALESCE((SELECT MAX(last_block) FROM mobile_hotspot_infos), -1)::bigint as block"#
           }
           "iot" => {
-            let query = if use_max {
-              r#"SELECT COALESCE((SELECT MAX(last_block) FROM iot_hotspot_infos), -1)::bigint as block"#
-            } else {
-              r#"SELECT COALESCE((SELECT MIN(last_block) FROM iot_hotspot_infos), -1)::bigint as block"#
-            };
-            let row = sqlx::query(query).fetch_one(&*pool).await?;
-            let block_value = row.get::<i64, _>("block");
-            (block_value, block_value >= 0)
+            r#"SELECT COALESCE((SELECT MAX(last_block) FROM iot_hotspot_infos), -1)::bigint as block"#
           }
           _ => {
             return Err(anyhow::anyhow!(
@@ -1059,68 +1036,41 @@ impl DatabaseClient {
             ))
           }
         };
+        let row = sqlx::query(query).fetch_one(&*pool).await?;
+        let block_value = row.get::<i64, _>("block");
 
-        let operation = if use_max { "Max" } else { "Min" };
         debug!(
-          "{} last_block for {} hotspots: {}",
-          operation, hotspot_type, block
+          "Max last_block for {} hotspots: {}",
+          hotspot_type, block_value
         );
 
-        if has_data {
-          Some(block as u64)
+        if block_value >= 0 {
+          Some(block_value as u64)
         } else {
           None
         }
       }
-      "construct_entity_ownership_changes" => {
-        let query = if use_max {
-          r#"SELECT COALESCE((SELECT MAX(last_block) FROM asset_owners), -1)::bigint as block"#
-        } else {
-          r#"SELECT COALESCE((SELECT MIN(last_block) FROM asset_owners WHERE last_block >= 0), -1)::bigint as block"#
+      "construct_entity_ownership_changes"
+      | "construct_welcome_pack_ownership"
+      | "construct_welcome_pack_reward_destination" => {
+        let table = match query_name {
+          "construct_entity_ownership_changes" => "asset_owners",
+          _ => "welcome_packs",
         };
-        let row = sqlx::query(query).fetch_one(&*pool).await?;
-
-        let block: i64 = row.get("block");
-        if block >= 0 {
-          Some(block as u64)
-        } else {
-          None
-        }
+        Self::max_block_for_table(&pool, table).await?
       }
       "construct_entity_reward_destination_changes" => {
-        let query = if use_max {
+        let row = sqlx::query(
           r#"
           SELECT GREATEST(
             COALESCE((SELECT MAX(last_block) FROM asset_owners), -1),
             COALESCE((SELECT MAX(last_block) FROM recipients), -1),
             COALESCE((SELECT MAX(last_block) FROM mini_fanouts), -1)
           )::bigint as block
-          "#
-        } else {
-          r#"
-          SELECT LEAST(
-            COALESCE((SELECT MIN(last_block) FROM asset_owners WHERE last_block >= 0), -1),
-            COALESCE((SELECT MIN(last_block) FROM recipients WHERE last_block >= 0), -1),
-            COALESCE((SELECT MIN(last_block) FROM mini_fanouts WHERE last_block >= 0), -1)
-          )::bigint as block
-          "#
-        };
-        let row = sqlx::query(query).fetch_one(&*pool).await?;
-
-        let block: i64 = row.get("block");
-        if block >= 0 {
-          Some(block as u64)
-        } else {
-          None
-        }
-      }
-      "construct_welcome_pack_ownership" | "construct_welcome_pack_reward_destination" => {
-        let query = if use_max {
-          r#"SELECT COALESCE((SELECT MAX(last_block) FROM welcome_packs), -1)::bigint as block"#
-        } else {
-          r#"SELECT COALESCE((SELECT MIN(last_block) FROM welcome_packs WHERE last_block >= 0), -1)::bigint as block"#
-        };
-        let row = sqlx::query(query).fetch_one(&*pool).await?;
+          "#,
+        )
+        .fetch_one(&*pool)
+        .await?;
 
         let block: i64 = row.get("block");
         if block >= 0 {
@@ -1135,15 +1085,36 @@ impl DatabaseClient {
       }
     };
 
-    let operation = if use_max { "Max" } else { "Min" };
     debug!(
-      "{} last_block for query '{}': {:?}",
-      operation, query_name, block
+      "Max last_block for query '{}': {:?}",
+      query_name, block
     );
 
     Ok(block)
   }
 
+  async fn max_block_for_table(pool: &PgPool, table: &str) -> Result<Option<u64>> {
+    let query =
+      format!(r#"SELECT COALESCE((SELECT MAX(last_block) FROM {table}), -1)::bigint as block"#);
+    let row = sqlx::query(&query).fetch_one(pool).await?;
+    let block: i64 = row.get("block");
+    Ok(if block >= 0 { Some(block as u64) } else { None })
+  }
+
+  async fn next_block_for_table(
+    pool: &PgPool,
+    table: &str,
+    after_block: u64,
+  ) -> Result<Option<u64>> {
+    let query =
+      format!(r#"SELECT MIN(last_block)::bigint as block FROM {table} WHERE last_block > $1"#);
+    let row = sqlx::query(&query)
+      .bind(after_block as i64)
+      .fetch_one(pool)
+      .await?;
+    Ok(row.get::<Option<i64>, _>("block").map(|b| b as u64))
+  }
+
   async fn get_next_data_block_after(
     &self,
     job: &PollingJob,
@@ -1187,17 +1158,14 @@ impl DatabaseClient {
 
         row.get::<Option<i64>, _>("block").map(|b| b as u64)
       }
-      "construct_entity_ownership_changes" => {
-        let query = r#"
-          SELECT MIN(last_block)::bigint as block FROM asset_owners WHERE last_block > $1
-        "#;
-
-        let row = sqlx::query(query)
-          .bind(after_block as i64)
-          .fetch_one(&*pool)
-          .await?;
-
-        row.get::<Option<i64>, _>("block").map(|b| b as u64)
+      "construct_entity_ownership_changes"
+      | "construct_welcome_pack_ownership"
+      | "construct_welcome_pack_reward_destination" => {
+        let table = match job.query_name.as_str() {
+          "construct_entity_ownership_changes" => "asset_owners",
+          _ => "welcome_packs",
+        };
+        Self::next_block_for_table(&pool, table, after_block).await?
       }
       "construct_entity_reward_destination_changes" => {
         let query = r#"
@@ -1217,17 +1185,6 @@ impl DatabaseClient {
 
         row.get::<Option<i64>, _>("block").map(|b| b as u64)
       }
-      "construct_welcome_pack_ownership" | "construct_welcome_pack_reward_destination" => {
-        let query =
-          r#"SELECT MIN(last_block)::bigint as block FROM welcome_packs WHERE last_block > $1"#;
-
-        let row = sqlx::query(query)
-          .bind(after_block as i64)
-          .fetch_one(&*pool)
-          .await?;
-
-        row.get::<Option<i64>, _>("block").map(|b| b as u64)
-      }
       _ => {
         debug!(
           "Unknown query name for next block lookup: {}",

utils/atomic-data-publisher/tests/common/mod.rs

@@ -177,36 +177,37 @@ pub async fn seed_mini_fanout(
 
 // --- Query runners ---
 
-pub async fn run_ownership_query(
+async fn run_block_range_query(
   pool: &PgPool,
+  sql: &str,
   last_processed_block: i64,
   max_block: i64,
 ) -> Vec<Value> {
-  sqlx::query(OWNERSHIP_SQL)
+  sqlx::query(sql)
     .bind(last_processed_block)
     .bind(max_block)
     .fetch_all(pool)
     .await
-    .expect("run ownership query")
+    .expect("run query")
     .into_iter()
     .map(|row| row.get::<Value, _>("atomic_data"))
     .collect()
 }
 
+pub async fn run_ownership_query(
+  pool: &PgPool,
+  last_processed_block: i64,
+  max_block: i64,
+) -> Vec<Value> {
+  run_block_range_query(pool, OWNERSHIP_SQL, last_processed_block, max_block).await
+}
+
 pub async fn run_reward_query(
   pool: &PgPool,
   last_processed_block: i64,
   max_block: i64,
 ) -> Vec<Value> {
-  sqlx::query(REWARD_SQL)
-    .bind(last_processed_block)
-    .bind(max_block)
-    .fetch_all(pool)
-    .await
-    .expect("run reward query")
-    .into_iter()
-    .map(|row| row.get::<Value, _>("atomic_data"))
-    .collect()
+  run_block_range_query(pool, REWARD_SQL, last_processed_block, max_block).await
 }
 
 pub async fn run_hotspot_query(
@@ -232,29 +233,13 @@ pub async fn run_wp_ownership_query(
   last_processed_block: i64,
   max_block: i64,
 ) -> Vec<Value> {
-  sqlx::query(WP_OWNERSHIP_SQL)
-    .bind(last_processed_block)
-    .bind(max_block)
-    .fetch_all(pool)
-    .await
-    .expect("run wp ownership query")
-    .into_iter()
-    .map(|row| row.get::<Value, _>("atomic_data"))
-    .collect()
+  run_block_range_query(pool, WP_OWNERSHIP_SQL, last_processed_block, max_block).await
 }
 
 pub async fn run_wp_reward_query(
   pool: &PgPool,
   last_processed_block: i64,
   max_block: i64,
 ) -> Vec<Value> {
-  sqlx::query(WP_REWARD_SQL)
-    .bind(last_processed_block)
-    .bind(max_block)
-    .fetch_all(pool)
-    .await
-    .expect("run wp reward query")
-    .into_iter()
-    .map(|row| row.get::<Value, _>("atomic_data"))
-    .collect()
+  run_block_range_query(pool, WP_REWARD_SQL, last_processed_block, max_block).await
 }