tcsenpai/pensieve
1
0
Fork 0
mirror of https://github.com/tcsenpai/pensieve.git synced 2025-06-07 19:55:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: concurrency bug

Browse Source
This commit is contained in:
arkohut 2024-10-09 11:34:37 +08:00
parent 50a73de847
commit 3a8dc920bc
1 changed files with 18 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
memos/models.py
View File

@ -28,6 +28,7 @@ import json
from .embedding import get_embeddings from .embedding import get_embeddings
from sqlite_vec import serialize_float32 from sqlite_vec import serialize_float32
import asyncio import asyncio
import threading
class Base(DeclarativeBase): class Base(DeclarativeBase):
@ -197,12 +198,16 @@ def load_extension(dbapi_conn, connection_record):
# load vector ext # load vector ext
sqlite_vec.load(dbapi_conn) sqlite_vec.load(dbapi_conn)
# Set WAL mode after loading extensions
dbapi_conn.execute("PRAGMA journal_mode=WAL")
def init_database(): def init_database():
"""Initialize the database.""" """Initialize the database."""
db_path = get_database_path() db_path = get_database_path()
engine = create_engine(f"sqlite:///{db_path}") engine = create_engine(f"sqlite:///{db_path}")
# Use a single event listener for both extension loading and WAL mode setting
event.listen(engine, "connect", load_extension) event.listen(engine, "connect", load_extension)
try: try:
@ -432,12 +437,22 @@ def delete_fts_and_vec(mapper, connection, target):
) )
# Update the event listener to use asyncio def run_async(coro):
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
return loop.run_until_complete(coro)
def update_fts_and_vec_sync(mapper, connection, target): def update_fts_and_vec_sync(mapper, connection, target):
asyncio.run(update_fts_and_vec(mapper, connection, target)) def run_in_thread():
run_async(update_fts_and_vec(mapper, connection, target))
thread = threading.Thread(target=run_in_thread)
thread.start()
thread.join()
# Replace the old event listener with the new sync version # Replace the old event listener with the new sync version
event.listen(EntityModel, "after_insert", update_fts_and_vec_sync) event.listen(EntityModel, "after_insert", update_fts_and_vec_sync)
event.listen(EntityModel, "after_update", update_fts_and_vec_sync) event.listen(EntityModel, "after_update", update_fts_and_vec_sync)
event.listen(EntityModel, "after_delete", delete_fts_and_vec) event.listen(EntityModel, "after_delete", delete_fts_and_vec)