11from pymongo import MongoClient
2- from pymongo import MongoClient
2+ from pymongo . errors import DuplicateKeyError
33import time
44from bson .objectid import ObjectId
55from bson .errors import InvalidId
66import os
77import xml .etree .ElementTree as ET
8- import random
8+ import secrets
99import string
1010from dotenv import load_dotenv
1111load_dotenv ()
@@ -14,42 +14,42 @@ class WorkFlowModel:
1414 def __init__ (self ) -> None :
1515 self .collection = MongoClient (os .getenv ('MongoURL' ))[
1616 os .getenv ('dbName' )][os .getenv ('tableName' )]
17+ self .collection .create_index ('serverID' , unique = True )
1718
1819 def get_random_string (self , length ):
1920 letters = string .ascii_letters + string .digits
20- return '' .join (random .choice (letters ) for i in range (length ))
21+ return '' .join (secrets .choice (letters ) for i in range (length ))
2122
2223 def insert (self , graphml , latestHash ):
23- serverID = ""
2424 while (True ):
2525 serverID = self .get_random_string (6 )
26- if (not self .collection .find_one ({'serverID' : serverID })):
27- break
28- self .collection .insert_one (
29- {'graphml' : graphml , 'latestHash' : latestHash , 'serverID' : serverID })
30- return serverID
26+ try :
27+ self .collection .insert_one (
28+ {'graphml' : graphml , 'latestHash' : latestHash , 'serverID' : serverID })
29+ return serverID
30+ except DuplicateKeyError :
31+ continue
3132
3233 def get (self , serverID ):
3334 cl = self .collection .find_one ({'serverID' : serverID })
3435 if not cl :
35- return False , 'Record Not Found'
36+ return None
3637 return cl ['graphml' ]
3738
3839 def update (self , serverID , graphml , latestHash , allHash ):
39- existingRecord = self .collection .find_one ({'serverID' : serverID })
40+ existingRecord = self .collection .find_one_and_update (
41+ {'serverID' : serverID , 'latestHash' : {'$in' : allHash }},
42+ {"$set" : {'graphml' : graphml , 'latestHash' : latestHash }})
4043 if existingRecord is None :
41- return False , 'serverID do not exists.'
42- latestExistingHash = existingRecord ['latestHash' ]
43- if latestExistingHash not in allHash :
44+ if not self .collection .find_one ({'serverID' : serverID }):
45+ return False , 'serverID do not exists.'
4446 return False , 'Can not update as provided graph do not has latest changes.'
45- self .collection .update_one ({'serverID' : serverID }, {
46- "$set" : {'graphml' : graphml , 'latestHash' : latestHash }})
4747 return True , latestHash
4848
4949 def forceUpdate (self , serverID , graphml , latestHash ):
50- existingRecord = self .collection .find_one ({'serverID' : serverID })
50+ existingRecord = self .collection .find_one_and_update (
51+ {'serverID' : serverID },
52+ {"$set" : {'graphml' : graphml , 'latestHash' : latestHash }})
5153 if existingRecord is None :
5254 return False , 'serverID do not exists.'
53- self .collection .update_one ({'serverID' : serverID },
54- {"$set" : {'graphml' : graphml , 'latestHash' : latestHash }})
5555 return True , latestHash
0 commit comments