chore(api): migrate event handlers to use Session(db.engine) (#35234)

api/events/event_handlers/create_document_index.py

@@ -6,9 +6,9 @@ import click
 from sqlalchemy import select
 from werkzeug.exceptions import NotFound
 
+from core.db.session_factory import session_factory
 from core.indexing_runner import DocumentIsPausedError, IndexingRunner
 from events.document_index_event import document_index_created
-from extensions.ext_database import db
 from libs.datetime_utils import naive_utc_now
 from models.dataset import Document
 from models.enums import IndexingStatus
@@ -22,24 +22,25 @@ def handle(sender, **kwargs):
     document_ids = kwargs.get("document_ids", [])
     documents = []
     start_at = time.perf_counter()
-    for document_id in document_ids:
-        logger.info(click.style(f"Start process document: {document_id}", fg="green"))
+    with session_factory.create_session() as session:
+        for document_id in document_ids:
+            logger.info(click.style(f"Start process document: {document_id}", fg="green"))
 
-        document = db.session.scalar(
-            select(Document).where(
-                Document.id == document_id,
-                Document.dataset_id == dataset_id,
+            document = session.scalar(
+                select(Document).where(
+                    Document.id == document_id,
+                    Document.dataset_id == dataset_id,
+                )
             )
-        )
 
-        if not document:
-            raise NotFound("Document not found")
+            if not document:
+                raise NotFound("Document not found")
 
-        document.indexing_status = IndexingStatus.PARSING
-        document.processing_started_at = naive_utc_now()
-        documents.append(document)
-        db.session.add(document)
-    db.session.commit()
+            document.indexing_status = IndexingStatus.PARSING
+            document.processing_started_at = naive_utc_now()
+            documents.append(document)
+            session.add(document)
+        session.commit()
 
     with contextlib.suppress(Exception):
         try:

api/events/event_handlers/create_installed_app_when_app_created.py

@@ -1,5 +1,5 @@
+from core.db.session_factory import session_factory
 from events.app_event import app_was_created
-from extensions.ext_database import db
 from models.model import InstalledApp
 
 
@@ -12,5 +12,6 @@ def handle(sender, **kwargs):
         app_id=app.id,
         app_owner_tenant_id=app.tenant_id,
     )
-    db.session.add(installed_app)
-    db.session.commit()
+    with session_factory.create_session() as session:
+        session.add(installed_app)
+        session.commit()

api/events/event_handlers/create_site_record_when_app_created.py

@@ -1,5 +1,5 @@
+from core.db.session_factory import session_factory
 from events.app_event import app_was_created
-from extensions.ext_database import db
 from models.enums import CustomizeTokenStrategy
 from models.model import Site
 
@@ -22,6 +22,6 @@ def handle(sender, **kwargs):
             created_by=app.created_by,
             updated_by=app.updated_by,
         )
-
-        db.session.add(site)
-        db.session.commit()
+        with session_factory.create_session() as session:
+            session.add(site)
+            session.commit()