From 4fc77de49b01cf168f0543a3a5aaffc365bcbecd Mon Sep 17 00:00:00 2001
From: Nasr <larko236@outlook.com>
Date: Fri, 7 Feb 2025 19:18:04 +0800
Subject: [PATCH] feat(torii-grpc): correct keys clause models predicate

---
 crates/torii/grpc/src/server/mod.rs | 71 +++++++++++++++++++----------
 1 file changed, 47 insertions(+), 24 deletions(-)

diff --git a/crates/torii/grpc/src/server/mod.rs b/crates/torii/grpc/src/server/mod.rs
index 1c32d78bda..90af12ecbb 100644
--- a/crates/torii/grpc/src/server/mod.rs
+++ b/crates/torii/grpc/src/server/mod.rs
@@ -445,15 +445,33 @@ impl DojoWorld {
         entity_updated_after: Option<String>,
     ) -> Result<(Vec<proto::types::Entity>, u32), Error> {
         let keys_pattern = build_keys_pattern(keys_clause)?;
+        let model_selectors: Vec<Felt> =
+            keys_clause.models.iter().map(|model| compute_selector_from_tag(model)).collect();
 
-        let where_clause = format!(
-            "{table}.keys REGEXP ? {}",
-            if entity_updated_after.is_some() {
-                format!("AND {table}.updated_at >= ?")
-            } else {
-                String::new()
-            }
-        );
+        // Build WHERE clause that only applies the keys pattern when the model matches
+        let where_clause = if model_selectors.is_empty() {
+            format!(
+                "{table}.keys REGEXP ? {}",
+                if entity_updated_after.is_some() {
+                    format!("AND {table}.updated_at >= ?")
+                } else {
+                    String::new()
+                }
+            )
+        } else {
+            format!(
+                "({table}.keys REGEXP ? AND {model_relation_table}.model_id IN ({models})) OR {model_relation_table}.model_id NOT IN ({models}) {}",
+                if entity_updated_after.is_some() {
+                    format!("AND {table}.updated_at >= ?")
+                } else {
+                    String::new()
+                },
+                models = model_selectors.iter()
+                    .map(|selector| format!("'{:#x}'", selector))
+                    .collect::<Vec<_>>()
+                    .join(", "),
+            )
+        };
 
         let mut bind_values = vec![keys_pattern];
         if let Some(entity_updated_after) = entity_updated_after.clone() {
@@ -704,7 +722,7 @@ impl DojoWorld {
         entity_updated_after: Option<String>,
     ) -> Result<(Vec<proto::types::Entity>, u32), Error> {
         let (where_clause, bind_values) =
-            build_composite_clause(table, &composite, entity_updated_after)?;
+            build_composite_clause(table, model_relation_table, &composite, entity_updated_after)?;
 
         let entity_models =
             entity_models.iter().map(|model| compute_selector_from_tag(model)).collect::<Vec<_>>();
@@ -1157,6 +1175,7 @@ fn build_keys_pattern(clause: &proto::types::KeysClause) -> Result<String, Error
 // builds a composite clause for a query
 fn build_composite_clause(
     table: &str,
+    model_relation_table: &str,
     composite: &proto::types::CompositeClause,
     entity_updated_after: Option<String>,
 ) -> Result<(String, Vec<String>), Error> {
@@ -1181,20 +1200,20 @@ fn build_composite_clause(
             ClauseType::Keys(keys) => {
                 let keys_pattern = build_keys_pattern(keys)?;
                 bind_values.push(keys_pattern);
-                where_clauses.push(format!("({table}.keys REGEXP ?)"));
-
-                // Add model checks for specified models
+                let model_selectors: Vec<Felt> =
+                    keys.models.iter().map(|model| compute_selector_from_tag(model)).collect();
 
-                // NOTE: disabled since we are now using the top level entity models
-
-                // for model in &keys.models {
-                //     let (namespace, model_name) = model
-                //         .split_once('-')
-                //         .ok_or(QueryError::InvalidNamespacedModel(model.clone()))?;
-                //     let model_id = compute_selector_from_names(namespace, model_name);
-
-                //     having_clauses.push(format!("INSTR(model_ids, '{:#x}') > 0", model_id));
-                // }
+                if model_selectors.is_empty() {
+                    where_clauses.push(format!("({table}.keys REGEXP ?)"));
+                } else {
+                    where_clauses.push(format!(
+                        "({table}.keys REGEXP ? AND {model_relation_table}.model_id IN ({models})) OR {model_relation_table}.model_id NOT IN ({models})",
+                        models = model_selectors.iter()
+                            .map(|selector| format!("'{:#x}'", selector))
+                            .collect::<Vec<_>>()
+                            .join(", "),
+                    ));
+                }
             }
             ClauseType::Member(member) => {
                 let comparison_operator = ComparisonOperator::from_repr(member.operator as usize)
@@ -1236,8 +1255,12 @@ fn build_composite_clause(
             }
             ClauseType::Composite(nested) => {
                 // Handle nested composite by recursively building the clause
-                let (nested_where, nested_values) =
-                    build_composite_clause(table, nested, entity_updated_after.clone())?;
+                let (nested_where, nested_values) = build_composite_clause(
+                    table,
+                    model_relation_table,
+                    nested,
+                    entity_updated_after.clone(),
+                )?;
 
                 if !nested_where.is_empty() {
                     where_clauses.push(nested_where);