Add balancer v2 indexing to pool-indexer

crates/pool-indexer/src/db/balancer_v2.rs

@@ -0,0 +1,116 @@
+use {
+    crate::{db::bytes_to_addr, indexer::balancer_v2::NewBalancerPool},
+    alloy_primitives::Address,
+    anyhow::{Context, Result},
+    bigdecimal::BigDecimal,
+    sqlx::{PgPool, Postgres, Row, Transaction},
+};
+
+/// Inserts discovered pools and their tokens. Pools are written before tokens
+/// to satisfy the FK; both `ON CONFLICT DO NOTHING` so re-indexing a pool is a
+/// no-op.
+pub async fn insert_pools(
+    tx: &mut Transaction<'_, Postgres>,
+    factory: &Address,
+    pools: &[NewBalancerPool],
+) -> Result<()> {
+    if pools.is_empty() {
+        return Ok(());
+    }
+
+    let mut pool_ids: Vec<&[u8]> = Vec::with_capacity(pools.len());
+    let mut addresses: Vec<&[u8]> = Vec::with_capacity(pools.len());
+    let mut pool_types: Vec<&str> = Vec::with_capacity(pools.len());
+    let mut created_blocks: Vec<i64> = Vec::with_capacity(pools.len());
+    let mut tok_pool_ids: Vec<&[u8]> = Vec::new();
+    let mut positions: Vec<i32> = Vec::new();
+    let mut tokens: Vec<&[u8]> = Vec::new();
+    let mut decimals: Vec<Option<i16>> = Vec::new();
+    let mut weights: Vec<Option<BigDecimal>> = Vec::new();
+    for pool in pools {
+        pool_ids.push(pool.pool_id.as_slice());
+        addresses.push(pool.address.as_slice());
+        pool_types.push(pool.pool_type.as_str());
+        created_blocks.push(pool.created_block.cast_signed());
+        for token in &pool.tokens {
+            tok_pool_ids.push(pool.pool_id.as_slice());
+            positions.push(i32::try_from(token.position).unwrap_or(i32::MAX));
+            tokens.push(token.address.as_slice());
+            decimals.push(token.decimals.map(i16::from));
+            weights.push(token.weight.clone());
+        }
+    }
+
+    sqlx::query(
+        "INSERT INTO balancer_v2_pools (pool_id, address, factory, pool_type, created_block)
+         SELECT t.pid, t.addr, $1, t.ptype, t.cblk
+         FROM UNNEST($2::BYTEA[], $3::BYTEA[], $4::TEXT[], $5::INT8[])
+              AS t(pid, addr, ptype, cblk)
+         ON CONFLICT (pool_id) DO NOTHING",
+    )
+    .bind(factory.as_slice())
+    .bind(pool_ids)
+    .bind(addresses)
+    .bind(pool_types)
+    .bind(created_blocks)
+    .execute(&mut **tx)
+    .await
+    .context("insert balancer pools")?;
+
+    sqlx::query(
+        "INSERT INTO balancer_v2_pool_tokens (pool_id, position, token, decimals, weight)
+         SELECT t.pid, t.pos, t.tok, t.dec, t.wgt
+         FROM UNNEST($1::BYTEA[], $2::INT4[], $3::BYTEA[], $4::INT2[], $5::NUMERIC[])
+              AS t(pid, pos, tok, dec, wgt)
+         ON CONFLICT (pool_id, position) DO NOTHING",
+    )
+    .bind(tok_pool_ids)
+    .bind(positions)
+    .bind(tokens)
+    .bind(decimals)
+    .bind(weights)
+    .execute(&mut **tx)
+    .await
+    .context("insert balancer pool tokens")?;
+
+    Ok(())
+}
+
+/// Distinct token addresses with no `decimals` recorded yet.
+pub async fn get_tokens_missing_decimals(pool: &PgPool) -> Result<Vec<Address>> {
+    let rows =
+        sqlx::query("SELECT DISTINCT token FROM balancer_v2_pool_tokens WHERE decimals IS NULL")
+            .fetch_all(pool)
+            .await
+            .context("get_tokens_missing_decimals")?;
+
+    rows.into_iter()
+        .map(|r| bytes_to_addr(r.get("token")))
+        .collect()
+}
+
+/// Sets `decimals` for every token row matching one of the inputs. Pass `-1`
+/// for "tried, failed" so the next backfill's `IS NULL` filter still skips it.
+pub async fn batch_set_token_decimals(
+    tx: &mut Transaction<'_, Postgres>,
+    entries: &[(Address, i16)],
+) -> Result<()> {
+    if entries.is_empty() {
+        return Ok(());
+    }
+    let tokens: Vec<&[u8]> = entries.iter().map(|(t, _)| t.as_slice()).collect();
+    let decimals: Vec<i16> = entries.iter().map(|(_, d)| *d).collect();
+
+    sqlx::query(
+        "UPDATE balancer_v2_pool_tokens p
+         SET decimals = i.dec
+         FROM UNNEST($1::BYTEA[], $2::INT2[]) AS i(tok, dec)
+         WHERE p.token = i.tok AND p.decimals IS NULL",
+    )
+    .bind(tokens)
+    .bind(decimals)
+    .execute(&mut **tx)
+    .await
+    .context("batch_set_token_decimals")?;
+    Ok(())
+}

crates/pool-indexer/src/db/mod.rs

@@ -1 +1,46 @@
+pub mod balancer_v2;
 pub mod uniswap_v3;
+
+use {
+    alloy_primitives::Address,
+    anyhow::{Context, Result},
+    sqlx::{PgPool, Postgres, Row, Transaction},
+};
+
+/// Decodes a Postgres `BYTEA` column into an [`Address`].
+pub(crate) fn bytes_to_addr(b: Vec<u8>) -> Result<Address> {
+    Address::try_from(b.as_slice()).context("invalid address bytes")
+}
+
+/// Highest block scanned for a factory's `PoolCreated` events. Shared by both
+/// indexers: `pool_indexer_checkpoints` is keyed by factory address, which is
+/// unique across protocols, so their rows never collide.
+pub async fn get_checkpoint(pool: &PgPool, factory: &Address) -> Result<Option<u64>> {
+    let row = sqlx::query(
+        "SELECT block_number FROM pool_indexer_checkpoints WHERE contract_address = $1",
+    )
+    .bind(factory.as_slice())
+    .fetch_optional(pool)
+    .await
+    .context("get_checkpoint")?;
+
+    Ok(row.map(|r| r.get::<i64, _>("block_number").cast_unsigned()))
+}
+
+pub async fn set_checkpoint(
+    tx: &mut Transaction<'_, Postgres>,
+    factory: &Address,
+    block_number: u64,
+) -> Result<()> {
+    sqlx::query(
+        "INSERT INTO pool_indexer_checkpoints (contract_address, block_number)
+         VALUES ($1, $2)
+         ON CONFLICT (contract_address) DO UPDATE SET block_number = EXCLUDED.block_number",
+    )
+    .bind(factory.as_slice())
+    .bind(block_number.cast_signed())
+    .execute(&mut **tx)
+    .await
+    .context("set_checkpoint")?;
+    Ok(())
+}

crates/pool-indexer/src/db/uniswap_v3.rs

@@ -1,5 +1,8 @@
 use {
-    crate::indexer::uniswap_v3::{LiquidityUpdateData, NewPoolData, PoolStateData, TickDeltaData},
+    crate::{
+        db::bytes_to_addr,
+        indexer::uniswap_v3::{LiquidityUpdateData, NewPoolData, PoolStateData, TickDeltaData},
+    },
     alloy_primitives::Address,
     anyhow::{Context, Result},
     bigdecimal::BigDecimal,
@@ -9,10 +12,6 @@ use {
     std::collections::BTreeSet,
 };
 
-fn bytes_to_addr(b: Vec<u8>) -> Result<Address> {
-    Address::try_from(b.as_slice()).context("invalid address bytes")
-}
-
 fn address_bytes_list(addresses: &[Address]) -> Vec<&[u8]> {
     addresses.iter().map(|address| address.as_slice()).collect()
 }
@@ -21,36 +20,6 @@ fn decode_pool_rows(rows: Vec<PgRow>) -> Result<Vec<PoolRow>> {
     rows.into_iter().map(PoolRow::try_from).collect()
 }
 
-pub async fn get_checkpoint(pool: &PgPool, contract: &Address) -> Result<Option<u64>> {
-    let row = sqlx::query(
-        "SELECT block_number FROM pool_indexer_checkpoints WHERE contract_address = $1",
-    )
-    .bind(contract.as_slice())
-    .fetch_optional(pool)
-    .await
-    .context("get_checkpoint")?;
-
-    Ok(row.map(|r| r.get::<i64, _>("block_number").cast_unsigned()))
-}
-
-pub async fn set_checkpoint(
-    tx: &mut Transaction<'_, Postgres>,
-    contract: &Address,
-    block_number: u64,
-) -> Result<()> {
-    sqlx::query(
-        "INSERT INTO pool_indexer_checkpoints (contract_address, block_number)
-         VALUES ($1, $2)
-         ON CONFLICT (contract_address) DO UPDATE SET block_number = EXCLUDED.block_number",
-    )
-    .bind(contract.as_slice())
-    .bind(block_number.cast_signed())
-    .execute(&mut **tx)
-    .await
-    .context("set_checkpoint")?;
-    Ok(())
-}
-
 pub async fn insert_pools(
     tx: &mut Transaction<'_, Postgres>,
     factory: &Address,