estel/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

replace db with sa to get typing support (#23240)

Browse Source
This commit is contained in:
Asuka Minato
2025-08-03 00:54:23 +09:00
committed by GitHub
parent ff9fd0cdb2
commit 58608f51da
17 changed files with 623 additions and 613 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
api/services/plugin/data_migration.py
View File

@@ -2,6 +2,7 @@ import json
import logging
import click
import sqlalchemy as sa
from core.plugin.entities.plugin import GenericProviderID, ModelProviderID, ToolProviderID
from models.engine import db
@@ -38,7 +39,7 @@ class PluginDataMigration:
where {provider_column_name} not like '%/%' and {provider_column_name} is not null and {provider_column_name} != ''
limit 1000"""
with db.engine.begin() as conn:
rs = conn.execute(db.text(sql))
rs = conn.execute(sa.text(sql))
current_iter_count = 0
for i in rs:
@@ -94,7 +95,7 @@ limit 1000"""
:provider_name
{update_retrieval_model_sql}
where id = :record_id"""
conn.execute(db.text(sql), params)
conn.execute(sa.text(sql), params)
click.echo(
click.style(
f"[{processed_count}] Migrated [{table_name}] {record_id} ({provider_name})",
@@ -148,7 +149,7 @@ limit 1000"""
params = {"last_id": last_id or ""}
with db.engine.begin() as conn:
rs = conn.execute(db.text(sql), params)
rs = conn.execute(sa.text(sql), params)
current_iter_count = 0
batch_updates = []
@@ -193,7 +194,7 @@ limit 1000"""
SET {provider_column_name} = :updated_value
WHERE id = :record_id
"""
conn.execute(db.text(update_sql), [{"updated_value": u, "record_id": r} for u, r in batch_updates])
conn.execute(sa.text(update_sql), [{"updated_value": u, "record_id": r} for u, r in batch_updates])
click.echo(
click.style(
f"[{processed_count}] Batch migrated [{len(batch_updates)}] records from [{table_name}]",

3
api/services/plugin/plugin_migration.py
View File

@@ -9,6 +9,7 @@ from typing import Any, Optional
from uuid import uuid4
import click
import sqlalchemy as sa
import tqdm
from flask import Flask, current_app
from sqlalchemy.orm import Session
@@ -197,7 +198,7 @@ class PluginMigration:
"""
with Session(db.engine) as session:
rs = session.execute(
db.text(f"SELECT DISTINCT {column} FROM {table} WHERE tenant_id = :tenant_id"), {"tenant_id": tenant_id}
sa.text(f"SELECT DISTINCT {column} FROM {table} WHERE tenant_id = :tenant_id"), {"tenant_id": tenant_id}
)
result = []
for row in rs: