mirror of
https://github.com/langgenius/dify.git
synced 2026-04-13 12:00:19 -04:00
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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user