chore: all model.query replace to db.session.query (#19521)

api/controllers/console/workspace/workspace.py

 from flask import request
 from flask_login import current_user
 from flask_restful import Resource, fields, inputs, marshal, marshal_with, reqparse
+from sqlalchemy import select
 from werkzeug.exceptions import Unauthorized
 
 import services
         parser.add_argument("limit", type=inputs.int_range(1, 100), required=False, default=20, location="args")
         args = parser.parse_args()
 
-        tenants = Tenant.query.order_by(Tenant.created_at.desc()).paginate(
-            page=args["page"], per_page=args["limit"], error_out=False
-        )
+        stmt = select(Tenant).order_by(Tenant.created_at.desc())
+        tenants = db.paginate(select=stmt, page=args["page"], per_page=args["limit"], error_out=False)
         has_more = False
 
         if tenants.has_next:
         parser.add_argument("replace_webapp_logo", type=str, location="json")
         args = parser.parse_args()
 
-        tenant = Tenant.query.filter(Tenant.id == current_user.current_tenant_id).one_or_404()
+        tenant = db.get_or_404(Tenant, current_user.current_tenant_id)
 
         custom_config_dict = {
             "remove_webapp_brand": args["remove_webapp_brand"],
         parser.add_argument("name", type=str, required=True, location="json")
         args = parser.parse_args()
 
-        tenant = Tenant.query.filter(Tenant.id == current_user.current_tenant_id).one_or_404()
+        tenant = db.get_or_404(Tenant, current_user.current_tenant_id)
         tenant.name = args["name"]
         db.session.commit()
 

api/core/rag/extractor/notion_extractor.py

 
     @classmethod
     def _get_access_token(cls, tenant_id: str, notion_workspace_id: str) -> str:
-        data_source_binding = DataSourceOauthBinding.query.filter(
-            db.and_(
-                DataSourceOauthBinding.tenant_id == tenant_id,
-                DataSourceOauthBinding.provider == "notion",
-                DataSourceOauthBinding.disabled == False,
-                DataSourceOauthBinding.source_info["workspace_id"] == f'"{notion_workspace_id}"',
+        data_source_binding = (
+            db.session.query(DataSourceOauthBinding)
+            .filter(
+                db.and_(
+                    DataSourceOauthBinding.tenant_id == tenant_id,
+                    DataSourceOauthBinding.provider == "notion",
+                    DataSourceOauthBinding.disabled == False,
+                    DataSourceOauthBinding.source_info["workspace_id"] == f'"{notion_workspace_id}"',
+                )
             )
-        ).first()
+            .first()
+        )
 
         if not data_source_binding:
             raise Exception(

api/libs/oauth_data_source.py

             "total": len(pages),
         }
         # save data source binding
-        data_source_binding = DataSourceOauthBinding.query.filter(
-            db.and_(
-                DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
-                DataSourceOauthBinding.provider == "notion",
-                DataSourceOauthBinding.access_token == access_token,
+        data_source_binding = (
+            db.session.query(DataSourceOauthBinding)
+            .filter(
+                db.and_(
+                    DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
+                    DataSourceOauthBinding.provider == "notion",
+                    DataSourceOauthBinding.access_token == access_token,
+                )
             )
-        ).first()
+            .first()
+        )
         if data_source_binding:
             data_source_binding.source_info = source_info
             data_source_binding.disabled = False
             "total": len(pages),
         }
         # save data source binding
-        data_source_binding = DataSourceOauthBinding.query.filter(
-            db.and_(
-                DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
-                DataSourceOauthBinding.provider == "notion",
-                DataSourceOauthBinding.access_token == access_token,
+        data_source_binding = (
+            db.session.query(DataSourceOauthBinding)
+            .filter(
+                db.and_(
+                    DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
+                    DataSourceOauthBinding.provider == "notion",
+                    DataSourceOauthBinding.access_token == access_token,
+                )
             )
-        ).first()
+            .first()
+        )
         if data_source_binding:
             data_source_binding.source_info = source_info
             data_source_binding.disabled = False
 
     def sync_data_source(self, binding_id: str):
         # save data source binding
-        data_source_binding = DataSourceOauthBinding.query.filter(
-            db.and_(
-                DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
-                DataSourceOauthBinding.provider == "notion",
-                DataSourceOauthBinding.id == binding_id,
-                DataSourceOauthBinding.disabled == False,
+        data_source_binding = (
+            db.session.query(DataSourceOauthBinding)
+            .filter(
+                db.and_(
+                    DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
+                    DataSourceOauthBinding.provider == "notion",
+                    DataSourceOauthBinding.id == binding_id,
+                    DataSourceOauthBinding.disabled == False,
+                )
             )
-        ).first()
+            .first()
+        )
         if data_source_binding:
             # get all authorized pages
             pages = self.get_authorized_pages(data_source_binding.access_token)

api/schedule/mail_clean_document_notify_task.py

             if plan != "sandbox":
                 knowledge_details = []
                 # check tenant
-                tenant = Tenant.query.filter(Tenant.id == tenant_id).first()
+                tenant = db.session.query(Tenant).filter(Tenant.id == tenant_id).first()
                 if not tenant:
                     continue
                 # check current owner

api/services/account_service.py

         """Link account integrate"""
         try:
             # Query whether there is an existing binding record for the same provider
-            account_integrate: Optional[AccountIntegrate] = AccountIntegrate.query.filter_by(
-                account_id=account.id, provider=provider
-            ).first()
+            account_integrate: Optional[AccountIntegrate] = (
+                db.session.query(AccountIntegrate).filter_by(account_id=account.id, provider=provider).first()
+            )
 
             if account_integrate:
                 # If it exists, update the record
 
     @staticmethod
     def get_custom_config(tenant_id: str) -> dict:
-        tenant = Tenant.query.filter(Tenant.id == tenant_id).one_or_404()
+        tenant = db.get_or_404(Tenant, tenant_id)
 
         return cast(dict, tenant.custom_config_dict)
 

api/services/annotation_service.py

 
 import pandas as pd
 from flask_login import current_user
-from sqlalchemy import or_
+from sqlalchemy import or_, select
 from werkzeug.datastructures import FileStorage
 from werkzeug.exceptions import NotFound
 
         if not app:
             raise NotFound("App not found")
         if keyword:
-            annotations = (
-                MessageAnnotation.query.filter(MessageAnnotation.app_id == app_id)
+            stmt = (
+                select(MessageAnnotation)
+                .filter(MessageAnnotation.app_id == app_id)
                 .filter(
                     or_(
                         MessageAnnotation.question.ilike("%{}%".format(keyword)),
                     )
                 )
                 .order_by(MessageAnnotation.created_at.desc(), MessageAnnotation.id.desc())
-                .paginate(page=page, per_page=limit, max_per_page=100, error_out=False)
             )
         else:
-            annotations = (
-                MessageAnnotation.query.filter(MessageAnnotation.app_id == app_id)
+            stmt = (
+                select(MessageAnnotation)
+                .filter(MessageAnnotation.app_id == app_id)
                 .order_by(MessageAnnotation.created_at.desc(), MessageAnnotation.id.desc())
-                .paginate(page=page, per_page=limit, max_per_page=100, error_out=False)
             )
+        annotations = db.paginate(select=stmt, page=page, per_page=limit, max_per_page=100, error_out=False)
         return annotations.items, annotations.total
 
     @classmethod
         if not annotation:
             raise NotFound("Annotation not found")
 
-        annotation_hit_histories = (
-            AppAnnotationHitHistory.query.filter(
+        stmt = (
+            select(AppAnnotationHitHistory)
+            .filter(
                 AppAnnotationHitHistory.app_id == app_id,
                 AppAnnotationHitHistory.annotation_id == annotation_id,
             )
             .order_by(AppAnnotationHitHistory.created_at.desc())
-            .paginate(page=page, per_page=limit, max_per_page=100, error_out=False)
+        )
+        annotation_hit_histories = db.paginate(
+            select=stmt, page=page, per_page=limit, max_per_page=100, error_out=False
         )
         return annotation_hit_histories.items, annotation_hit_histories.total
 

api/services/dataset_service.py

                             exist_document[data_source_info["notion_page_id"]] = document.id
                     for notion_info in notion_info_list:
                         workspace_id = notion_info.workspace_id
-                        data_source_binding = DataSourceOauthBinding.query.filter(
-                            db.and_(
-                                DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
-                                DataSourceOauthBinding.provider == "notion",
-                                DataSourceOauthBinding.disabled == False,
-                                DataSourceOauthBinding.source_info["workspace_id"] == f'"{workspace_id}"',
+                        data_source_binding = (
+                            db.session.query(DataSourceOauthBinding)
+                            .filter(
+                                db.and_(
+                                    DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
+                                    DataSourceOauthBinding.provider == "notion",
+                                    DataSourceOauthBinding.disabled == False,
+                                    DataSourceOauthBinding.source_info["workspace_id"] == f'"{workspace_id}"',
+                                )
                             )
-                        ).first()
+                            .first()
+                        )
                         if not data_source_binding:
                             raise ValueError("Data source binding not found.")
                         for page in notion_info.pages:
                 notion_info_list = document_data.data_source.info_list.notion_info_list
                 for notion_info in notion_info_list:
                     workspace_id = notion_info.workspace_id
-                    data_source_binding = DataSourceOauthBinding.query.filter(
-                        db.and_(
-                            DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
-                            DataSourceOauthBinding.provider == "notion",
-                            DataSourceOauthBinding.disabled == False,
-                            DataSourceOauthBinding.source_info["workspace_id"] == f'"{workspace_id}"',
+                    data_source_binding = (
+                        db.session.query(DataSourceOauthBinding)
+                        .filter(
+                            db.and_(
+                                DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
+                                DataSourceOauthBinding.provider == "notion",
+                                DataSourceOauthBinding.disabled == False,
+                                DataSourceOauthBinding.source_info["workspace_id"] == f'"{workspace_id}"',
+                            )
                         )
-                    ).first()
+                        .first()
+                    )
                     if not data_source_binding:
                         raise ValueError("Data source binding not found.")
                     for page in notion_info.pages:

api/tasks/document_indexing_sync_task.py

         page_id = data_source_info["notion_page_id"]
         page_type = data_source_info["type"]
         page_edited_time = data_source_info["last_edited_time"]
-        data_source_binding = DataSourceOauthBinding.query.filter(
-            db.and_(
-                DataSourceOauthBinding.tenant_id == document.tenant_id,
-                DataSourceOauthBinding.provider == "notion",
-                DataSourceOauthBinding.disabled == False,
-                DataSourceOauthBinding.source_info["workspace_id"] == f'"{workspace_id}"',
+        data_source_binding = (
+            db.session.query(DataSourceOauthBinding)
+            .filter(
+                db.and_(
+                    DataSourceOauthBinding.tenant_id == document.tenant_id,
+                    DataSourceOauthBinding.provider == "notion",
+                    DataSourceOauthBinding.disabled == False,
+                    DataSourceOauthBinding.source_info["workspace_id"] == f'"{workspace_id}"',
+                )
             )
-        ).first()
+            .first()
+        )
         if not data_source_binding:
             raise ValueError("Data source binding not found.")