From 0eeb8ae6d697d38e5204bc408a399be3023ded12 Mon Sep 17 00:00:00 2001
From: Larko <59736843+Larkooo@users.noreply.github.com>
Date: Thu, 12 Dec 2024 12:16:05 +0700
Subject: [PATCH] fix: keys composite clause (#2798)

* fix: keys composite clause

* count

* c

* f
---
 crates/torii/grpc/src/server/mod.rs | 50 ++++++++++++++++++++++-------
 1 file changed, 39 insertions(+), 11 deletions(-)

diff --git a/crates/torii/grpc/src/server/mod.rs b/crates/torii/grpc/src/server/mod.rs
index bb04979b83..113832b1e8 100644
--- a/crates/torii/grpc/src/server/mod.rs
+++ b/crates/torii/grpc/src/server/mod.rs
@@ -764,16 +764,31 @@ impl DojoWorld {
         let (where_clause, having_clause, join_clause, bind_values) =
             build_composite_clause(table, model_relation_table, &composite)?;
 
-        let count_query = format!(
-            r#"
-            SELECT COUNT(DISTINCT [{table}].id)
-            FROM [{table}]
-            JOIN {model_relation_table} ON [{table}].id = {model_relation_table}.entity_id
-            {join_clause}
-            {where_clause}
-            {having_clause}
-            "#,
-        );
+        let count_query = if !having_clause.is_empty() {
+            format!(
+                r#"
+                SELECT COUNT(*) FROM (
+                    SELECT [{table}].id
+                    FROM [{table}]
+                    JOIN {model_relation_table} ON [{table}].id = {model_relation_table}.entity_id
+                    {join_clause}
+                    {where_clause}
+                    GROUP BY [{table}].id
+                    {having_clause}
+                ) as filtered_count
+                "#
+            )
+        } else {
+            format!(
+                r#"
+                SELECT COUNT(DISTINCT [{table}].id)
+                FROM [{table}]
+                JOIN {model_relation_table} ON [{table}].id = {model_relation_table}.entity_id
+                {join_clause}
+                {where_clause}
+                "#
+            )
+        };
 
         let mut count_query = sqlx::query_scalar::<_, u32>(&count_query);
         for value in &bind_values {
@@ -1146,7 +1161,7 @@ fn build_keys_pattern(clause: &proto::types::KeysClause) -> Result<String, Error
     let mut keys_pattern = format!("^{}", keys.join("/"));
 
     if clause.pattern_matching == proto::types::PatternMatching::VariableLen as i32 {
-        keys_pattern += &format!("({})*", KEY_PATTERN);
+        keys_pattern += &format!("/({})*", KEY_PATTERN);
     }
     keys_pattern += "/$";
 
@@ -1184,6 +1199,19 @@ fn build_composite_clause(
                 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
+                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(group_concat({model_relation_table}.model_id), '{:#x}') > 0",
+                        model_id
+                    ));
+                }
             }
             ClauseType::Member(member) => {
                 let comparison_operator = ComparisonOperator::from_repr(member.operator as usize)