fix/db: improve DB query complexity

CHANGELOG.md

@@ -20,6 +20,7 @@
 - Added `GetLimits` endpoint to the RPC server ([#1410](https://github.com/0xMiden/miden-node/pull/1410)).
 - Added gRPC-Web probe support to the `miden-network-monitor` binary ([#1484](https://github.com/0xMiden/miden-node/pull/1484)).
 - Add DB schema change check ([#1268](https://github.com/0xMiden/miden-node/pull/1485)).
+- Improve DB query performance for account queries ([#1496](https://github.com/0xMiden/miden-node/pull/1496).
 
 ### Changes
 

crates/store/src/db/models/queries/accounts/at_block.rs

@@ -2,14 +2,7 @@ use std::collections::BTreeMap;
 
 use diesel::prelude::Queryable;
 use diesel::query_dsl::methods::SelectDsl;
-use diesel::{
-    BoolExpressionMethods,
-    ExpressionMethods,
-    OptionalExtension,
-    QueryDsl,
-    RunQueryDsl,
-    SqliteConnection,
-};
+use diesel::{ExpressionMethods, OptionalExtension, QueryDsl, RunQueryDsl, SqliteConnection};
 use miden_protocol::account::{
     AccountHeader,
     AccountId,
@@ -125,6 +118,14 @@ pub(crate) fn select_account_header_at_block(
 // ================================================================================================
 
 /// Query vault assets at a specific block by finding the most recent update for each `vault_key`.
+///
+/// Fetches all vault entries for this account up to and including the target block, ordered by
+/// `vault_key` ascending and `block_num` descending. This allows efficient in-memory deduplication
+/// by taking only the first entry for each `vault_key`.
+///
+/// Note: Diesel's query DSL doesn't support joining with derived tables (subqueries with GROUP BY),
+/// so we use the ordering-based approach which is efficient due to the database index on
+/// `(account_id, vault_key, block_num)`.
 pub(crate) fn select_account_vault_at_block(
     conn: &mut SqliteConnection,
     account_id: AccountId,
@@ -135,50 +136,36 @@ pub(crate) fn select_account_vault_at_block(
     let account_id_bytes = account_id.to_bytes();
     let block_num_sql = block_num.to_raw_sql();
 
-    // Since Diesel doesn't support composite keys in subqueries easily, we use a two-step approach:
-    // Step 1: Get max block_num for each vault_key
-    let latest_blocks_per_vault_key = Vec::from_iter(
-        QueryDsl::select(
-            t::table
-                .filter(t::account_id.eq(&account_id_bytes))
-                .filter(t::block_num.le(block_num_sql))
-                .group_by(t::vault_key),
-            (t::vault_key, diesel::dsl::max(t::block_num)),
-        )
-        .load::<(Vec<u8>, Option<i64>)>(conn)?
-        .into_iter()
-        .filter_map(|(key, maybe_block)| maybe_block.map(|block| (key, block))),
-    );
-
-    if latest_blocks_per_vault_key.is_empty() {
-        return Ok(Vec::new());
-    }
+    // Query all vault entries for this account up to and including this block.
+    // Order by vault_key ascending, then block_num descending so the first entry
+    // for each vault_key is the latest one.
+    let all_entries: Vec<(Vec<u8>, Option<Vec<u8>>)> = SelectDsl::select(
+        t::table
+            .filter(t::account_id.eq(&account_id_bytes))
+            .filter(t::block_num.le(block_num_sql))
+            .order((t::vault_key.asc(), t::block_num.desc())),
+        (t::vault_key, t::asset),
+    )
+    .load(conn)?;
 
-    // Step 2: Fetch the full rows matching (vault_key, block_num) pairs
+    // For each vault_key, keep only the latest entry (first one due to ordering)
     let mut assets = Vec::new();
-    for (vault_key_bytes, max_block) in latest_blocks_per_vault_key {
-        // TODO we should not make a query per vault key, but query many at once or
-        // or find an alternative approach
-        let result: Option<Option<Vec<u8>>> = QueryDsl::select(
-            t::table.filter(
-                t::account_id
-                    .eq(&account_id_bytes)
-                    .and(t::vault_key.eq(&vault_key_bytes))
-                    .and(t::block_num.eq(max_block)),
-            ),
-            t::asset,
-        )
-        .first(conn)
-        .optional()?;
-        if let Some(Some(asset_bytes)) = result {
-            let asset = Asset::read_from_bytes(&asset_bytes)?;
+    let mut last_vault_key: Option<&[u8]> = None;
+
+    for (vault_key_bytes, asset_bytes) in &all_entries {
+        // Skip if we've already processed this vault_key
+        if last_vault_key == Some(vault_key_bytes.as_slice()) {
+            continue;
+        }
+        last_vault_key = Some(vault_key_bytes.as_slice());
+
+        // Only include if the asset exists (not deleted)
+        if let Some(asset_bytes) = asset_bytes {
+            let asset = Asset::read_from_bytes(asset_bytes)?;
             assets.push(asset);
         }
     }
 
-    // Sort by vault_key for consistent ordering
-    assets.sort_by_key(Asset::vault_key);
-
     Ok(assets)
 }
 
@@ -218,18 +205,19 @@ pub(crate) fn select_account_storage_at_block(
     let header = AccountStorageHeader::read_from_bytes(&blob)?;
 
     // Query all map values for this account up to and including this block.
-    // For each (slot_name, key), we need the latest value at or before block_num.
-    // First, get all entries up to block_num
-    let map_values: Vec<(i64, String, Vec<u8>, Vec<u8>)> =
-        SelectDsl::select(t::table, (t::block_num, t::slot_name, t::key, t::value))
-            .filter(t::account_id.eq(&account_id_bytes).and(t::block_num.le(block_num_sql)))
+    // Order by (slot_name, key) ascending, then block_num descending so the first entry
+    // for each (slot_name, key) pair is the latest one.
+    let map_values: Vec<(String, Vec<u8>, Vec<u8>)> =
+        SelectDsl::select(t::table, (t::slot_name, t::key, t::value))
+            .filter(t::account_id.eq(&account_id_bytes))
+            .filter(t::block_num.le(block_num_sql))
             .order((t::slot_name.asc(), t::key.asc(), t::block_num.desc()))
             .load(conn)?;
 
-    // For each (slot_name, key) pair, keep only the latest entry (highest block_num)
+    // For each (slot_name, key) pair, keep only the latest entry (first one due to ordering)
     let mut latest_map_entries: BTreeMap<(StorageSlotName, Word), Word> = BTreeMap::new();
 
-    for (_, slot_name_str, key_bytes, value_bytes) in map_values {
+    for (slot_name_str, key_bytes, value_bytes) in map_values {
         let slot_name: StorageSlotName = slot_name_str.parse().map_err(|_| {
             DatabaseError::DataCorrupted(format!("Invalid slot name: {slot_name_str}"))
         })?;

crates/store/src/db/models/queries/accounts/tests.rs

@@ -550,3 +550,180 @@ fn test_upsert_accounts_with_empty_storage() {
         "Storage header blob should exist even for empty storage"
     );
 }
+
+// VAULT AT BLOCK HISTORICAL QUERY TESTS
+// ================================================================================================
+
+/// Tests that querying vault at an older block returns the correct historical state,
+/// even when the same `vault_key` has been updated in later blocks.
+///
+/// This tests the deduplication logic that relies on ordering by (`vault_key` ASC, `block_num` DESC).
+#[test]
+fn test_select_account_vault_at_block_historical_with_updates() {
+    use assert_matches::assert_matches;
+    use miden_protocol::asset::{AssetVaultKey, FungibleAsset};
+    use miden_protocol::testing::account_id::ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET;
+
+    let mut conn = setup_test_db();
+    let (account, _) = create_test_account_with_storage();
+    let account_id = account.id();
+
+    // Faucet ID is needed for creating FungibleAssets
+    let faucet_id = AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).unwrap();
+
+    let block_1 = BlockNumber::from_epoch(0);
+    let block_2 = BlockNumber::from_epoch(1);
+    let block_3 = BlockNumber::from_epoch(2);
+
+    insert_block_header(&mut conn, block_1);
+    insert_block_header(&mut conn, block_2);
+    insert_block_header(&mut conn, block_3);
+
+    // Insert account at block 1
+    let delta = AccountDelta::try_from(account.clone()).unwrap();
+    let account_update = BlockAccountUpdate::new(
+        account_id,
+        account.commitment(),
+        AccountUpdateDetails::Delta(delta),
+    );
+    upsert_accounts(&mut conn, &[account_update], block_1).expect("upsert_accounts failed");
+
+    // Insert vault asset at block 1: vault_key_1 = 1000 tokens
+    let vault_key_1 = AssetVaultKey::new_unchecked(Word::from([
+        Felt::new(1),
+        Felt::new(0),
+        Felt::new(0),
+        Felt::new(0),
+    ]));
+    let asset_v1 = Asset::Fungible(FungibleAsset::new(faucet_id, 1000).unwrap());
+
+    insert_account_vault_asset(&mut conn, account_id, block_1, vault_key_1, Some(asset_v1))
+        .expect("insert vault asset failed");
+
+    // Update vault asset at block 2: vault_key_1 = 2000 tokens (updated value)
+    let asset_v2 = Asset::Fungible(FungibleAsset::new(faucet_id, 2000).unwrap());
+    insert_account_vault_asset(&mut conn, account_id, block_2, vault_key_1, Some(asset_v2))
+        .expect("insert vault asset update failed");
+
+    // Add a second vault_key at block 2
+    let vault_key_2 = AssetVaultKey::new_unchecked(Word::from([
+        Felt::new(2),
+        Felt::new(0),
+        Felt::new(0),
+        Felt::new(0),
+    ]));
+    let asset_key2 = Asset::Fungible(FungibleAsset::new(faucet_id, 500).unwrap());
+    insert_account_vault_asset(&mut conn, account_id, block_2, vault_key_2, Some(asset_key2))
+        .expect("insert second vault asset failed");
+
+    // Update vault_key_1 again at block 3: vault_key_1 = 3000 tokens
+    let asset_v3 = Asset::Fungible(FungibleAsset::new(faucet_id, 3000).unwrap());
+    insert_account_vault_asset(&mut conn, account_id, block_3, vault_key_1, Some(asset_v3))
+        .expect("insert vault asset update 2 failed");
+
+    // Query at block 1: should only see vault_key_1 with 1000 tokens
+    let assets_at_block_1 = select_account_vault_at_block(&mut conn, account_id, block_1)
+        .expect("Query at block 1 should succeed");
+
+    assert_eq!(assets_at_block_1.len(), 1, "Should have 1 asset at block 1");
+    assert_matches!(&assets_at_block_1[0], Asset::Fungible(f) if f.amount() == 1000);
+
+    // Query at block 2: should see vault_key_1 with 2000 tokens AND vault_key_2 with 500 tokens
+    let assets_at_block_2 = select_account_vault_at_block(&mut conn, account_id, block_2)
+        .expect("Query at block 2 should succeed");
+
+    assert_eq!(assets_at_block_2.len(), 2, "Should have 2 assets at block 2");
+
+    // Find the amounts (order may vary)
+    let amounts: Vec<u64> = assets_at_block_2
+        .iter()
+        .map(|a| assert_matches!(a, Asset::Fungible(f) => f.amount()))
+        .collect();
+
+    assert!(amounts.contains(&2000), "Block 2 should have vault_key_1 with 2000 tokens");
+    assert!(amounts.contains(&500), "Block 2 should have vault_key_2 with 500 tokens");
+
+    // Query at block 3: should see vault_key_1 with 3000 tokens AND vault_key_2 with 500 tokens
+    let assets_at_block_3 = select_account_vault_at_block(&mut conn, account_id, block_3)
+        .expect("Query at block 3 should succeed");
+
+    assert_eq!(assets_at_block_3.len(), 2, "Should have 2 assets at block 3");
+
+    let amounts: Vec<u64> = assets_at_block_3
+        .iter()
+        .map(|a| assert_matches!(a, Asset::Fungible(f) => f.amount()))
+        .collect();
+
+    assert!(amounts.contains(&3000), "Block 3 should have vault_key_1 with 3000 tokens");
+    assert!(amounts.contains(&500), "Block 3 should have vault_key_2 with 500 tokens");
+}
+
+/// Tests that deleted vault assets (asset = None) are correctly excluded from results,
+/// and that the deduplication handles deletion entries properly.
+#[test]
+fn test_select_account_vault_at_block_with_deletion() {
+    use assert_matches::assert_matches;
+    use miden_protocol::asset::{AssetVaultKey, FungibleAsset};
+    use miden_protocol::testing::account_id::ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET;
+
+    let mut conn = setup_test_db();
+    let (account, _) = create_test_account_with_storage();
+    let account_id = account.id();
+
+    // Faucet ID is needed for creating FungibleAssets
+    let faucet_id = AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).unwrap();
+
+    let block_1 = BlockNumber::from_epoch(0);
+    let block_2 = BlockNumber::from_epoch(1);
+    let block_3 = BlockNumber::from_epoch(2);
+
+    insert_block_header(&mut conn, block_1);
+    insert_block_header(&mut conn, block_2);
+    insert_block_header(&mut conn, block_3);
+
+    // Insert account at block 1
+    let delta = AccountDelta::try_from(account.clone()).unwrap();
+    let account_update = BlockAccountUpdate::new(
+        account_id,
+        account.commitment(),
+        AccountUpdateDetails::Delta(delta),
+    );
+    upsert_accounts(&mut conn, &[account_update], block_1).expect("upsert_accounts failed");
+
+    // Insert vault asset at block 1
+    let vault_key = AssetVaultKey::new_unchecked(Word::from([
+        Felt::new(1),
+        Felt::new(0),
+        Felt::new(0),
+        Felt::new(0),
+    ]));
+    let asset = Asset::Fungible(FungibleAsset::new(faucet_id, 1000).unwrap());
+
+    insert_account_vault_asset(&mut conn, account_id, block_1, vault_key, Some(asset))
+        .expect("insert vault asset failed");
+
+    // Delete the vault asset at block 2 (insert with asset = None)
+    insert_account_vault_asset(&mut conn, account_id, block_2, vault_key, None)
+        .expect("delete vault asset failed");
+
+    // Re-add the vault asset at block 3 with different amount
+    let asset_v3 = Asset::Fungible(FungibleAsset::new(faucet_id, 2000).unwrap());
+    insert_account_vault_asset(&mut conn, account_id, block_3, vault_key, Some(asset_v3))
+        .expect("re-add vault asset failed");
+
+    // Query at block 1: should see the asset
+    let assets_at_block_1 = select_account_vault_at_block(&mut conn, account_id, block_1)
+        .expect("Query at block 1 should succeed");
+    assert_eq!(assets_at_block_1.len(), 1, "Should have 1 asset at block 1");
+
+    // Query at block 2: should NOT see the asset (it was deleted)
+    let assets_at_block_2 = select_account_vault_at_block(&mut conn, account_id, block_2)
+        .expect("Query at block 2 should succeed");
+    assert!(assets_at_block_2.is_empty(), "Should have no assets at block 2 (deleted)");
+
+    // Query at block 3: should see the re-added asset with new amount
+    let assets_at_block_3 = select_account_vault_at_block(&mut conn, account_id, block_3)
+        .expect("Query at block 3 should succeed");
+    assert_eq!(assets_at_block_3.len(), 1, "Should have 1 asset at block 3");
+    assert_matches!(&assets_at_block_3[0], Asset::Fungible(f) if f.amount() == 2000);
+}