refactor: migrate db.session.query to select in infra layer (#33694)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/schedule/check_upgradable_plugin_task.py

@@ -3,6 +3,7 @@ import math
 import time
 
 import click
+from sqlalchemy import select
 
 import app
 from core.helper.marketplace import fetch_global_plugin_manifest
@@ -28,17 +29,15 @@ def check_upgradable_plugin_task():
     now_seconds_of_day = time.time() % 86400 - 30  # we assume the tz is UTC
     click.echo(click.style(f"Now seconds of day: {now_seconds_of_day}", fg="green"))
 
-    strategies = (
-        db.session.query(TenantPluginAutoUpgradeStrategy)
-        .where(
+    strategies = db.session.scalars(
+        select(TenantPluginAutoUpgradeStrategy).where(
             TenantPluginAutoUpgradeStrategy.upgrade_time_of_day >= now_seconds_of_day,
             TenantPluginAutoUpgradeStrategy.upgrade_time_of_day
             < now_seconds_of_day + AUTO_UPGRADE_MINIMAL_CHECKING_INTERVAL,
             TenantPluginAutoUpgradeStrategy.strategy_setting
             != TenantPluginAutoUpgradeStrategy.StrategySetting.DISABLED,
         )
-        .all()
-    )
+    ).all()
 
     total_strategies = len(strategies)
     click.echo(click.style(f"Total strategies: {total_strategies}", fg="green"))

api/schedule/clean_embedding_cache_task.py

@@ -2,7 +2,7 @@ import datetime
 import time
 
 import click
-from sqlalchemy import text
+from sqlalchemy import select, text
 from sqlalchemy.exc import SQLAlchemyError
 
 import app
@@ -19,14 +19,12 @@ def clean_embedding_cache_task():
     thirty_days_ago = datetime.datetime.now() - datetime.timedelta(days=clean_days)
     while True:
         try:
-            embedding_ids = (
-                db.session.query(Embedding.id)
+            embedding_ids = db.session.scalars(
+                select(Embedding.id)
                 .where(Embedding.created_at < thirty_days_ago)
                 .order_by(Embedding.created_at.desc())
                 .limit(100)
-                .all()
-            )
-            embedding_ids = [embedding_id[0] for embedding_id in embedding_ids]
+            ).all()
         except SQLAlchemyError:
             raise
         if embedding_ids:

api/schedule/clean_unused_datasets_task.py

@@ -3,7 +3,7 @@ import time
 from typing import TypedDict
 
 import click
-from sqlalchemy import func, select
+from sqlalchemy import func, select, update
 from sqlalchemy.exc import SQLAlchemyError
 
 import app
@@ -51,7 +51,7 @@ def clean_unused_datasets_task():
             try:
                 # Subquery for counting new documents
                 document_subquery_new = (
-                    db.session.query(Document.dataset_id, func.count(Document.id).label("document_count"))
+                    select(Document.dataset_id, func.count(Document.id).label("document_count"))
                     .where(
                         Document.indexing_status == "completed",
                         Document.enabled == True,
@@ -64,7 +64,7 @@ def clean_unused_datasets_task():
 
                 # Subquery for counting old documents
                 document_subquery_old = (
-                    db.session.query(Document.dataset_id, func.count(Document.id).label("document_count"))
+                    select(Document.dataset_id, func.count(Document.id).label("document_count"))
                     .where(
                         Document.indexing_status == "completed",
                         Document.enabled == True,
@@ -142,8 +142,8 @@ def clean_unused_datasets_task():
                             index_processor.clean(dataset, None)
 
                             # Update document
-                            db.session.query(Document).filter_by(dataset_id=dataset.id).update(
-                                {Document.enabled: False}
+                            db.session.execute(
+                                update(Document).where(Document.dataset_id == dataset.id).values(enabled=False)
                             )
                             db.session.commit()
                             click.echo(click.style(f"Cleaned unused dataset {dataset.id} from db success!", fg="green"))

api/schedule/create_tidb_serverless_task.py

@@ -1,6 +1,7 @@
 import time
 
 import click
+from sqlalchemy import func, select
 
 import app
 from configs import dify_config
@@ -20,7 +21,7 @@ def create_tidb_serverless_task():
         try:
             # check the number of idle tidb serverless
             idle_tidb_serverless_number = (
-                db.session.query(TidbAuthBinding).where(TidbAuthBinding.active == False).count()
+                db.session.scalar(select(func.count(TidbAuthBinding.id)).where(TidbAuthBinding.active == False)) or 0
             )
             if idle_tidb_serverless_number >= tidb_serverless_number:
                 break

api/schedule/mail_clean_document_notify_task.py

@@ -49,16 +49,18 @@ def mail_clean_document_notify_task():
             if plan != CloudPlan.SANDBOX:
                 knowledge_details = []
                 # check tenant
-                tenant = db.session.query(Tenant).where(Tenant.id == tenant_id).first()
+                tenant = db.session.scalar(select(Tenant).where(Tenant.id == tenant_id))
                 if not tenant:
                     continue
                 # check current owner
-                current_owner_join = (
-                    db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, role="owner").first()
+                current_owner_join = db.session.scalar(
+                    select(TenantAccountJoin)
+                    .where(TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.role == "owner")
+                    .limit(1)
                 )
                 if not current_owner_join:
                     continue
-                account = db.session.query(Account).where(Account.id == current_owner_join.account_id).first()
+                account = db.session.scalar(select(Account).where(Account.id == current_owner_join.account_id))
                 if not account:
                     continue
 
@@ -71,7 +73,7 @@ def mail_clean_document_notify_task():
                     )
 
                 for dataset_id, document_ids in dataset_auto_dataset_map.items():
-                    dataset = db.session.query(Dataset).where(Dataset.id == dataset_id).first()
+                    dataset = db.session.scalar(select(Dataset).where(Dataset.id == dataset_id))
                     if dataset:
                         document_count = len(document_ids)
                         knowledge_details.append(rf"Knowledge base {dataset.name}: {document_count} documents")