Close getredash#1199: support for nested fields in MongoDB results

.gitignore

@@ -1,3 +1,4 @@
+.venv
 .cache
 .coverage.*
 .coveralls.yml

redash/query_runner/mongodb.py

@@ -74,6 +74,49 @@ def parse_query_json(query):
     return query_data
 
 
+def _get_column_by_name(columns, column_name):
+    for c in columns:
+        if "name" in c and c["name"] == column_name:
+            return c
+
+    return None
+
+
+def parse_results(results):
+    rows = []
+    columns = []
+
+    for row in results:
+        parsed_row = {}
+
+        for key in row:
+            if isinstance(row[key], dict):
+                for inner_key in row[key]:
+                    column_name = '{}.{}'.format(key, inner_key)
+                    if _get_column_by_name(columns, column_name) is None:
+                        columns.append({
+                            "name": column_name,
+                            "friendly_name": column_name,
+                            "type": TYPES_MAP.get(type(row[key][inner_key]), TYPE_STRING)
+                        })
+
+                    parsed_row[column_name] = row[key][inner_key]
+
+            else:
+                if _get_column_by_name(columns, key) is None:
+                    columns.append({
+                        "name": key,
+                        "friendly_name": key,
+                        "type": TYPES_MAP.get(type(row[key]), TYPE_STRING)
+                    })
+
+                parsed_row[key] = row[key]
+
+        rows.append(parsed_row)
+
+    return rows, columns
+
+
 class MongoDB(BaseQueryRunner):
     @classmethod
     def configuration_schema(cls):
@@ -113,13 +156,6 @@ def __init__(self, configuration):
 
         self.is_replica_set = True if "replicaSetName" in self.configuration and self.configuration["replicaSetName"] else False
 
-    def _get_column_by_name(self, columns, column_name):
-        for c in columns:
-            if "name" in c and c["name"] == column_name:
-                return c
-
-        return None
-
     def _get_db(self):
         if self.is_replica_set:
             db_connection = pymongo.MongoReplicaSetClient(self.configuration["connectionString"], replicaSet=self.configuration["replicaSetName"])
@@ -259,21 +295,14 @@ def run_query(self, query, user):
 
             rows.append({ "count" : cursor })
         else:
-            for r in cursor:
-                for k in r:
-                    if self._get_column_by_name(columns, k) is None:
-                        columns.append({
-                            "name": k,
-                            "friendly_name": k,
-                            "type": TYPES_MAP.get(type(r[k]), TYPE_STRING)
-                        })
-
-                rows.append(r)
+            rows, columns = parse_results(cursor)
 
         if f:
             ordered_columns = []
             for k in sorted(f, key=f.get):
-                ordered_columns.append(self._get_column_by_name(columns, k))
+                column = _get_column_by_name(columns, k)
+                if column:
+                    ordered_columns.append(column)
 
             columns = ordered_columns
 

tests/query_runner/test_mongodb.py

@@ -2,7 +2,7 @@
 import json
 from unittest import TestCase
 from pytz import utc
-from redash.query_runner.mongodb import parse_query_json
+from redash.query_runner.mongodb import parse_query_json, parse_results, _get_column_by_name
 
 from redash.utils import parse_human_time
 
@@ -105,4 +105,42 @@ def test_supports_relative_timestamps(self):
         self.assertEqual(query_data['ts'], one_hour_ago)
 
 
-
+class TestMongoResults(TestCase):
+    def test_parses_regular_results(self):
+        raw_results = [
+            {'column': 1, 'column2': 'test'},
+            {'column': 2, 'column2': 'test', 'column3': 'hello'}
+        ]
+        rows, columns = parse_results(raw_results)
+
+        for i, row in enumerate(rows):
+            self.assertDictEqual(row, raw_results[i])
+
+        self.assertIsNotNone(_get_column_by_name(columns, 'column'))
+        self.assertIsNotNone(_get_column_by_name(columns, 'column2'))
+        self.assertIsNotNone(_get_column_by_name(columns, 'column3'))
+
+    def test_parses_nested_results(self):
+        raw_results = [
+            {'column': 1, 'column2': 'test', 'nested': {
+                'a': 1,
+                'b': 'str'
+            }},
+            {'column': 2, 'column2': 'test', 'column3': 'hello', 'nested': {
+                'a': 2,
+                'b': 'str2',
+                'c': 'c'
+            }}
+        ]
+
+        rows, columns = parse_results(raw_results)
+
+        self.assertDictEqual(rows[0], { 'column': 1, 'column2': 'test', 'nested.a': 1, 'nested.b': 'str' })
+        self.assertDictEqual(rows[1], { 'column': 2, 'column2': 'test', 'column3': 'hello', 'nested.a': 2, 'nested.b': 'str2', 'nested.c': 'c' })
+
+        self.assertIsNotNone(_get_column_by_name(columns, 'column'))
+        self.assertIsNotNone(_get_column_by_name(columns, 'column2'))
+        self.assertIsNotNone(_get_column_by_name(columns, 'column3'))
+        self.assertIsNotNone(_get_column_by_name(columns, 'nested.a'))
+        self.assertIsNotNone(_get_column_by_name(columns, 'nested.b'))
+        self.assertIsNotNone(_get_column_by_name(columns, 'nested.c'))