From 8da946fe9b55bc0320318f03d3f3d630cc1db6f3 Mon Sep 17 00:00:00 2001
From: Samir Patel <48686912+54mir@users.noreply.github.com>
Date: Mon, 21 Mar 2022 11:46:13 -0500
Subject: [PATCH] allow empty values to be provided for vars in Row call

queries can become arbitrarily long when variables are used.
This is especially the case when a variable is defined as
'select distinct field from table'
and a user wants to use all the values in a Row call
(which is effectively disabling any condition on the field).

This change allows users to select no values for a variable
associated with a Row call to disable the condition. If that
variable is the only condition (query expands to nothing)
then it interprets it as an All call.
---
 pql/ast.go      |  32 ++++++++++++++-
 pql/ast_test.go | 107 ++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 134 insertions(+), 5 deletions(-)

diff --git a/pql/ast.go b/pql/ast.go
index f61ebba95..2735f4939 100644
--- a/pql/ast.go
+++ b/pql/ast.go
@@ -29,6 +29,9 @@ func (q *Query) ExpandVars(vars map[string]interface{}) (*Query, error) {
 		if err != nil {
 			return nil, err
 		}
+		if len(newCalls) == 0 {
+			return nil, fmt.Errorf("no values to use for variable expansion")
+		}
 		other.Calls = append(other.Calls, newCalls...)
 	}
 
@@ -1011,10 +1014,32 @@ func (c *Call) ExpandVars(vars map[string]interface{}) ([]*Call, error) {
 				if err != nil {
 					return nil, err
 				}
-				if len(newArg) != 1 {
+
+				switch len(newArg) {
+				case 0:
+					if call.Name == "Row" {
+						other.Args[key] = &Call{Name: "All"}
+					} else {
+						return nil, fmt.Errorf("variable: non-Row calls require values be supplied, got: %+v", newArg)
+					}
+				case 1:
+					other.Args[key] = newArg[0]
+				default:
 					return nil, fmt.Errorf("variable: requires single value for argument, got: %+v", newArg)
 				}
-				other.Args[key] = newArg[0]
+
+			}
+		}
+
+		// if the call had children, but due to variable expansion, it now has none - then the user
+		// did not select any values for any variables. If the child was originally a Row call, we equate
+		// this to an All() (i.e. the user chooses to apply no conditions to the query).
+		if len(c.Children) > 0 && len(other.Children) == 0 {
+			if c.Children[0].Name == "Row" {
+				r := Call{Name: "All"}
+				other.Children = []*Call{&r}
+			} else {
+				return nil, fmt.Errorf("variable: non-Row calls require values be supplied")
 			}
 		}
 		return []*Call{&other}, nil
@@ -1025,6 +1050,9 @@ func (c *Call) ExpandVars(vars map[string]interface{}) ([]*Call, error) {
 func (c *Call) expandVars(name string, values []interface{}) []*Call {
 	switch c.Name {
 	case "Row":
+		if len(values) == 0 {
+			return []*Call{}
+		}
 		union := &Call{Name: "Union"}
 		for i := range values {
 			r := Call{Name: "Row", Args: CopyArgs(c.Args)}
diff --git a/pql/ast_test.go b/pql/ast_test.go
index d77fe623b..e96a5f4ef 100644
--- a/pql/ast_test.go
+++ b/pql/ast_test.go
@@ -68,18 +68,37 @@ func TestQuery_ExpandVars(t *testing.T) {
 			output: `Count(Union(Row(animal="cat"), Row(animal="dog"), Row(animal="pig")))`,
 			vars:   map[string]interface{}{"var1": []interface{}{"cat", "dog", "pig"}},
 		},
+		{
+			name:   "ExpandRowEQInterior-NoValues",
+			input:  `count(row(animal=$var1))`,
+			output: `Count(All())`,
+			vars:   map[string]interface{}{"var1": []interface{}{}},
+		},
 		{
 			name:   "ExpandRowEQExterior",
 			input:  `row(animal=$var1)`,
 			output: `Union(Row(animal="cat"))`,
 			vars:   map[string]interface{}{"var1": []interface{}{"cat"}},
 		},
+		{
+			name:    "ExpandRowEQExterior-NoValues",
+			input:   `row(animal=$var1)`,
+			output:  ``,
+			vars:    map[string]interface{}{"var1": []interface{}{}},
+			wantErr: true,
+		},
 		{
 			name:   "ExpandRowGT",
 			input:  `count(row(num>$var1))`,
 			output: `Count(Union(Row(num>5), Row(num>10)))`,
 			vars:   map[string]interface{}{"var1": []interface{}{5, 10}},
 		},
+		{
+			name:   "ExpandRowGT-NoValues",
+			input:  `count(row(num>$var1))`,
+			output: `Count(All())`,
+			vars:   map[string]interface{}{"var1": []interface{}{}},
+		},
 		{
 			name:   "ExpandRowNOT",
 			input:  `count(row(num!=$var1))`,
@@ -98,18 +117,52 @@ func TestQuery_ExpandVars(t *testing.T) {
 			output: `GroupBy(Rows(_field="cat"), Rows(_field="dog"), limit=5)`,
 			vars:   map[string]interface{}{"var1": []interface{}{"cat", "dog"}},
 		},
+		{
+			name:    "ExpandRowsInterior-NoValues",
+			input:   `GroupBy(rows($var1), limit=5)`,
+			output:  `GroupBy(), limit=5)`,
+			vars:    map[string]interface{}{"var1": []interface{}{}},
+			wantErr: true,
+		},
 		{
 			name:   "ExpandRowsExterior",
 			input:  `rows($var1)`,
 			output: `Rows(_field="cat")` + "\n" + `Rows(_field="dog")`,
 			vars:   map[string]interface{}{"var1": []interface{}{"cat", "dog"}},
 		},
+		{
+			name:    "ExpandRowsExterior-NoValues",
+			input:   `rows($var1)`,
+			output:  ``,
+			vars:    map[string]interface{}{"var1": []interface{}{}},
+			wantErr: true,
+		},
 		{
 			name:   "ExpandRowAndRows",
 			input:  `GroupBy(Rows($animal), limit=7, filter=Row(size=$size))`,
 			output: `GroupBy(Rows(_field="cat"), Rows(_field="dog"), filter=Union(Row(size="lg"), Row(size="md")), limit=7)`,
 			vars:   map[string]interface{}{"animal": []interface{}{"cat", "dog"}, "size": []interface{}{"lg", "md"}},
 		},
+		{
+			name:    "ExpandRowAndRows-NoValues1",
+			input:   `GroupBy(Rows($animal), limit=7, filter=Row(size=$size))`,
+			output:  `GroupBy(filter=All(), limit=7)`,
+			vars:    map[string]interface{}{"animal": []interface{}{}, "size": []interface{}{}},
+			wantErr: true,
+		},
+		{
+			name:   "ExpandRowAndRows-NoValues2",
+			input:  `GroupBy(Rows($animal), limit=7, filter=Row(size=$size))`,
+			output: `GroupBy(Rows(_field="cat"), Rows(_field="dog"), filter=All(), limit=7)`,
+			vars:   map[string]interface{}{"animal": []interface{}{"cat", "dog"}, "size": []interface{}{}},
+		},
+		{
+			name:    "ExpandRowAndRows-NoValues3",
+			input:   `GroupBy(Rows($animal), limit=7, filter=Row(size=$size))`,
+			output:  `GroupBy(filter=Union(Row(size="lg"), Row(size="md")), limit=7)`,
+			vars:    map[string]interface{}{"animal": []interface{}{}, "size": []interface{}{"lg", "md"}},
+			wantErr: true,
+		},
 		{
 			name:    "ExpandBad",
 			input:   `$animal`,
@@ -128,31 +181,79 @@ func TestQuery_ExpandVars(t *testing.T) {
 			output: `Intersect(ConstRow(columns=[5,10]), Union(Row(animal="cat"), Row(animal="dog")))`,
 			vars:   map[string]interface{}{"var1": []interface{}{"cat", "dog"}, "var2": []interface{}{5, 10}},
 		},
+		{
+			name:   "ExpandAsCSV-NoValues",
+			input:  `Intersect(ConstRow(columns=$var2), Row(animal=$var1))`,
+			output: `Intersect(ConstRow(columns=[]))`,
+			vars:   map[string]interface{}{"var1": []interface{}{}, "var2": []interface{}{}},
+		},
 		{
 			name:   "ExpandPercentile",
 			input:  `Percentile(field="bytes", nth=99.0, filter=Row(level=$animal))`,
 			output: `Percentile(field="bytes", filter=Union(Row(level="cat"), Row(level="dog")), nth=99)`,
 			vars:   map[string]interface{}{"animal": []interface{}{"cat", "dog"}, "columns": []interface{}{5, 10}},
 		},
+		{
+			name:   "ExpandPercentile-NoValues",
+			input:  `Percentile(field="bytes", nth=99.0, filter=Row(level=$animal))`,
+			output: `Percentile(field="bytes", filter=All(), nth=99)`,
+			vars:   map[string]interface{}{"animal": []interface{}{}, "columns": []interface{}{}},
+		},
+		{
+			name:   "Union-NoValues-1",
+			input:  `Count(Union(row(x=$var1), row(y=$var2)), limit=5)`,
+			output: `Count(Union(Union(Row(y="cat"), Row(y="dog"))), limit=5)`,
+			vars:   map[string]interface{}{"var1": []interface{}{}, "var2": []interface{}{"cat", "dog"}},
+		},
+		{
+			name:    "Rows-NoValues-1",
+			input:   `groupby(rows($var1))`,
+			output:  `GroupBy()`,
+			vars:    map[string]interface{}{"var1": []interface{}{}},
+			wantErr: true,
+		},
+		{
+			name:   "Extract-NoValues-1",
+			input:  `Extract(Limit(Row(animals=$var1), limit=1000), Rows($var2))`,
+			output: `Extract(Limit(All(), limit=1000))`,
+			vars:   map[string]interface{}{"var1": []interface{}{}, "var2": []interface{}{}},
+		},
+		{
+			name:   "Extract-NoValues-2",
+			input:  `Extract(Limit(Row(animals=$var1), limit=1000), Rows($var2))`,
+			output: `Extract(Limit(Union(Row(animals="a"), Row(animals="b")), limit=1000))`,
+			vars:   map[string]interface{}{"var1": []interface{}{"a", "b"}, "var2": []interface{}{}},
+		},
+		{
+			name:   "Extract-NoValues-3",
+			input:  `Extract(Limit(Row(animals=$var1), limit=1000), Rows($var2))`,
+			output: `Extract(Limit(All(), limit=1000), Rows(_field="a"), Rows(_field="b"))`,
+			vars:   map[string]interface{}{"var1": []interface{}{}, "var2": []interface{}{"a", "b"}},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			q, err := pql.NewParser(strings.NewReader(tt.input)).Parse()
 			if err != nil {
 				if !tt.wantErr {
-					t.Errorf("Parse error = %v, wantErr %v", err, tt.wantErr)
+					t.Errorf("Parse error = %v", err)
 				}
 				return
 			}
 
 			got, err := q.ExpandVars(tt.vars)
 			if err != nil {
-				t.Errorf("Query.ExpandVars() error = %v", err)
+				if !tt.wantErr {
+					t.Errorf("Query.ExpandVars() error = %v", err)
+				}
 				return
 			}
-			if tt.output != got.String() {
+			if tt.output != got.String() && !tt.wantErr {
 				t.Errorf("got %v, want %v", got, tt.output)
 			}
+			if tt.wantErr {
+				t.Error("test succeeded, but expected error")
+			}
 		})
 	}
 }