Skip to content

Commit 1a90ad5

Browse files
fixpngfixpng
authored
Fix: mongodb 用不了count统计 (#2991)
1 parent c647f07 commit 1a90ad5

File tree

2 files changed

+261
-1
lines changed

2 files changed

+261
-1
lines changed

‎sql/engines/mongo.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1091,7 +1091,10 @@ def query(self, db_name=None, sql="", limit_num=0, close_conn=True, **kwargs):
10911091
query_skip = int(query_dict["skip"])
10921092
find_cmd += f".skip({query_skip})"
10931093
if "count" in query_dict:
1094-
find_cmd += ".count()"
1094+
if condition:
1095+
find_cmd = "collection.count_documents(condition)"
1096+
else:
1097+
find_cmd = "collection.count_documents({})"
10951098
if "explain" in query_dict:
10961099
find_cmd += ".explain()"
10971100

‎sql/engines/test_mongo.py

Lines changed: 257 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from sql.engines.mongo import MongoEngine
44

55

6+
# 创建MongoEngine测试实例的fixture
67
@pytest.fixture
78
def mongo_engine():
89
engine = MongoEngine()
@@ -15,6 +16,7 @@ def mongo_engine():
1516
return engine
1617

1718

19+
# 测试带load参数的命令生成
1820
def test_build_cmd_with_load(mongo_engine):
1921
# Call the method with is_load=True
2022
cmd = mongo_engine._build_cmd(
@@ -35,6 +37,7 @@ def test_build_cmd_with_load(mongo_engine):
3537
assert cmd == expected_cmd
3638

3739

40+
# 测试不带load参数的命令生成
3841
def test_build_cmd_without_load(mongo_engine):
3942
# Call the method with is_load=False
4043
cmd = mongo_engine._build_cmd(
@@ -55,6 +58,7 @@ def test_build_cmd_without_load(mongo_engine):
5558
assert cmd == expected_cmd
5659

5760

61+
# 测试无认证信息时命令生成
5862
def test_build_cmd_without_auth(mongo_engine):
5963
# Set user and password to None
6064
mongo_engine.user = None
@@ -79,6 +83,7 @@ def test_build_cmd_without_auth(mongo_engine):
7983
assert cmd == expected_cmd
8084

8185

86+
# 测试无认证信息且带load参数的命令生成
8287
def test_build_cmd_with_load_without_auth(mongo_engine):
8388
# Set user and password to None
8489
mongo_engine.user = None
@@ -101,3 +106,255 @@ def test_build_cmd_with_load_without_auth(mongo_engine):
101106

102107
# Assertions
103108
assert cmd == expected_cmd
109+
110+
111+
# 参数化测试多种Mongo命令的生成
112+
@pytest.mark.parametrize(
113+
"params,expected_cmd",
114+
[
115+
# 基础 find 查询
116+
(
117+
dict(
118+
db_name="test_db",
119+
auth_db="admin",
120+
slave_ok="",
121+
sql="db.test_collection.find({})",
122+
is_load=False,
123+
),
124+
"mongo --quiet -u test_user -p 'test_password' localhost:27017/admin <<\\EOF\n"
125+
"db=db.getSiblingDB('test_db');db.test_collection.find({})\nEOF",
126+
),
127+
# find 带条件
128+
(
129+
dict(
130+
db_name="test_db",
131+
auth_db="admin",
132+
slave_ok="",
133+
sql="db.test_collection.find({'name':'archery'})",
134+
is_load=False,
135+
),
136+
"mongo --quiet -u test_user -p 'test_password' localhost:27017/admin <<\\EOF\n"
137+
"db=db.getSiblingDB('test_db');db.test_collection.find({'name':'archery'})\nEOF",
138+
),
139+
# aggregate 查询
140+
(
141+
dict(
142+
db_name="test_db",
143+
auth_db="admin",
144+
slave_ok="",
145+
sql="db.test_collection.aggregate([{'$match':{'age':{'$gt':18}}}])",
146+
is_load=False,
147+
),
148+
"mongo --quiet -u test_user -p 'test_password' localhost:27017/admin <<\\EOF\n"
149+
"db=db.getSiblingDB('test_db');db.test_collection.aggregate([{'$match':{'age':{'$gt':18}}}])\nEOF",
150+
),
151+
# count 查询
152+
(
153+
dict(
154+
db_name="test_db",
155+
auth_db="admin",
156+
slave_ok="",
157+
sql="db.test_collection.count({'status':'active'})",
158+
is_load=False,
159+
),
160+
"mongo --quiet -u test_user -p 'test_password' localhost:27017/admin <<\\EOF\n"
161+
"db=db.getSiblingDB('test_db');db.test_collection.count({'status':'active'})\nEOF",
162+
),
163+
# explain 查询
164+
(
165+
dict(
166+
db_name="test_db",
167+
auth_db="admin",
168+
slave_ok="",
169+
sql="db.test_collection.find({'score':{'$gte':90}}).explain()",
170+
is_load=False,
171+
),
172+
"mongo --quiet -u test_user -p 'test_password' localhost:27017/admin <<\\EOF\n"
173+
"db=db.getSiblingDB('test_db');db.test_collection.find({'score':{'$gte':90}}).explain()\nEOF",
174+
),
175+
# find 带 sort/limit/skip
176+
(
177+
dict(
178+
db_name="test_db",
179+
auth_db="admin",
180+
slave_ok="",
181+
sql="db.test_collection.find({}).sort({'age':-1}).limit(10).skip(5)",
182+
is_load=False,
183+
),
184+
"mongo --quiet -u test_user -p 'test_password' localhost:27017/admin <<\\EOF\n"
185+
"db=db.getSiblingDB('test_db');db.test_collection.find({}).sort({'age':-1}).limit(10).skip(5)\nEOF",
186+
),
187+
# findOne 查询
188+
(
189+
dict(
190+
db_name="test_db",
191+
auth_db="admin",
192+
slave_ok="",
193+
sql="db.test_collection.findOne({'_id':ObjectId('507f1f77bcf86cd799439011')})",
194+
is_load=False,
195+
),
196+
"mongo --quiet -u test_user -p 'test_password' localhost:27017/admin <<\\EOF\n"
197+
"db=db.getSiblingDB('test_db');db.test_collection.findOne({'_id':ObjectId('507f1f77bcf86cd799439011')})\nEOF",
198+
),
199+
# insertOne
200+
(
201+
dict(
202+
db_name="test_db",
203+
auth_db="admin",
204+
slave_ok="",
205+
sql="db.test_collection.insertOne({'name':'archery','age':20})",
206+
is_load=False,
207+
),
208+
"mongo --quiet -u test_user -p 'test_password' localhost:27017/admin <<\\EOF\n"
209+
"db=db.getSiblingDB('test_db');db.test_collection.insertOne({'name':'archery','age':20})\nEOF",
210+
),
211+
# updateMany
212+
(
213+
dict(
214+
db_name="test_db",
215+
auth_db="admin",
216+
slave_ok="",
217+
sql="db.test_collection.updateMany({'status':'active'},{'$set':{'score':100}})",
218+
is_load=False,
219+
),
220+
"mongo --quiet -u test_user -p 'test_password' localhost:27017/admin <<\\EOF\n"
221+
"db=db.getSiblingDB('test_db');db.test_collection.updateMany({'status':'active'},{'$set':{'score':100}})\nEOF",
222+
),
223+
# deleteMany
224+
(
225+
dict(
226+
db_name="test_db",
227+
auth_db="admin",
228+
slave_ok="",
229+
sql="db.test_collection.deleteMany({'expired':true})",
230+
is_load=False,
231+
),
232+
"mongo --quiet -u test_user -p 'test_password' localhost:27017/admin <<\\EOF\n"
233+
"db=db.getSiblingDB('test_db');db.test_collection.deleteMany({'expired':true})\nEOF",
234+
),
235+
# createIndex
236+
(
237+
dict(
238+
db_name="test_db",
239+
auth_db="admin",
240+
slave_ok="",
241+
sql="db.test_collection.createIndex({'name':1},{'background':true})",
242+
is_load=False,
243+
),
244+
"mongo --quiet -u test_user -p 'test_password' localhost:27017/admin <<\\EOF\n"
245+
"db=db.getSiblingDB('test_db');db.test_collection.createIndex({'name':1},{'background':true})\nEOF",
246+
),
247+
# dropIndex
248+
(
249+
dict(
250+
db_name="test_db",
251+
auth_db="admin",
252+
slave_ok="",
253+
sql="db.test_collection.dropIndex('name_1')",
254+
is_load=False,
255+
),
256+
"mongo --quiet -u test_user -p 'test_password' localhost:27017/admin <<\\EOF\n"
257+
"db=db.getSiblingDB('test_db');db.test_collection.dropIndex('name_1')\nEOF",
258+
),
259+
# createCollection
260+
(
261+
dict(
262+
db_name="test_db",
263+
auth_db="admin",
264+
slave_ok="",
265+
sql="db.createCollection('new_collection')",
266+
is_load=False,
267+
),
268+
"mongo --quiet -u test_user -p 'test_password' localhost:27017/admin <<\\EOF\n"
269+
"db=db.getSiblingDB('test_db');db.createCollection('new_collection')\nEOF",
270+
),
271+
],
272+
)
273+
def test_build_cmd_various_queries(mongo_engine, params, expected_cmd):
274+
# 测试多种Mongo命令的生成
275+
cmd = mongo_engine._build_cmd(**params)
276+
assert cmd == expected_cmd
277+
278+
279+
# 测试带slave_ok参数的命令生成
280+
def test_build_cmd_with_slave_ok(mongo_engine):
281+
cmd = mongo_engine._build_cmd(
282+
db_name="test_db",
283+
auth_db="admin",
284+
slave_ok="rs.slaveOk();",
285+
sql="db.test_collection.find()",
286+
is_load=False,
287+
)
288+
expected_cmd = (
289+
"mongo --quiet -u test_user -p 'test_password' localhost:27017/admin <<\\EOF\n"
290+
"db=db.getSiblingDB('test_db');rs.slaveOk();db.test_collection.find()\nEOF"
291+
)
292+
assert cmd == expected_cmd
293+
294+
295+
# 测试db_name为空时命令生成
296+
def test_build_cmd_with_empty_db_name(mongo_engine):
297+
cmd = mongo_engine._build_cmd(
298+
db_name="",
299+
auth_db="admin",
300+
slave_ok="",
301+
sql="db.test_collection.find()",
302+
is_load=False,
303+
)
304+
expected_cmd = (
305+
"mongo --quiet -u test_user -p 'test_password' localhost:27017/admin <<\\EOF\n"
306+
"db=db.getSiblingDB('');db.test_collection.find()\nEOF"
307+
)
308+
assert cmd == expected_cmd
309+
310+
311+
# 测试user和password为None时命令生成
312+
def test_build_cmd_with_none_user_password(mongo_engine):
313+
mongo_engine.user = None
314+
mongo_engine.password = None
315+
cmd = mongo_engine._build_cmd(
316+
db_name="test_db",
317+
auth_db="admin",
318+
slave_ok="",
319+
sql="db.test_collection.find()",
320+
is_load=False,
321+
)
322+
expected_cmd = (
323+
"mongo --quiet localhost:27017/admin <<\\EOF\n"
324+
"db=db.getSiblingDB('test_db');db.test_collection.find()\nEOF"
325+
)
326+
assert cmd == expected_cmd
327+
328+
329+
# 测试带load和slave_ok参数的命令生成
330+
def test_build_cmd_with_load_and_slave_ok(mongo_engine):
331+
cmd = mongo_engine._build_cmd(
332+
db_name="test_db",
333+
auth_db="admin",
334+
slave_ok="rs.slaveOk();",
335+
tempfile_="/tmp/test.js",
336+
is_load=True,
337+
)
338+
expected_cmd = (
339+
"mongo --quiet -u test_user -p 'test_password' localhost:27017/admin <<\\EOF\n"
340+
"db=db.getSiblingDB('test_db');rs.slaveOk();load('/tmp/test.js')\nEOF"
341+
)
342+
assert cmd == expected_cmd
343+
344+
345+
# 测试无认证信息且带load和slave_ok参数的命令生成
346+
def test_build_cmd_with_load_without_auth_and_slave_ok(mongo_engine):
347+
mongo_engine.user = None
348+
mongo_engine.password = None
349+
cmd = mongo_engine._build_cmd(
350+
db_name="test_db",
351+
auth_db="admin",
352+
slave_ok="rs.slaveOk();",
353+
tempfile_="/tmp/test.js",
354+
is_load=True,
355+
)
356+
expected_cmd = (
357+
"mongo --quiet localhost:27017/admin <<\\EOF\n"
358+
"db=db.getSiblingDB('test_db');rs.slaveOk();load('/tmp/test.js')\nEOF"
359+
)
360+
assert cmd == expected_cmd

0 commit comments

Comments
 (0)