Catch aggregate calls in constraints and indexes in ql compiler (#7343)

- Catches aggregate calls in constraints and indexes earlier in the compilation process
    - Raises errors in migration create which were previously only caught on migrate
- Fixes abstract constraints not raising errors when using aggregate functions
- Adds more information in errors about aggregate functions

Given the schema:
```
  type User {
     name: str;
  }

  type HasTwoOrFewerUsers {
     users: User;
     constraint expression on (count(.users) <= 2);
  }
```

The following error is produced:
```
error: cannot use aggregate functions or operators in a non-aggregating constraint
   │
25 │      constraint expression on (count(.users) <= 2);
   │                                ^^^^^^^^^^^^^ error
```

edb/ir/utils.py

@@ -48,6 +48,7 @@
 from edb.common import ordered
 
 from edb.edgeql import qltypes as ft
+from edb.schema import name as sn
 
 from . import ast as irast
 from . import typeutils
@@ -489,13 +490,47 @@ def find_potentially_visible(
     return visible_paths
 
 
-def contains_set_of_op(ir: irast.Base) -> bool:
+def is_singleton_set_of_call(
+    call: irast.Call
+) -> bool:
+    # Some set functions and operators are allowed in singleton mode
+    # as long as their inputs are singletons
+
+    return call.func_shortname in {
+        sn.QualName('std', 'IN'),
+        sn.QualName('std', 'NOT IN'),
+        sn.QualName('std', 'EXISTS'),
+        sn.QualName('std', '??'),
+        sn.QualName('std', 'IF'),
+    }
+
+
+def has_set_of_param(
+    call: irast.Call,
+) -> bool:
+    return any(
+        arg.param_typemod == ft.TypeModifier.SetOfType
+        for arg in call.args.values()
+    )
+
+
+def returns_set_of(
+    call: irast.Call,
+) -> bool:
+    return call.typemod == ft.TypeModifier.SetOfType
+
+
+def find_set_of_op(
+    ir: irast.Base,
+    has_multi_param: bool,
+) -> Optional[irast.Call]:
     def flt(n: irast.Call) -> bool:
-        return any(
-            arg.param_typemod == ft.TypeModifier.SetOfType
-            for arg in n.args.values()
+        return (
+            (has_multi_param or not is_singleton_set_of_call(n))
+            and (has_set_of_param(n) or returns_set_of(n))
         )
-    return bool(ast.find_children(ir, irast.Call, flt, terminate_early=True))
+    calls = ast.find_children(ir, irast.Call, flt, terminate_early=True)
+    return next(iter(calls or []), None)
 
 
 T = TypeVar('T')

edb/pgsql/compiler/expr.py

@@ -433,10 +433,16 @@ def compile_OperatorCall(
                 dispatch.compile(r_expr, ctx=ctx),
             ],
         )
-    elif expr.typemod is ql_ft.TypeModifier.SetOfType:
+    elif irutils.is_singleton_set_of_call(expr):
+        pass
+    elif irutils.returns_set_of(expr):
         raise errors.UnsupportedFeatureError(
-            f'set returning operator {expr.func_shortname} is not supported '
-            f'in singleton expressions')
+            f"set returning operator '{expr.func_shortname}' is not supported "
+            f"in singleton expressions")
+    elif irutils.has_set_of_param(expr):
+        raise errors.UnsupportedFeatureError(
+            f"aggregate operator '{expr.func_shortname}' is not supported "
+            f"in singleton expressions")
 
     args, maybe_null = _compile_call_args(expr, ctx=ctx)
     return _wrap_call(
@@ -683,9 +689,15 @@ def compile_FunctionCall(
             f'unimplemented function for singleton mode: {fname}'
         )
 
-    if expr.typemod is ql_ft.TypeModifier.SetOfType:
+    if irutils.is_singleton_set_of_call(expr):
+        pass
+    elif irutils.returns_set_of(expr):
         raise errors.UnsupportedFeatureError(
             'set returning functions are not supported in simple expressions')
+    elif irutils.has_set_of_param(expr):
+        raise errors.UnsupportedFeatureError(
+            f"aggregate function '{expr.func_shortname}' is not supported "
+            f"in singleton expressions")
 
     args, maybe_null = _compile_call_args(expr, ctx=ctx)
 

edb/schema/constraints.py

@@ -931,12 +931,20 @@ def _populate_concrete_constraint_attrs(
                     span=sourcectx
                 )
 
-            if has_any_multi and ir_utils.contains_set_of_op(
-                    final_subjectexpr.irast):
-                raise errors.InvalidConstraintDefinitionError(
-                    "cannot use aggregate functions or operators "
-                    "in a non-aggregating constraint",
-                    span=sourcectx
+            if set_of_op := ir_utils.find_set_of_op(
+                final_subjectexpr.irast,
+                has_any_multi,
+            ):
+                label = (
+                    'function'
+                    if isinstance(set_of_op, irast.FunctionCall) else
+                    'operator'
+                )
+                op_name = str(set_of_op.func_shortname)
+                raise errors.UnsupportedFeatureError(
+                    f"cannot use SET OF {label} '{op_name}' "
+                    f"in a constraint",
+                    span=set_of_op.span
                 )
 
             if (

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 SET OF {label} '{op_name}' "
+                    f"in an index expression",
+                    span=set_of_op.span
                 )
 
             # compile the expression to sql to preempt errors downstream

tests/test_constraints.py

@@ -1400,9 +1400,9 @@ async def test_constraints_ddl_08(self):
             """)
 
         async with self.assertRaisesRegexTx(
-            edgedb.InvalidConstraintDefinitionError,
-            "cannot use aggregate functions or operators in a "
-            "non-aggregating constraint",
+            edgedb.UnsupportedFeatureError,
+            "cannot use SET OF operator 'std::EXISTS' "
+            "in a constraint",
         ):
             await self.con.execute("""
                 ALTER TYPE ObjCnstr2 {
@@ -1411,9 +1411,9 @@ async def test_constraints_ddl_08(self):
             """)
 
         async with self.assertRaisesRegexTx(
-            edgedb.InvalidConstraintDefinitionError,
-            "cannot use aggregate functions or operators in a "
-            "non-aggregating constraint",
+            edgedb.UnsupportedFeatureError,
+            "cannot use SET OF operator 'std::EXISTS' "
+            "in a constraint",
         ):
             await self.con.execute("""
                 ALTER TYPE ObjCnstr2 {
@@ -2119,3 +2119,147 @@ async def test_constraints_abstract_object_03(self):
             await self.con.execute("""
                 update ChatBase set { messages := 'hello world' };
             """)
+
+    async def test_constraints_singleton_set_ops_01(self):
+        await self.con.execute(
+            """
+            create type X {
+                create property a -> int64 {
+                    create constraint expression on (
+                        __subject__ in {1}
+                    );
+                }
+            };
+        """)
+
+        async with self.assertRaisesRegexTx(
+            edgedb.UnsupportedFeatureError,
+            "cannot use SET OF operator 'std::IN' "
+            "in a constraint"
+        ):
+            await self.con.execute(
+                """
+                create type Y {
+                    create multi property a -> int64 {
+                        create constraint expression on (
+                            __subject__ in {1}
+                        );
+                    }
+                };
+            """)
+
+    async def test_constraints_singleton_set_ops_02(self):
+        await self.con.execute(
+            """
+            create type X {
+                create property a -> int64 {
+                    create constraint expression on (
+                        __subject__ not in {1}
+                    );
+                }
+            };
+        """)
+
+        async with self.assertRaisesRegexTx(
+            edgedb.UnsupportedFeatureError,
+            "cannot use SET OF operator 'std::NOT IN' "
+            "in a constraint"
+        ):
+            await self.con.execute(
+                """
+                create type Y {
+                    create multi property a -> int64 {
+                        create constraint expression on (
+                            __subject__ not in {1}
+                        );
+                    }
+                };
+            """)
+
+    async def test_constraints_singleton_set_ops_03(self):
+        await self.con.execute(
+            """
+            create type X {
+                create property a -> int64 {
+                    create constraint expression on (
+                        exists(__subject__)
+                    );
+                }
+            };
+        """)
+
+        async with self.assertRaisesRegexTx(
+            edgedb.UnsupportedFeatureError,
+            "cannot use SET OF operator 'std::EXISTS' "
+            "in a constraint"
+        ):
+            await self.con.execute(
+                """
+                create type Y {
+                    create multi property a -> int64 {
+                        create constraint expression on (
+                            exists(__subject__)
+                        );
+                    }
+                };
+            """)
+
+    async def test_constraints_singleton_set_ops_04(self):
+        await self.con.execute(
+            """
+            create type X {
+                create property a -> int64 {
+                    create constraint expression on (
+                        __subject__ ?? 1 = 0
+                    );
+                }
+            };
+        """)
+
+        async with self.assertRaisesRegexTx(
+            edgedb.UnsupportedFeatureError,
+            r"cannot use SET OF operator 'std::\?\?' "
+            r"in a constraint"
+        ):
+            await self.con.execute(
+                """
+                create type Y {
+                    create multi property a -> int64 {
+                        create constraint expression on (
+                            __subject__ ?? 1 = 0
+                        );
+                    }
+                };
+            """)
+
+    async def test_constraints_singleton_set_ops_05(self):
+        await self.con.execute(
+            """
+            create type X {
+                create property a -> tuple<bool, int64> {
+                    create constraint expression on (
+                        __subject__.1 < 0
+                        if __subject__.0 else
+                        __subject__.1 >= 0
+                    );
+                }
+            };
+        """)
+
+        async with self.assertRaisesRegexTx(
+            edgedb.UnsupportedFeatureError,
+            "cannot use SET OF operator 'std::IF' "
+            "in a constraint"
+        ):
+            await self.con.execute(
+                """
+                create type Y {
+                    create multi property a -> tuple<bool, int64> {
+                        create constraint expression on (
+                            __subject__.1 < 0
+                            if __subject__.0 else
+                            __subject__.1 >= 0
+                        );
+                    }
+                };
+            """)