add field from result set in mongodb

sql/engines/mongo.py

@@ -825,6 +825,8 @@ def parse_tuple(self, cursor, db_name, tb_name, projection=None):
             result = self.get_all_columns_by_tb(db_name=db_name, tb_name=tb_name)
             columns = result.rows
         columns.insert(0, "mongodballdata")  # 隐藏JSON结果列
+        columns = self.fill_query_columns(cursor, columns)
+
         for ro in cursor:
             json_col = json.dumps(ro, ensure_ascii=False, indent=2, separators=(",", ":"))
             row.insert(0, json_col)
@@ -851,3 +853,13 @@ def parse_tuple(self, cursor, db_name, tb_name, projection=None):
             rows.append(tuple(row))
             row.clear()
         return tuple(rows), columns
+
+    @staticmethod
+    def fill_query_columns(cursor, columns):
+        """补充结果集中`get_all_columns_by_tb`未获取的字段"""
+        cols = columns
+        for ro in cursor:
+            for key in ro.keys():
+                if key not in cols:
+                    cols.append(key)
+        return cols

sql/engines/tests.py

@@ -1653,3 +1653,11 @@ def test_execute(self, mock_get_master, mock_exec_cmd):
         check_result = self.engine.execute("some_db", sql)
         mock_get_master.assert_called_once()
         self.assertEqual(check_result.rows[0].__dict__["errlevel"], 0)
+
+    def test_fill_query_columns(self):
+        columns = ["_id", "title", "tags", "likes"]
+        cursor = [{"_id": {"$oid": "5f10162029684728e70045ab"}, "title": "MongoDB", "text": "archery", "likes": 100},
+                  {"_id": {"$oid": "7f10162029684728e70045ab"}, "author": "archery"}]
+        cols = self.engine.fill_query_columns(cursor, columns=columns)
+        print(cols)
+        self.assertEqual(cols, ["_id", "title", "tags", "likes", "text", "author"])