Catch aggregate calls earlier in indexes.

edgedb · May 16, 2024 · 74ac34e · 74ac34e
1 parent 0690ac0
commit 74ac34e
Show file tree

Hide file tree

Showing 5 changed files with 50 additions and 23 deletions.
diff --git a/edb/ir/utils.py b/edb/ir/utils.py
@@ -490,15 +490,6 @@ def find_potentially_visible(
     return visible_paths
 
 
-def contains_set_of_op(ir: irast.Base) -> bool:
-    def flt(n: irast.Call) -> bool:
-        return any(
-            arg.param_typemod == ft.TypeModifier.SetOfType
-            for arg in n.args.values()
-        )
-    return bool(ast.find_children(ir, irast.Call, flt, terminate_early=True))
-
-
 def is_singleton_set_of_call(
     call: irast.Call
 ) -> bool:

diff --git a/edb/schema/indexes.py b/edb/schema/indexes.py
@@ -798,6 +798,7 @@ def compile_expr_field(
         value: s_expr.Expression,
         track_schema_ref_exprs: bool=False,
     ) -> s_expr.CompiledExpression:
+        from edb.ir import ast as irast
         from edb.ir import utils as irutils
 
         if field.name in {'expr', 'except_expr'}:
@@ -850,11 +851,20 @@ def compile_expr_field(
                     has_multi = True
                     break
 
-            if has_multi and irutils.contains_set_of_op(expr.irast):
+            if set_of_op := irutils.find_set_of_op(
+                expr.irast,
+                has_multi,
+            ):
+                label = (
+                    'function'
+                    if isinstance(set_of_op, irast.FunctionCall) else
+                    'operator'
+                )
+                op_name = str(set_of_op.func_shortname)
                 raise errors.SchemaDefinitionError(
-                    "cannot use aggregate functions or operators "
-                    "in an index expression",
-                    span=self.span,
+                    f"cannot use aggregate {label} '{op_name}' "
+                    f"in an index expression",
+                    span=set_of_op.span
                 )
 
             # compile the expression to sql to preempt errors downstream

diff --git a/tests/test_edgeql_ddl.py b/tests/test_edgeql_ddl.py
@@ -13217,16 +13217,28 @@ async def test_edgeql_ddl_index_07(self):
 
     async def test_edgeql_ddl_index_08(self):
         async with self.assertRaisesRegexTx(
-                edgedb.UnsupportedFeatureError,
-                'set returning operator std::DISTINCT is not supported '
-                'in singleton expressions'):
+                edgedb.SchemaDefinitionError,
+                "cannot use aggregate operator 'std::DISTINCT' "
+                "in an index expression"):
             await self.con.execute(r"""
                 CREATE TYPE default::IndexNonSingletonTest {
                     CREATE PROPERTY has_bad_index: std::str;
                     CREATE INDEX ON (DISTINCT (.has_bad_index));
                 };
             """)
 
+    async def test_edgeql_ddl_index_09(self):
+        async with self.assertRaisesRegexTx(
+                edgedb.SchemaDefinitionError,
+                "cannot use aggregate function 'std::count' "
+                "in an index expression"):
+            await self.con.execute(r"""
+                CREATE TYPE default::IndexNonSingletonTest {
+                    CREATE PROPERTY has_bad_index: std::str;
+                    CREATE INDEX ON (std::count (.has_bad_index));
+                };
+            """)
+
     async def test_edgeql_ddl_abstract_index_01(self):
         for _ in range(2):
             await self.con.execute('''

diff --git a/tests/test_indexes.py b/tests/test_indexes.py
@@ -173,8 +173,8 @@ async def test_index_05(self):
 
         async with self.assertRaisesRegexTx(
             edgedb.SchemaDefinitionError,
-            "cannot use aggregate functions or operators in an "
-            "index expression",
+            "cannot use aggregate operator 'std::EXISTS' "
+            "in an index expression",
         ):
             await self.con.execute(
                 """
@@ -186,8 +186,8 @@ async def test_index_05(self):
 
         async with self.assertRaisesRegexTx(
             edgedb.SchemaDefinitionError,
-            "cannot use aggregate functions or operators in an "
-            "index expression",
+            "cannot use aggregate function 'std::count' "
+            "in an index expression",
         ):
             await self.con.execute(
                 """

diff --git a/tests/test_schema.py b/tests/test_schema.py
@@ -2874,9 +2874,9 @@ def test_schema_constraint_singleton_05b(self):
         """
 
     @tb.must_fail(
-        errors.UnsupportedFeatureError,
-        'set returning operator std::DISTINCT is not supported '
-        'in singleton expressions',
+        errors.SchemaDefinitionError,
+        "cannot use aggregate operator 'std::DISTINCT' "
+        "in an index expression",
     )
     def test_schema_index_non_singleton_01(self):
         """
@@ -2887,6 +2887,20 @@ def test_schema_index_non_singleton_01(self):
         }
         """
 
+    @tb.must_fail(
+        errors.SchemaDefinitionError,
+        "cannot use aggregate function 'std::count' "
+        "in an index expression",
+    )
+    def test_schema_index_non_singleton_02(self):
+        """
+        type IndexNonSingletonTest {
+            property has_bad_index -> str;
+
+            index on (count(.has_bad_index))
+        }
+        """
+
     @tb.must_fail(
         errors.SchemaError,
         "cannot create union \\(test::X | test::Y\\) with property 'a' using "