feat(multi): uniformly distribute addresses for quick search (#1290)

* feat: uniformly distribute addresses for quick search

* refactor: rename parameters

* fix: incorrect case for dapp client query parameter

* feat: add deterministic ordering for address quick search

* feat: apply priority ordering for all quick search entities

Author: lok52

multichain-aggregator/Cargo.lock

@@ -3,6 +3,7 @@ use api_client_framework::{
 };
 use reqwest::Method;
 use serde::{Deserialize, Serialize};
+use serde_with::{formats::CommaSeparator, serde_as, StringWithSeparator};
 use url::Url;
 
 pub fn new_client(url: Url) -> Result<Client, Error> {
@@ -17,11 +18,14 @@ pub mod search_dapps {
         pub params: SearchDappsParams,
     }
 
+    #[serde_as]
     #[derive(Serialize, Clone, Debug, Default, PartialEq)]
+    #[serde(rename_all = "camelCase")]
     pub struct SearchDappsParams {
         pub title: Option<String>,
         pub categories: Option<String>,
-        pub chain_ids: Option<String>,
+        #[serde_as(as = "StringWithSeparator::<CommaSeparator, i64>")]
+        pub chain_ids: Vec<i64>,
     }
 
     impl Endpoint for SearchDapps {

multichain-aggregator/multichain-aggregator-logic/src/clients/dapp.rs

@@ -4,6 +4,7 @@ use api_client_framework::{
 };
 use reqwest::Method;
 use serde::{Deserialize, Serialize};
+use serde_with::{formats::CommaSeparator, serde_as, StringWithSeparator};
 use url::Url;
 
 pub fn new_client(url: Url) -> Result<Client, Error> {
@@ -41,11 +42,14 @@ pub mod search_token_infos {
         pub params: SearchTokenInfosParams,
     }
 
+    #[serde_as]
     #[serde_with::skip_serializing_none]
     #[derive(Serialize, Clone, Debug, Default, PartialEq)]
     pub struct SearchTokenInfosParams {
         pub query: String,
-        pub chain_id: Option<ChainId>,
+        #[serde_as(as = "StringWithSeparator::<CommaSeparator, ChainId>")]
+        #[serde(skip_serializing_if = "Vec::is_empty")]
+        pub chain_id: Vec<ChainId>,
         pub page_size: Option<u32>,
         pub page_token: Option<String>,
     }

multichain-aggregator/multichain-aggregator-logic/src/clients/token_info.rs

@@ -7,7 +7,10 @@ use entity::{
 use regex::Regex;
 use sea_orm::{
     prelude::Expr,
-    sea_query::{Alias, ColumnRef, CommonTableExpression, IntoIden, OnConflict, Query, WithClause},
+    sea_query::{
+        Alias, ColumnRef, CommonTableExpression, IntoIden, OnConflict, Query, WindowStatement,
+        WithClause,
+    },
     ActiveValue::NotSet,
     ColumnTrait, ConnectionTrait, DbErr, EntityTrait, FromQueryResult, IntoSimpleExpr, Iterable,
     Order, QuerySelect,
@@ -46,50 +49,147 @@ where
     Ok(())
 }
 
-pub async fn list<C>(
-    db: &C,
-    address: Option<AddressAlloy>,
-    query: Option<String>,
-    chain_id: Option<ChainId>,
-    token_types: Option<Vec<db_enum::TokenType>>,
-    page_size: u64,
-    page_token: Option<(AddressAlloy, ChainId)>,
-) -> Result<(Vec<Model>, Option<(AddressAlloy, ChainId)>), DbErr>
-where
-    C: ConnectionTrait,
-{
+fn prepare_address_search_cte(
+    contract_name_query: Option<String>,
+    cte_name: impl IntoIden,
+) -> CommonTableExpression {
     // Materialize addresses CTE when searching by contract_name.
     // Otherwise, query planner chooses a suboptimal plan.
-    let is_cte_materialized = query.is_some();
-    let addresses_cte_iden = Alias::new("addresses").into_iden();
-    let addresses_cte = CommonTableExpression::new()
+    // If query is not provided, this CTE will be folded by the optimizer.
+    let is_cte_materialized = contract_name_query.is_some();
+
+    CommonTableExpression::new()
         .query(
             QuerySelect::query(&mut Entity::find())
-                .apply_if(query, |q, query| {
+                .apply_if(contract_name_query, |q, query| {
                     let ts_query = prepare_ts_query(&query);
                     q.and_where(Expr::cust_with_expr(
                         "to_tsvector('english', contract_name) @@ to_tsquery($1)",
                         ts_query,
                     ));
                 })
+                // Apply a hard limit in case we materialize the CTE
+                .apply_if(is_cte_materialized.then_some(10_000), |q, limit| {
+                    q.limit(limit);
+                })
                 .to_owned(),
         )
         .materialized(is_cte_materialized)
-        .table_name(addresses_cte_iden.clone())
+        .table_name(cte_name)
+        .to_owned()
+}
+
+pub async fn uniform_chain_search<C>(
+    db: &C,
+    contract_name_query: String,
+    token_types: Option<Vec<db_enum::TokenType>>,
+    chain_ids: Vec<ChainId>,
+) -> Result<Vec<Model>, DbErr>
+where
+    C: ConnectionTrait,
+{
+    if chain_ids.is_empty() {
+        return Ok(vec![]);
+    }
+
+    let ts_rank_ordering = Expr::cust_with_expr(
+        "ts_rank(to_tsvector('english', contract_name), to_tsquery($1))",
+        prepare_ts_query(&contract_name_query),
+    );
+
+    let addresses_cte_iden = Alias::new("addresses").into_iden();
+    let addresses_cte =
+        prepare_address_search_cte(Some(contract_name_query), addresses_cte_iden.clone());
+
+    let row_number = Expr::custom_keyword(Alias::new("ROW_NUMBER()"));
+    let ranked_addresses_iden = Alias::new("ranked_addresses").into_iden();
+    let ranked_addresses_cte = CommonTableExpression::new()
+        .query(
+            Query::select()
+                .column(ColumnRef::TableAsterisk(addresses_cte_iden.clone()))
+                .expr_window_as(
+                    row_number,
+                    WindowStatement::partition_by(Column::ChainId)
+                        .order_by_expr(ts_rank_ordering, Order::Desc)
+                        .order_by(Column::Hash, Order::Asc)
+                        .to_owned(),
+                    Alias::new("rn"),
+                )
+                .from(addresses_cte_iden.clone())
+                .and_where(Column::ChainId.is_in(chain_ids.clone()))
+                .apply_if(token_types, |q, token_types| {
+                    if !token_types.is_empty() {
+                        q.and_where(Column::TokenType.is_in(token_types));
+                    } else {
+                        q.and_where(Column::TokenType.is_null());
+                    }
+                })
+                .to_owned(),
+        )
+        .table_name(ranked_addresses_iden.clone())
+        .to_owned();
+
+    let limit = chain_ids.len() as u64;
+    let base_select = Query::select()
+        .column(ColumnRef::Asterisk)
+        .from(ranked_addresses_iden)
+        .and_where(Expr::col(Alias::new("rn")).eq(1))
+        .order_by_expr(
+            Expr::cust_with_exprs(
+                "array_position($1, $2)",
+                [chain_ids.into(), Expr::col(Column::ChainId).into()],
+            ),
+            Order::Asc,
+        )
+        .limit(limit)
         .to_owned();
 
+    let query = WithClause::new()
+        .cte(addresses_cte)
+        .cte(ranked_addresses_cte)
+        .to_owned()
+        .query(base_select);
+
+    let addresses = Model::find_by_statement(db.get_database_backend().build(&query))
+        .all(db)
+        .await?;
+
+    Ok(addresses)
+}
+
+pub async fn list<C>(
+    db: &C,
+    address: Option<AddressAlloy>,
+    contract_name_query: Option<String>,
+    chain_ids: Option<Vec<ChainId>>,
+    token_types: Option<Vec<db_enum::TokenType>>,
+    page_size: u64,
+    page_token: Option<(AddressAlloy, ChainId)>,
+) -> Result<(Vec<Model>, Option<(AddressAlloy, ChainId)>), DbErr>
+where
+    C: ConnectionTrait,
+{
+    let addresses_cte_iden = Alias::new("addresses").into_iden();
+    let addresses_cte = prepare_address_search_cte(contract_name_query, addresses_cte_iden.clone());
+
     let base_select = Query::select()
         .column(ColumnRef::Asterisk)
         .from(addresses_cte_iden)
-        .apply_if(chain_id, |q, chain_id| {
-            q.and_where(Column::ChainId.eq(chain_id));
+        .apply_if(chain_ids, |q, chain_ids| {
+            if !chain_ids.is_empty() {
+                q.and_where(Column::ChainId.is_in(chain_ids));
+            }
+        })
+        .apply_if(token_types, |q, token_types| {
+            if !token_types.is_empty() {
+                q.and_where(Column::TokenType.is_in(token_types));
+            } else {
+                q.and_where(Column::TokenType.is_null());
+            }
         })
         .apply_if(address, |q, address| {
             q.and_where(Column::Hash.eq(address.as_slice()));
         })
-        .apply_if(token_types, |q, token_types| {
-            q.and_where(Column::TokenType.is_in(token_types));
-        })
         .apply_if(page_token, |q, page_token| {
             q.and_where(
                 Expr::tuple([

multichain-aggregator/multichain-aggregator-logic/src/repository/addresses.rs

@@ -3,7 +3,7 @@ use crate::types::{block_ranges::BlockRange, ChainId};
 use entity::block_ranges::{ActiveModel, Column, Entity, Model};
 use sea_orm::{
     prelude::Expr, sea_query::OnConflict, ActiveValue::NotSet, ColumnTrait, ConnectionTrait, DbErr,
-    EntityTrait, QueryFilter,
+    EntityTrait, QueryFilter, QueryTrait,
 };
 
 pub async fn upsert_many<C>(db: &C, block_ranges: Vec<BlockRange>) -> Result<(), DbErr>
@@ -56,13 +56,21 @@ where
 pub async fn list_matching_block_ranges_paginated<C>(
     db: &C,
     block_number: u64,
+    chain_ids: Option<Vec<ChainId>>,
     page_size: u64,
     page_token: Option<ChainId>,
 ) -> Result<(Vec<Model>, Option<ChainId>), DbErr>
 where
     C: ConnectionTrait,
 {
     let mut c = Entity::find()
+        .apply_if(chain_ids, |q, chain_ids| {
+            if !chain_ids.is_empty() {
+                q.filter(Column::ChainId.is_in(chain_ids))
+            } else {
+                q
+            }
+        })
         .filter(Column::MinBlockNumber.lte(block_number))
         .filter(Column::MaxBlockNumber.gte(block_number))
         .cursor_by(Column::ChainId);

multichain-aggregator/multichain-aggregator-logic/src/repository/block_ranges.rs

@@ -40,7 +40,7 @@ pub async fn list<C>(
     db: &C,
     hash: BlockHash,
     hash_type: Option<db_enum::HashType>,
-    chain_id: Option<ChainId>,
+    chain_ids: Option<Vec<ChainId>>,
     page_size: u64,
     page_token: Option<ChainId>,
 ) -> Result<(Vec<Model>, Option<ChainId>), DbErr>
@@ -52,8 +52,12 @@ where
         .apply_if(hash_type, |q, hash_type| {
             q.filter(Column::HashType.eq(hash_type))
         })
-        .apply_if(chain_id, |q, chain_id| {
-            q.filter(Column::ChainId.eq(chain_id))
+        .apply_if(chain_ids, |q, chain_ids| {
+            if !chain_ids.is_empty() {
+                q.filter(Column::ChainId.is_in(chain_ids))
+            } else {
+                q
+            }
         })
         .cursor_by(Column::ChainId);