From d350ca32cc65edf8510952b3ae522ba12e2d506a Mon Sep 17 00:00:00 2001 From: cpzt Date: Thu, 14 Oct 2021 15:58:07 +0800 Subject: [PATCH] add field from result set in mongodb --- sql/engines/mongo.py | 12 ++++++++++++ sql/engines/tests.py | 8 ++++++++ 2 files changed, 20 insertions(+) diff --git a/sql/engines/mongo.py b/sql/engines/mongo.py index 49ca125b21..a47e324c4b 100644 --- a/sql/engines/mongo.py +++ b/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 diff --git a/sql/engines/tests.py b/sql/engines/tests.py index 4669780a65..0ea7b365aa 100644 --- a/sql/engines/tests.py +++ b/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"])