Merge pull request #607 from 1llu5ion/fix/606-window-expression-arithmetic

Fix chained operators on window expressions

Author: go-jet

internal/jet/bool_expression.go

@@ -101,6 +101,13 @@ type boolExpressionWrapper struct {
 	Expression
 }
 
+func (b *boolExpressionWrapper) setRoot(root Expression) {
+	b.Expression.setRoot(root)
+	if boolRoot, ok := root.(BoolExpression); ok {
+		b.boolInterfaceImpl.root = boolRoot
+	}
+}
+
 func newBoolExpressionWrap(expression Expression) BoolExpression {
 	boolExpressionWrap := &boolExpressionWrapper{Expression: expression}
 	boolExpressionWrap.boolInterfaceImpl.root = boolExpressionWrap

internal/jet/float_expression.go

@@ -101,6 +101,13 @@ type floatExpressionWrapper struct {
 	Expression
 }
 
+func (f *floatExpressionWrapper) setRoot(root Expression) {
+	f.Expression.setRoot(root)
+	if floatRoot, ok := root.(FloatExpression); ok {
+		f.floatInterfaceImpl.root = floatRoot
+	}
+}
+
 func newFloatExpressionWrap(expression Expression) FloatExpression {
 	floatExpressionWrap := &floatExpressionWrapper{Expression: expression}
 	floatExpressionWrap.floatInterfaceImpl.root = floatExpressionWrap

internal/jet/integer_expression.go

@@ -140,6 +140,13 @@ type integerExpressionWrapper struct {
 	Expression
 }
 
+func (i *integerExpressionWrapper) setRoot(root Expression) {
+	i.Expression.setRoot(root)
+	if integerRoot, ok := root.(IntegerExpression); ok {
+		i.integerInterfaceImpl.root = integerRoot
+	}
+}
+
 func newIntExpressionWrap(expression Expression) IntegerExpression {
 	intExpressionWrap := &integerExpressionWrapper{Expression: expression}
 	intExpressionWrap.integerInterfaceImpl.root = intExpressionWrap

internal/jet/window_func_test.go

@@ -19,3 +19,17 @@ func TestWindowFunctions(t *testing.T) {
 	assertClauseSerialize(t, ORDER_BY(table1Col1).RANGE(PRECEDING(UNBOUNDED), CURRENT_ROW),
 		"(ORDER BY table1.col1 RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)")
 }
+
+func TestWindowFunctionChainedOperators(t *testing.T) {
+	assertClauseSerialize(t,
+		ROW_NUMBER().OVER(ORDER_BY(table1ColInt)).ADD(Int(1)),
+		"(ROW_NUMBER() OVER (ORDER BY table1.col_int) + $1)", int64(1))
+	assertClauseSerialize(t,
+		SUMf(table1ColFloat).OVER(PARTITION_BY(table1ColInt)).SUB(
+			SUMf(table2ColFloat).OVER(PARTITION_BY(table1ColInt)),
+		),
+		"(SUM(table1.col_float) OVER (PARTITION BY table1.col_int) - SUM(table2.col_float) OVER (PARTITION BY table1.col_int))")
+	assertClauseSerialize(t,
+		BOOL_AND(table1ColBool).OVER(PARTITION_BY(table1ColInt)).AND(table2ColBool),
+		"(BOOL_AND(table1.col_bool) OVER (PARTITION BY table1.col_int) AND table2.col_bool)")
+}