From 3a8dc920bc040fb9c56d7ee7493244c572413b31 Mon Sep 17 00:00:00 2001
From: arkohut <39525455+arkohut@users.noreply.github.com>
Date: Wed, 9 Oct 2024 11:34:37 +0800
Subject: [PATCH] fix: concurrency bug

---
 memos/models.py | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/memos/models.py b/memos/models.py
index e46e4e5..c451980 100644
--- a/memos/models.py
+++ b/memos/models.py
@@ -28,6 +28,7 @@ import json
 from .embedding import get_embeddings
 from sqlite_vec import serialize_float32
 import asyncio
+import threading
 
 
 class Base(DeclarativeBase):
@@ -197,12 +198,16 @@ def load_extension(dbapi_conn, connection_record):
     # load vector ext
     sqlite_vec.load(dbapi_conn)
 
+    # Set WAL mode after loading extensions
+    dbapi_conn.execute("PRAGMA journal_mode=WAL")
+
 
 def init_database():
     """Initialize the database."""
     db_path = get_database_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)
 
     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):
-    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
 event.listen(EntityModel, "after_insert", 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)
\ No newline at end of file