From dcfb3454c1974cc856b25b5f7b7d58d2415e1978 Mon Sep 17 00:00:00 2001
From: Dmitry Shmulevich <dmitsh@users.noreply.github.com>
Date: Tue, 15 Mar 2022 19:05:53 -0700
Subject: [PATCH] mongodb: fix querying non-string values

Signed-off-by: Dmitry Shmulevich <dmitsh@users.noreply.github.com>
---
 state/mongodb/mongodb_query.go      | 22 +++++++++++++----
 state/mongodb/mongodb_query_test.go |  4 ++++
 tests/state/query/q6.json           | 37 +++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 4 deletions(-)
 create mode 100644 tests/state/query/q6.json

diff --git a/state/mongodb/mongodb_query.go b/state/mongodb/mongodb_query.go
index b6d5d2269..7df4d2d1c 100644
--- a/state/mongodb/mongodb_query.go
+++ b/state/mongodb/mongodb_query.go
@@ -37,7 +37,12 @@ type Query struct {
 
 func (q *Query) VisitEQ(f *query.EQ) (string, error) {
 	// { <key>: <val> }
-	return fmt.Sprintf(`{ "value.%s": %q }`, f.Key, f.Val), nil
+	switch v := f.Val.(type) {
+	case string:
+		return fmt.Sprintf(`{ "value.%s": %q }`, f.Key, v), nil
+	default:
+		return fmt.Sprintf(`{ "value.%s": %v }`, f.Key, v), nil
+	}
 }
 
 func (q *Query) VisitIN(f *query.IN) (string, error) {
@@ -45,9 +50,18 @@ func (q *Query) VisitIN(f *query.IN) (string, error) {
 	if len(f.Vals) == 0 {
 		return "", fmt.Errorf("empty IN operator for key %q", f.Key)
 	}
-	str := fmt.Sprintf(`{ "value.%s": { "$in": [ %q`, f.Key, f.Vals[0])
-	for _, v := range f.Vals[1:] {
-		str += fmt.Sprintf(", %q", v)
+	str := fmt.Sprintf(`{ "value.%s": { "$in": [ `, f.Key)
+
+	for i := 0; i < len(f.Vals); i++ {
+		if i > 0 {
+			str += ", "
+		}
+		switch v := f.Vals[i].(type) {
+		case string:
+			str += fmt.Sprintf("%q", v)
+		default:
+			str += fmt.Sprintf("%v", v)
+		}
 	}
 	str += " ] } }"
 
diff --git a/state/mongodb/mongodb_query_test.go b/state/mongodb/mongodb_query_test.go
index 1a6e7f3e7..e5b985efc 100644
--- a/state/mongodb/mongodb_query_test.go
+++ b/state/mongodb/mongodb_query_test.go
@@ -44,6 +44,10 @@ func TestMongoQuery(t *testing.T) {
 			input: "../../tests/state/query/q4.json",
 			query: `{ "$or": [ { "value.person.org": "A" }, { "$and": [ { "value.person.org": "B" }, { "value.state": { "$in": [ "CA", "WA" ] } } ] } ] }`,
 		},
+		{
+			input: "../../tests/state/query/q6.json",
+			query: `{ "$or": [ { "value.person.id": 123 }, { "$and": [ { "value.person.org": "B" }, { "value.person.id": { "$in": [ 567, 890 ] } } ] } ] }`,
+		},
 	}
 	for _, test := range tests {
 		data, err := ioutil.ReadFile(test.input)
diff --git a/tests/state/query/q6.json b/tests/state/query/q6.json
new file mode 100644
index 000000000..da53b8b10
--- /dev/null
+++ b/tests/state/query/q6.json
@@ -0,0 +1,37 @@
+{
+    "filter": {
+        "OR": [
+            {
+                "EQ": {
+                    "person.id": 123
+                }
+            },
+            {
+                "AND": [
+                    {
+                        "EQ": {
+                            "person.org": "B"
+                        }
+                    },
+                    {
+                        "IN": {
+                            "person.id": [567, 890]
+                        }
+                    }
+                ]
+            }
+        ]
+    },
+    "sort": [
+        {
+            "key": "state",
+            "order": "DESC"
+        },
+        {
+            "key": "person.name"
+        }
+    ],
+    "page": {
+        "limit": 2
+    }
+}