From 46dd38d762e9a885468c027094f70c12107ed1fe Mon Sep 17 00:00:00 2001
From: David Lutterkort <lutter@watzmann.net>
Date: Mon, 26 Feb 2024 09:39:19 -0800
Subject: [PATCH] graph: Allow more SQL constructs in aggregates

This adds comparison operators, various 'is true|false|null|distinct from'
constructs, casts, and 'case when ..'
---
 graph/src/schema/input/sqlexpr.rs             | 64 +++++++++++--------
 .../test_schemas/ts_expr_simple.graphql       |  8 +++
 2 files changed, 45 insertions(+), 27 deletions(-)

diff --git a/graph/src/schema/input/sqlexpr.rs b/graph/src/schema/input/sqlexpr.rs
index 4f2d957846e..5e0d8c95f6c 100644
--- a/graph/src/schema/input/sqlexpr.rs
+++ b/graph/src/schema/input/sqlexpr.rs
@@ -109,19 +109,43 @@ impl<'a> VisitExpr<'a> {
             }
             Function(func) => self.visit_func(func),
             Value(_) => Ok(()),
+            Case {
+                operand,
+                conditions,
+                results,
+                else_result,
+            } => {
+                if let Some(operand) = operand {
+                    self.visit_expr(operand)?;
+                }
+                for condition in conditions {
+                    self.visit_expr(condition)?;
+                }
+                for result in results {
+                    self.visit_expr(result)?;
+                }
+                if let Some(else_result) = else_result {
+                    self.visit_expr(else_result)?;
+                }
+                Ok(())
+            }
+            Cast {
+                expr,
+                data_type: _,
+                format: _,
+            } => self.visit_expr(expr),
+            Nested(expr) | IsFalse(expr) | IsNotFalse(expr) | IsTrue(expr) | IsNotTrue(expr)
+            | IsNull(expr) | IsNotNull(expr) => self.visit_expr(expr),
+            IsDistinctFrom(expr1, expr2) | IsNotDistinctFrom(expr1, expr2) => {
+                self.visit_expr(expr1)?;
+                self.visit_expr(expr2)?;
+                Ok(())
+            }
             CompoundIdentifier(_) => self.nope("CompoundIdentifier"),
             JsonAccess { .. } => self.nope("JsonAccess"),
             CompositeAccess { .. } => self.nope("CompositeAccess"),
-            IsFalse(_) => self.nope("IsFalse"),
-            IsNotFalse(_) => self.nope("IsNotFalse"),
-            IsTrue(_) => self.nope("IsTrue"),
-            IsNotTrue(_) => self.nope("IsNotTrue"),
-            IsNull(_) => self.nope("IsNull"),
-            IsNotNull(_) => self.nope("IsNotNull"),
             IsUnknown(_) => self.nope("IsUnknown"),
             IsNotUnknown(_) => self.nope("IsNotUnknown"),
-            IsDistinctFrom(_, _) => self.nope("IsDistinctFrom"),
-            IsNotDistinctFrom(_, _) => self.nope("IsNotDistinctFrom"),
             InList { .. } => self.nope("InList"),
             InSubquery { .. } => self.nope("InSubquery"),
             InUnnest { .. } => self.nope("InUnnest"),
@@ -133,7 +157,6 @@ impl<'a> VisitExpr<'a> {
             AnyOp { .. } => self.nope("AnyOp"),
             AllOp { .. } => self.nope("AllOp"),
             Convert { .. } => self.nope("Convert"),
-            Cast { .. } => self.nope("Cast"),
             TryCast { .. } => self.nope("TryCast"),
             SafeCast { .. } => self.nope("SafeCast"),
             AtTimeZone { .. } => self.nope("AtTimeZone"),
@@ -145,12 +168,10 @@ impl<'a> VisitExpr<'a> {
             Trim { .. } => self.nope("Trim"),
             Overlay { .. } => self.nope("Overlay"),
             Collate { .. } => self.nope("Collate"),
-            Nested(_) => self.nope("Nested"),
             IntroducedString { .. } => self.nope("IntroducedString"),
             TypedString { .. } => self.nope("TypedString"),
             MapAccess { .. } => self.nope("MapAccess"),
             AggregateExpressionWithFilter { .. } => self.nope("AggregateExpressionWithFilter"),
-            Case { .. } => self.nope("Case"),
             Exists { .. } => self.nope("Exists"),
             Subquery(_) => self.nope("Subquery"),
             ArraySubquery(_) => self.nope("ArraySubquery"),
@@ -217,17 +238,9 @@ impl<'a> VisitExpr<'a> {
     fn check_binary_op(&mut self, op: &p::BinaryOperator) -> Result<(), ()> {
         use p::BinaryOperator::*;
         match op {
-            Plus | Minus | Multiply | Divide | Modulo => Ok(()),
+            Plus | Minus | Multiply | Divide | Modulo | PGExp | Gt | Lt | GtEq | LtEq
+            | Spaceship | Eq | NotEq | And | Or => Ok(()),
             StringConcat
-            | Gt
-            | Lt
-            | GtEq
-            | LtEq
-            | Spaceship
-            | Eq
-            | NotEq
-            | And
-            | Or
             | Xor
             | BitwiseOr
             | BitwiseAnd
@@ -238,7 +251,6 @@ impl<'a> VisitExpr<'a> {
             | PGBitwiseXor
             | PGBitwiseShiftLeft
             | PGBitwiseShiftRight
-            | PGExp
             | PGOverlap
             | PGRegexMatch
             | PGRegexIMatch
@@ -258,11 +270,9 @@ impl<'a> VisitExpr<'a> {
     fn check_unary_op(&mut self, op: &p::UnaryOperator) -> Result<(), ()> {
         use p::UnaryOperator::*;
         match op {
-            Plus | Minus => Ok(()),
-            Not | PGBitwiseNot | PGSquareRoot | PGCubeRoot | PGPostfixFactorial
-            | PGPrefixFactorial | PGAbs => {
-                self.not_supported(format!("unary operator {op} is not supported"))
-            }
+            Plus | Minus | Not => Ok(()),
+            PGBitwiseNot | PGSquareRoot | PGCubeRoot | PGPostfixFactorial | PGPrefixFactorial
+            | PGAbs => self.not_supported(format!("unary operator {op} is not supported")),
         }
     }
 }
diff --git a/graph/src/schema/test_schemas/ts_expr_simple.graphql b/graph/src/schema/test_schemas/ts_expr_simple.graphql
index 5a067077482..ed15c14ceb3 100644
--- a/graph/src/schema/test_schemas/ts_expr_simple.graphql
+++ b/graph/src/schema/test_schemas/ts_expr_simple.graphql
@@ -14,4 +14,12 @@ type Stats @aggregation(intervals: ["hour", "day"], source: "Data") {
   price0_sq: BigDecimal! @aggregate(fn: "sum", arg: "power(price0, 2)")
   sum_sq: BigDecimal! @aggregate(fn: "sum", arg: "price0 * price0")
   sum_sq_cross: BigDecimal! @aggregate(fn: "sum", arg: "price0 * price1")
+
+  max_some: BigDecimal!
+    @aggregate(
+      fn: "max"
+      arg: "case when price0 > price1 then price0 else 0 end"
+    )
+
+  max_cast: BigDecimal! @aggregate(fn: "sum", arg: "(price0/7)::int4")
 }