|
|
@ -64,14 +64,14 @@ class DB():
|
|
|
|
check_same_thread=False
|
|
|
|
check_same_thread=False
|
|
|
|
)
|
|
|
|
)
|
|
|
|
self.db.row_factory = sqlite3.Row
|
|
|
|
self.db.row_factory = sqlite3.Row
|
|
|
|
|
|
|
|
|
|
|
|
self.dblock = Lock()
|
|
|
|
|
|
|
|
atexit.register(lambda db : db.close(), self.db)
|
|
|
|
atexit.register(lambda db : db.close(), self.db)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
self.db_lock = Lock()
|
|
|
|
|
|
|
|
|
|
|
|
self._assure_tables()
|
|
|
|
self._assure_tables()
|
|
|
|
|
|
|
|
|
|
|
|
def _assure_tables(self):
|
|
|
|
def _assure_tables(self):
|
|
|
|
self.dblock.acquire()
|
|
|
|
self.db_lock.acquire()
|
|
|
|
with self.db:
|
|
|
|
with self.db:
|
|
|
|
self.db.execute("""CREATE TABLE IF NOT EXISTS Messages (
|
|
|
|
self.db.execute("""CREATE TABLE IF NOT EXISTS Messages (
|
|
|
|
count INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
|
|
count INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
|
@ -82,11 +82,11 @@ class DB():
|
|
|
|
json BLOB,
|
|
|
|
json BLOB,
|
|
|
|
processed BOOL DEFAULT FALSE
|
|
|
|
processed BOOL DEFAULT FALSE
|
|
|
|
)""")
|
|
|
|
)""")
|
|
|
|
self.dblock.release()
|
|
|
|
self.db_lock.release()
|
|
|
|
|
|
|
|
|
|
|
|
@validate_call
|
|
|
|
@validate_call
|
|
|
|
def add_message(self, sender:str, recipient:str, message:AgentMessage, processed:bool=False) -> int:
|
|
|
|
def add_message(self, sender:str, recipient:str, message:AgentMessage, processed:bool=False) -> int:
|
|
|
|
self.dblock.acquire()
|
|
|
|
self.db_lock.acquire()
|
|
|
|
with self.db:
|
|
|
|
with self.db:
|
|
|
|
self.db.execute(
|
|
|
|
self.db.execute(
|
|
|
|
"""INSERT INTO Messages (
|
|
|
|
"""INSERT INTO Messages (
|
|
|
@ -101,20 +101,21 @@ class DB():
|
|
|
|
"processed" : processed
|
|
|
|
"processed" : processed
|
|
|
|
})
|
|
|
|
})
|
|
|
|
new_count = self.db.execute("SELECT LAST_INSERT_ROWID() as last").fetchone()
|
|
|
|
new_count = self.db.execute("SELECT LAST_INSERT_ROWID() as last").fetchone()
|
|
|
|
self.dblock.release()
|
|
|
|
self.db_lock.release()
|
|
|
|
|
|
|
|
|
|
|
|
return new_count['last']
|
|
|
|
return new_count['last']
|
|
|
|
|
|
|
|
|
|
|
|
@validate_call
|
|
|
|
@validate_call
|
|
|
|
def set_processed(self, count:int, processed:bool=True) -> bool:
|
|
|
|
def set_processed(self, count:int, processed:bool=True) -> bool:
|
|
|
|
self.dblock.acquire()
|
|
|
|
self.db_lock.acquire()
|
|
|
|
with self.db:
|
|
|
|
with self.db:
|
|
|
|
try:
|
|
|
|
try:
|
|
|
|
self.db.execute("UPDATE Messages SET processed = ? WHERE count = ?", (processed, count))
|
|
|
|
self.db.execute("UPDATE Messages SET processed = ? WHERE count = ?", (processed, count))
|
|
|
|
return True
|
|
|
|
return True
|
|
|
|
except:
|
|
|
|
except:
|
|
|
|
return False
|
|
|
|
return False
|
|
|
|
self.dblock.release()
|
|
|
|
finally:
|
|
|
|
|
|
|
|
self.db_lock.release()
|
|
|
|
|
|
|
|
|
|
|
|
def __iter__(self) -> Generator[RowObject, None, None]:
|
|
|
|
def __iter__(self) -> Generator[RowObject, None, None]:
|
|
|
|
yield from self.iterate()
|
|
|
|
yield from self.iterate()
|
|
|
|