Skip to content

Commit 3d9cd25

Browse files
Boris-codeBoris-code
authored
Merge pull request Boris-code#246 from chang-xiao-feng/master
1、MongoDB 批量更新代码。
2 parents 033feba + da3adae commit 3d9cd25

2 files changed

Lines changed: 106 additions & 25 deletions

File tree

‎feapder/db/mongodb.py‎

Lines changed: 103 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
from urllib import parse
1313

1414
import pymongo
15-
from pymongo import MongoClient
15+
from pymongo import MongoClient, UpdateOne
1616
from pymongo.collection import Collection
1717
from pymongo.database import Database
1818
from pymongo.errors import DuplicateKeyError, BulkWriteError
@@ -23,14 +23,14 @@
2323

2424
class MongoDB:
2525
def __init__(
26-
self,
27-
ip=None,
28-
port=None,
29-
db=None,
30-
user_name=None,
31-
user_pass=None,
32-
url=None,
33-
**kwargs,
26+
self,
27+
ip=None,
28+
port=None,
29+
db=None,
30+
user_name=None,
31+
user_pass=None,
32+
url=None,
33+
**kwargs,
3434
):
3535
if url:
3636
self.client = MongoClient(url, **kwargs)
@@ -94,7 +94,7 @@ def get_collection(self, coll_name, **kwargs) -> Collection:
9494
return self.db.get_collection(coll_name, **kwargs)
9595

9696
def find(
97-
self, coll_name: str, condition: Optional[Dict] = None, limit: int = 0, **kwargs
97+
self, coll_name: str, condition: Optional[Dict] = None, limit: int = 0, **kwargs
9898
) -> List[Dict]:
9999
"""
100100
@summary:
@@ -133,13 +133,13 @@ def find(
133133
return dataset
134134

135135
def add(
136-
self,
137-
coll_name,
138-
data: Dict,
139-
replace=False,
140-
update_columns=(),
141-
update_columns_value=(),
142-
insert_ignore=False,
136+
self,
137+
coll_name,
138+
data: Dict,
139+
replace=False,
140+
update_columns=(),
141+
update_columns_value=(),
142+
insert_ignore=False,
143143
):
144144
"""
145145
添加单条数据
@@ -195,13 +195,13 @@ def add(
195195
return affect_count
196196

197197
def add_batch(
198-
self,
199-
coll_name: str,
200-
datas: List[Dict],
201-
replace=False,
202-
update_columns=(),
203-
update_columns_value=(),
204-
condition_fields: dict = None,
198+
self,
199+
coll_name: str,
200+
datas: List[Dict],
201+
replace=False,
202+
update_columns=(),
203+
update_columns_value=(),
204+
condition_fields: dict = None,
205205
):
206206
"""
207207
批量添加数据
@@ -331,6 +331,70 @@ def update(self, coll_name, data: Dict, condition: Dict, upsert: bool = False):
331331
else:
332332
return True
333333

334+
def update_many(self, coll_name, data: Dict, condition: Dict, upsert: bool = False):
335+
"""
336+
批量更新
337+
Args:
338+
coll_name: 集合名
339+
data: 单条数据 {"xxx":"xxx"}
340+
condition: 更新条件 {"_id": "xxxx"}
341+
upsert: 数据不存在则插入,默认为 False
342+
343+
Returns: True / False
344+
"""
345+
try:
346+
collection = self.get_collection(coll_name)
347+
collection.update_many(condition, {"$set": data}, upsert=upsert)
348+
except Exception as e:
349+
log.error(
350+
"""
351+
error:{}
352+
condition: {}
353+
""".format(
354+
e, condition
355+
)
356+
)
357+
return False
358+
else:
359+
return True
360+
361+
def update_batch(
362+
self,
363+
coll_name: str,
364+
update_data_list: List[Dict],
365+
condition_field: str,
366+
upsert: bool = False,
367+
):
368+
"""
369+
批量更新数据
370+
Args:
371+
coll_name: 集合名
372+
update_data_list: 更新数据列表
373+
condition_field: 更新条件字段
374+
upsert: 数据不存在则插入，默认为 False
375+
376+
Returns: 更新行数
377+
378+
"""
379+
if not update_data_list:
380+
return 0
381+
382+
collection = self.get_collection(coll_name)
383+
bulk_operations = []
384+
385+
for update_data in update_data_list:
386+
condition = {condition_field: update_data.get(condition_field)}
387+
update_operation = UpdateOne(
388+
condition, {"$set": update_data}, upsert=upsert
389+
)
390+
bulk_operations.append(update_operation)
391+
try:
392+
result = collection.bulk_write(bulk_operations, ordered=False)
393+
return result.modified_count + result.upserted_count
394+
except BulkWriteError as e:
395+
log.error(f"Bulk write error: {e.details}")
396+
return 0
397+
334398
def delete(self, coll_name, condition: Dict) -> bool:
335399
"""
336400
删除
@@ -401,7 +465,7 @@ def get_index_key(self, coll_name, index_name):
401465
return index_keys
402466

403467
def __get_update_condition(
404-
self, coll_name: str, data: dict, duplicate_errmsg: str
468+
self, coll_name: str, data: dict, duplicate_errmsg: str
405469
) -> dict:
406470
"""
407471
根据索引冲突的报错信息 获取更新条件
@@ -420,3 +484,17 @@ def __get_update_condition(
420484

421485
def __getattr__(self, name):
422486
return getattr(self.db, name)
487+
488+
489+
if __name__ == '__main__':
490+
update_data_list = [
491+
{"_id": "1", "status": 1},
492+
{"_id": "2", "status": 1}]
493+
mongo = MongoDB()
494+
updated_count = mongo.update_batch("your_table_name", update_data_list, "_id")
495+
print(f"Updated {updated_count} documents.")
496+
497+
id_list = ['1', '2']
498+
result = mongo.update_many("your_table_name",
499+
{"status": 1},
500+
{"_id": {"$in": id_list}})

‎feapder/db/redisdb.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -744,6 +744,9 @@ def hget_count(self, table):
744744
def hkeys(self, table):
745745
return self._redis.hkeys(table)
746746

747+
def hvals(self, key):
748+
return self._redis.hvals(key)
749+
747750
def setbit(
748751
self, table, offsets: Union[int, List[int]], values: Union[int, List[int]]
749752
):

0 commit comments

Comments
 (0)