crates/torii/core/src/model.rs (1)

414-453: Ohayo sensei! Consider reducing duplication of collect_columns logic.
This block replicates the existing approach from build_sql_query. You might refactor it into a shared helper function if feasible.

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

#!/bin/bash
# We'll examine usage of fetch_entities across the codebase to verify we handle chunk results properly
rg -A 5 'fetch_entities'

#!/bin/bash
# Let's examine the full context of fetch_entities implementation
ast-grep --pattern 'pub async fn fetch_entities($$$)'

# Also check if there are any chunk-related variables or logic
rg -B 3 -A 3 'chunk' crates/torii/core/src/model.rs

crates/torii/core/src/model.rs (1)

509-543: Consider adding batch size configuration for better performance tuning.

While the implementation is correct, consider making the chunk size configurable through a configuration parameter rather than hardcoding it. This would allow for performance tuning based on:

• Available system resources
• Specific database workload patterns
• Query complexity

-const MAX_JOINS: usize = 64;
+/// Maximum number of joins per chunk, defaulting to SQLite's limit of 64
+const DEFAULT_MAX_JOINS: usize = 64;
+
+pub struct FetchOptions {
+    /// Maximum number of joins per chunk
+    pub max_joins: usize,
+}
+
+impl Default for FetchOptions {
+    fn default() -> Self {
+        Self {
+            max_joins: DEFAULT_MAX_JOINS,
+        }
+    }
+}

crates/torii/grpc/src/server/mod.rs (5)

Line range hint 315-342: Improve WHERE clause construction readability.

Consider extracting the WHERE clause construction into a separate function for better readability and maintainability.

+fn build_where_clause(table: &str, ids: &[String], entity_updated_after: Option<&str>) -> String {
+    let ids_clause = ids.join(" OR ");
+    match entity_updated_after {
+        Some(_) => format!("{} AND {}.updated_at >= ?", ids_clause, table),
+        None => ids_clause,
+    }
+}
+
 let where_clause = match &hashed_keys {
     Some(hashed_keys) => {
-        let ids =
-            hashed_keys.hashed_keys.iter().map(|_| "{table}.id = ?").collect::<Vec<_>>();
-        format!(
-            "{} {}",
-            ids.join(" OR "),
-            if entity_updated_after.is_some() {
-                format!("AND {table}.updated_at >= ?")
-            } else {
-                String::new()
-            }
-        )
+        let ids = hashed_keys.hashed_keys.iter().map(|_| format!("{}.id = ?", table)).collect::<Vec<_>>();
+        build_where_clause(table, &ids, entity_updated_after.as_deref())
     }
     None => {
-        if entity_updated_after.is_some() {
-            format!("{table}.updated_at >= ?")
-        } else {
-            String::new()
-        }
+        entity_updated_after.map_or(String::new(), |_| format!("{}.updated_at >= ?", table))
     }
 };

---

365-396: Consider extracting common SQL query patterns.

The historical event messages query construction could be refactored to reuse common SQL patterns.

+const HISTORICAL_EVENTS_BASE_QUERY: &str = r#"
+    SELECT {columns}
+    FROM {table}
+    JOIN {model_relation_table} ON {table}.id = {model_relation_table}.entity_id
+    WHERE {where_clause}
+    GROUP BY {table}.event_id
+    ORDER BY {table}.event_id DESC
+"#;
+
-let count_query = format!(
-    r#"
-    SELECT COUNT(*) FROM {table}
-    JOIN {model_relation_table} ON {table}.id = {model_relation_table}.entity_id
-    WHERE {where_clause}
-    GROUP BY {table}.event_id
-"#
-);
+let count_query = HISTORICAL_EVENTS_BASE_QUERY.replace("{columns}", "COUNT(*)");

---

400-414: Enhance error handling for fetch_entities call.

While the parallel processing is efficient, consider adding more specific error handling for different failure scenarios.

-let (rows, total_count) = fetch_entities(
+let fetch_result = fetch_entities(
     &self.pool,
     &schemas,
     table,
     model_relation_table,
     entity_relation_column,
     if !where_clause.is_empty() { Some(&where_clause) } else { None },
     if !having_clause.is_empty() { Some(&having_clause) } else { None },
     order_by,
     limit,
     offset,
     bind_values,
 )
-.await?;
+.await
+.map_err(|e| match e {
+    Error::Sql(sqlx::Error::RowNotFound) => Error::from(QueryError::NoResults),
+    Error::Sql(sqlx::Error::Database(e)) if e.code().as_deref() == Some("HY000") => {
+        Error::from(QueryError::JoinLimitExceeded)
+    }
+    _ => e,
+})?;
+let (rows, total_count) = fetch_result;

---

Line range hint 636-672: Enhance error messages for schema fetching.

Consider adding more descriptive error messages when schemas cannot be fetched.

 let schemas = self
     .model_cache
     .models(&model_ids)
     .await?
     .into_iter()
     .map(|m| m.schema)
-    .collect::<Vec<_>>();
+    .collect::<Vec<_>>();
+
+if schemas.is_empty() {
+    return Err(Error::from(QueryError::InvalidModel(format!(
+        "No valid schemas found for models: {:?}",
+        model_ids.iter().map(|id| format!("{:#x}", id)).collect::<Vec<_>>()
+    ))));
+}

---

Line range hint 715-734: Consider adding query optimization hints.

For complex composite queries, consider adding SQL hints to help the query optimizer.

 let (rows, total_count) = fetch_entities(
     &self.pool,
     &schemas,
     table,
     model_relation_table,
     entity_relation_column,
-    if where_clause.is_empty() { None } else { Some(&where_clause) },
+    if where_clause.is_empty() { 
+        None 
+    } else { 
+        // Add INDEXED BY hint for better query performance
+        Some(&format!("{} INDEXED BY idx_{}_id", where_clause, table))
+    },
     if having_clause.is_empty() { None } else { Some(&having_clause) },
     order_by,
     limit,
     offset,
     bind_values,
 )
 .await?;

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

crates/torii/core/src/model.rs (4)

401-414: Ohayo sensei! Consider grouping parameters into a dedicated struct.

The function has 11 parameters which could be grouped into a dedicated struct for better maintainability and readability.

+#[derive(Debug)]
+pub struct EntityQueryParams {
+    pub schemas: Vec<Ty>,
+    pub table_name: String,
+    pub model_relation_table: String,
+    pub entity_relation_column: String,
+    pub where_clause: Option<String>,
+    pub having_clause: Option<String>,
+    pub order_by: Option<String>,
+    pub limit: Option<u32>,
+    pub offset: Option<u32>,
+    pub bind_values: Vec<String>,
+}

-pub async fn fetch_entities(
-    pool: &Pool<sqlx::Sqlite>,
-    schemas: &[Ty],
-    table_name: &str,
-    model_relation_table: &str,
-    entity_relation_column: &str,
-    where_clause: Option<&str>,
-    having_clause: Option<&str>,
-    order_by: Option<&str>,
-    limit: Option<u32>,
-    offset: Option<u32>,
-    bind_values: Vec<String>,
-) -> Result<(Vec<sqlx::sqlite::SqliteRow>, u32), Error>
+pub async fn fetch_entities(
+    pool: &Pool<sqlx::Sqlite>,
+    params: EntityQueryParams,
+) -> Result<(Vec<sqlx::sqlite::SqliteRow>, u32), Error>

---

456-460: Document the MAX_JOINS constant with SQLite-specific context.

The MAX_JOINS constant is crucial for handling SQLite's join limits, but its significance isn't documented.

-    const MAX_JOINS: usize = 64;
+    /// SQLite has a compile-time limit on the maximum number of joins in a single query.
+    /// This constant ensures we stay below this limit by chunking our queries.
+    /// See: https://www.sqlite.org/limits.html
+    const MAX_JOINS: usize = 64;

---

488-498: Optimize the count query for better performance.

The count query includes unnecessary columns in the subquery. For counting, we only need the ID column.

-            "SELECT COUNT(*) FROM (SELECT {}.id, group_concat({}.model_id) as model_ids FROM [{}] \
+            "SELECT COUNT(*) FROM (SELECT {}.id FROM [{}] \
             {} {} GROUP BY {}.id {})",
             table_name,
-            model_relation_table,
             table_name,

---

501-505: Add error context for database operations.

Enhance error handling by providing context about which query failed.

-        let chunk_count: u32 = count_stmt.fetch_one(pool).await?;
+        let chunk_count: u32 = count_stmt
+            .fetch_one(pool)
+            .await
+            .map_err(|e| Error::Database(format!("Failed to execute count query: {}", e)))?;

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro