refactor(api): migrate tools, account, workflow and plugin services to SQLAlchemy 2.0 (#34966)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
wdeveloper16
2026-04-12 03:45:27 +02:00
committed by GitHub
parent 440602f52a
commit 0841b4c663
9 changed files with 111 additions and 102 deletions

View File

@@ -13,6 +13,7 @@ import sqlalchemy as sa
import tqdm
from flask import Flask, current_app
from pydantic import TypeAdapter
from sqlalchemy import func, select
from sqlalchemy.orm import Session
from core.agent.entities import AgentToolEntity
@@ -66,7 +67,7 @@ class PluginMigration:
current_time = started_at
with Session(db.engine) as session:
total_tenant_count = session.query(Tenant.id).count()
total_tenant_count = session.scalar(select(func.count(Tenant.id))) or 0
click.echo(click.style(f"Total tenant count: {total_tenant_count}", fg="white"))
@@ -123,9 +124,12 @@ class PluginMigration:
tenant_count = 0
for test_interval in test_intervals:
tenant_count = (
session.query(Tenant.id)
.where(Tenant.created_at.between(current_time, current_time + test_interval))
.count()
session.scalar(
select(func.count(Tenant.id)).where(
Tenant.created_at.between(current_time, current_time + test_interval)
)
)
or 0
)
if tenant_count <= 100:
interval = test_interval
@@ -147,8 +151,8 @@ class PluginMigration:
batch_end = min(current_time + interval, ended_at)
rs = (
session.query(Tenant.id)
rs = session.execute(
select(Tenant.id)
.where(Tenant.created_at.between(current_time, batch_end))
.order_by(Tenant.created_at)
)
@@ -235,7 +239,7 @@ class PluginMigration:
Extract tool tables.
"""
with Session(db.engine) as session:
rs = session.query(BuiltinToolProvider).where(BuiltinToolProvider.tenant_id == tenant_id).all()
rs = session.scalars(select(BuiltinToolProvider).where(BuiltinToolProvider.tenant_id == tenant_id)).all()
result = []
for row in rs:
result.append(ToolProviderID(row.provider).plugin_id)
@@ -249,7 +253,7 @@ class PluginMigration:
"""
with Session(db.engine) as session:
rs = session.query(Workflow).where(Workflow.tenant_id == tenant_id).all()
rs = session.scalars(select(Workflow).where(Workflow.tenant_id == tenant_id)).all()
result = []
for row in rs:
graph = row.graph_dict
@@ -272,7 +276,7 @@ class PluginMigration:
Extract app tables.
"""
with Session(db.engine) as session:
apps = session.query(App).where(App.tenant_id == tenant_id).all()
apps = session.scalars(select(App).where(App.tenant_id == tenant_id)).all()
if not apps:
return []
@@ -280,7 +284,7 @@ class PluginMigration:
app.app_model_config_id for app in apps if app.is_agent or app.mode == AppMode.AGENT_CHAT
]
rs = session.query(AppModelConfig).where(AppModelConfig.id.in_(agent_app_model_config_ids)).all()
rs = session.scalars(select(AppModelConfig).where(AppModelConfig.id.in_(agent_app_model_config_ids))).all()
result = []
for row in rs:
agent_config = row.agent_mode_dict