opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-08 02:43:49 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: migrate session.query to select API in rag pipeline task files (#34648)

Browse Source
This commit is contained in:
Renzo 2026-04-07 00:56:19 -05:00 committed by GitHub
parent bceb0eee9b
commit e2ecd68556
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 10 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
api/tasks/rag_pipeline/priority_rag_pipeline_run_task.py
View File

@ -10,6 +10,7 @@ from typing import Any
import click import click
from celery import shared_task # type: ignore from celery import shared_task # type: ignore
from flask import current_app, g from flask import current_app, g
from sqlalchemy import select
from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.orm import Session, sessionmaker
from configs import dify_config from configs import dify_config
@ -118,20 +119,20 @@ def run_single_rag_pipeline_task(rag_pipeline_invoke_entity: Mapping[str, Any],
with Session(db.engine, expire_on_commit=False) as session: with Session(db.engine, expire_on_commit=False) as session:
# Load required entities # Load required entities
account = session.query(Account).where(Account.id == user_id).first() account = session.scalar(select(Account).where(Account.id == user_id).limit(1))
if not account: if not account:
raise ValueError(f"Account {user_id} not found") raise ValueError(f"Account {user_id} not found")
tenant = session.query(Tenant).where(Tenant.id == tenant_id).first() tenant = session.scalar(select(Tenant).where(Tenant.id == tenant_id).limit(1))
if not tenant: if not tenant:
raise ValueError(f"Tenant {tenant_id} not found") raise ValueError(f"Tenant {tenant_id} not found")
account.current_tenant = tenant account.current_tenant = tenant
pipeline = session.query(Pipeline).where(Pipeline.id == pipeline_id).first() pipeline = session.scalar(select(Pipeline).where(Pipeline.id == pipeline_id).limit(1))
if not pipeline: if not pipeline:
raise ValueError(f"Pipeline {pipeline_id} not found") raise ValueError(f"Pipeline {pipeline_id} not found")
workflow = session.query(Workflow).where(Workflow.id == pipeline.workflow_id).first() workflow = session.scalar(select(Workflow).where(Workflow.id == pipeline.workflow_id).limit(1))
if not workflow: if not workflow:
raise ValueError(f"Workflow {pipeline.workflow_id} not found") raise ValueError(f"Workflow {pipeline.workflow_id} not found")

9
api/tasks/rag_pipeline/rag_pipeline_run_task.py
View File

@ -11,6 +11,7 @@ from typing import Any
import click import click
from celery import group, shared_task from celery import group, shared_task
from flask import current_app, g from flask import current_app, g
from sqlalchemy import select
from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.orm import Session, sessionmaker
from configs import dify_config from configs import dify_config
@ -132,20 +133,20 @@ def run_single_rag_pipeline_task(rag_pipeline_invoke_entity: Mapping[str, Any],
with Session(db.engine) as session: with Session(db.engine) as session:
# Load required entities # Load required entities
account = session.query(Account).where(Account.id == user_id).first() account = session.scalar(select(Account).where(Account.id == user_id).limit(1))
if not account: if not account:
raise ValueError(f"Account {user_id} not found") raise ValueError(f"Account {user_id} not found")
tenant = session.query(Tenant).where(Tenant.id == tenant_id).first() tenant = session.scalar(select(Tenant).where(Tenant.id == tenant_id).limit(1))
if not tenant: if not tenant:
raise ValueError(f"Tenant {tenant_id} not found") raise ValueError(f"Tenant {tenant_id} not found")
account.current_tenant = tenant account.current_tenant = tenant
pipeline = session.query(Pipeline).where(Pipeline.id == pipeline_id).first() pipeline = session.scalar(select(Pipeline).where(Pipeline.id == pipeline_id).limit(1))
if not pipeline: if not pipeline:
raise ValueError(f"Pipeline {pipeline_id} not found") raise ValueError(f"Pipeline {pipeline_id} not found")
workflow = session.query(Workflow).where(Workflow.id == pipeline.workflow_id).first() workflow = session.scalar(select(Workflow).where(Workflow.id == pipeline.workflow_id).limit(1))
if not workflow: if not workflow:
raise ValueError(f"Workflow {pipeline.workflow_id} not found") raise ValueError(f"Workflow {pipeline.workflow_id} not found")