Bugfix: Adding stored query to array should expand the stored query. (#122)

Author: scudette

fixtures/multi_vql_queries.golden

@@ -1338,5 +1338,83 @@
       "StructValue.src_ip": "127.0.0.1",
       "StructValue.SrcIp": null
     }
+  ],
+  "088/000 Adding Stored Query: LET X = SELECT value AS Foo FROM range(start=1, end=4)": null,
+  "088/001 Adding Stored Query: SELECT (dict(Foo=12), ) + X, X + (dict(Foo=12), ) FROM scope()": [
+    {
+      "(dict(Foo=12), ) + X": [
+        {
+          "Foo": 12
+        },
+        {
+          "Foo": 1
+        },
+        {
+          "Foo": 2
+        },
+        {
+          "Foo": 3
+        },
+        {
+          "Foo": 4
+        }
+      ],
+      "X + (dict(Foo=12), )": [
+        {
+          "Foo": 1
+        },
+        {
+          "Foo": 2
+        },
+        {
+          "Foo": 3
+        },
+        {
+          "Foo": 4
+        },
+        {
+          "Foo": 12
+        }
+      ]
+    }
+  ],
+  "089/000 Adding Materialized Stored Query: LET X \u003c= SELECT value AS Foo FROM range(start=1, end=4)": null,
+  "089/001 Adding Materialized Stored Query: SELECT (dict(Foo=12), ) + X, X + (dict(Foo=12), ) FROM scope()": [
+    {
+      "(dict(Foo=12), ) + X": [
+        {
+          "Foo": 12
+        },
+        {
+          "Foo": 1
+        },
+        {
+          "Foo": 2
+        },
+        {
+          "Foo": 3
+        },
+        {
+          "Foo": 4
+        }
+      ],
+      "X + (dict(Foo=12), )": [
+        {
+          "Foo": 1
+        },
+        {
+          "Foo": 2
+        },
+        {
+          "Foo": 3
+        },
+        {
+          "Foo": 4
+        },
+        {
+          "Foo": 12
+        }
+      ]
+    }
   ]
 }

protocols/protocol_add.go

@@ -1,6 +1,7 @@
 package protocols
 
 import (
+	"context"
 	"reflect"
 
 	"www.velocidex.com/golang/vfilter/types"
@@ -94,8 +95,8 @@ func (self AddDispatcher) Add(scope types.Scope, a types.Any, b types.Any) types
 
 	// Handle array concatenation
 	if is_array(a) || is_array(b) {
-		a_slice := convertToSlice(a)
-		b_slice := convertToSlice(b)
+		a_slice := convertToSlice(scope, a)
+		b_slice := convertToSlice(scope, b)
 
 		return append(a_slice, b_slice...)
 	}
@@ -111,7 +112,7 @@ func (self *AddDispatcher) AddImpl(elements ...AddProtocol) {
 	}
 }
 
-func convertToSlice(a types.Any) []types.Any {
+func convertToSlice(scope types.Scope, a types.Any) []types.Any {
 	if is_array(a) {
 		a_slice := reflect.ValueOf(a)
 		result := make([]types.Any, 0, a_slice.Len())
@@ -120,5 +121,15 @@ func convertToSlice(a types.Any) []types.Any {
 		}
 		return result
 	}
+
+	stored_query, ok := a.(types.StoredQuery)
+	if ok {
+		var res []types.Any
+		for _, x := range MaterializeToArray(context.Background(), scope, stored_query) {
+			res = append(res, x)
+		}
+		return res
+	}
+
 	return []types.Any{a}
 }

vfilter_test.go

@@ -1334,6 +1334,17 @@ FROM scope()
 	{"Test struct associative", `
 SELECT StructValue.SrcIP, StructValue.src_ip, StructValue.SrcIp
 FROM scope()`},
+
+	{"Adding Stored Query", `
+LET X = SELECT value AS Foo FROM range(start=1, end=4)
+
+SELECT (dict(Foo=12), ) + X, X + (dict(Foo=12),) FROM scope()
+`},
+	{"Adding Materialized Stored Query", `
+LET X <= SELECT value AS Foo FROM range(start=1, end=4)
+
+SELECT (dict(Foo=12), ) + X, X + (dict(Foo=12),) FROM scope()
+`},
 }
 
 type _RangeArgs struct {
@@ -1533,7 +1544,7 @@ func TestMultiVQLQueries(t *testing.T) {
 	// Store the result in ordered dict so we have a consistent golden file.
 	result := ordereddict.NewDict()
 	for i, testCase := range multiVQLTest {
-		if false && i != 85 {
+		if false && i != 88 {
 			continue
 		}
 		scope := makeTestScope()