refactor: select in rag_pipeline by RenzoMXD · Pull Request #34495 · langgenius/dify

RenzoMXD · 2026-04-03T03:18:06Z

Summary

Migrate all 19 db.session.query() calls to SQLAlchemy 2.x select() style in rag_pipeline/rag_pipeline.py
Replace .where().first() with db.session.scalar(select().where().limit(1))
Use session.get() for primary key lookups on Pipeline, Workflow, and Document
Replace .count() with db.session.scalar(select(func.count()))
Replace .query(func.max()).scalar() with db.session.scalar(select(func.max()))
Replace dynamic .query().all() with db.session.scalars(stmt).all()
Update test mocks in test_rag_pipeline_service.py to match new patterns (unavoidable)

Test plan

All 249 rag_pipeline unit tests pass
Basedpyright type check passes (0 errors)

Part of #22668

github-actions · 2026-04-03T03:19:37Z

Pyrefly Diff

No changes detected.

github-actions · 2026-04-03T03:21:22Z

Pyrefly Diff

No changes detected.

Copilot

Pull request overview

This PR migrates ORM usage in the RAG pipeline service from SQLAlchemy 1.x Query patterns to SQLAlchemy 2.x select()/Session.scalar()/Session.scalars() patterns to improve typing and modernize query style (part of #22668).

Changes:

Refactor multiple db.session.query(...) call sites in rag_pipeline.py to select()-style statements, including get() for PK lookups and select(func.count(...)) for counts.
Replace .first() and .count() patterns with db.session.scalar(...) and db.session.scalars(...).all() equivalents.
Update unit test mocks in test_rag_pipeline_service.py to match the new session APIs.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 11 comments.

File	Description
api/services/rag_pipeline/rag_pipeline.py	Replaces legacy `Query` usage with SQLAlchemy 2.x `select()`/`scalar()`/`scalars()` and `session.get()` patterns in the RAG pipeline service.
api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py	Adjusts mocks/patches to align with the refactored `scalar()` / `get()` / `scalars()` access patterns.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-03T03:47:51Z

@@ -1656,8 +1602,7 @@ def _event_generator():

    document = SimpleNamespace(indexing_status="waiting", error=None)
    query = mocker.Mock()


Unused query mock left over after switching to db.session.get(...); remove it to avoid confusing test setup and to keep mocks minimal.

Suggested change

query = mocker.Mock()

Copilot · 2026-04-03T03:47:52Z



 def test_get_pipeline_raises_when_dataset_missing(mocker, rag_pipeline_service) -> None:
    query = mocker.Mock()


query = mocker.Mock() is now unused after patching db.session.scalar; remove the unused mock to keep the test focused and reduce noise.

Copilot · 2026-04-03T03:47:52Z

@@ -1742,8 +1682,7 @@ def test_get_pipeline_raises_when_dataset_missing(mocker, rag_pipeline_service)
 def test_get_pipeline_raises_when_pipeline_missing(mocker, rag_pipeline_service) -> None:
    dataset = SimpleNamespace(pipeline_id="p1")
    query = mocker.Mock()


query = mocker.Mock() is unused in this test after migrating away from session.query(...); it can be deleted.

Copilot · 2026-04-03T03:47:52Z

@@ -1783,8 +1722,7 @@ def test_get_pipeline_templates_builtin_en_us_no_fallback(mocker) -> None:
 def test_update_customized_pipeline_template_commits_when_name_empty(mocker) -> None:
    template = SimpleNamespace(name="old", description="old", icon={}, updated_by=None)
    query = mocker.Mock()


Unused query mock: the test patches db.session.scalar directly now, so this variable can be removed.

Copilot · 2026-04-03T03:47:52Z

@@ -2011,8 +1949,7 @@ def test_run_free_workflow_node_delegates_to_handle_result(mocker, rag_pipeline_
 def test_publish_customized_pipeline_template_raises_when_workflow_missing(mocker, rag_pipeline_service) -> None:
    pipeline = SimpleNamespace(id="p1", tenant_id="t1", workflow_id="wf-1")
    query = mocker.Mock()


query = mocker.Mock() is unused after switching this test to patch db.session.get; remove it to avoid dead code in the test.

Copilot · 2026-04-03T03:47:53Z

@@ -2069,8 +2001,8 @@ def test_retry_error_document_raises_when_workflow_missing(mocker, rag_pipeline_
    exec_log = SimpleNamespace(pipeline_id="p1")
    pipeline = SimpleNamespace(id="p1")
    query = mocker.Mock()


Unused query mock in this test after refactor to db.session.scalar/get; please remove it.

Copilot · 2026-04-03T03:47:53Z

@@ -2086,8 +2018,7 @@ def test_get_datasource_plugins_returns_empty_for_non_datasource_nodes(mocker, r
        graph_dict={"nodes": [{"id": "n1", "data": {"type": "start"}}]}, rag_pipeline_variables=[]
    )
    query = mocker.Mock()


query = mocker.Mock() is unused here now that the test patches db.session.scalar directly; remove to reduce unnecessary setup.

Suggested change

query = mocker.Mock()

Copilot · 2026-04-03T03:47:53Z

@@ -2250,8 +2181,7 @@ def test_get_datasource_plugins_handles_empty_datasource_data_and_non_published(
        rag_pipeline_variables=[{"variable": "v1", "belong_to_node_id": "shared"}],
    )
    query = mocker.Mock()


Unused query mock remains after switching this test to patch db.session.scalar; remove it to keep the test setup minimal.

Copilot · 2026-04-03T03:47:53Z

@@ -2291,8 +2221,7 @@ def test_get_datasource_plugins_extracts_user_inputs_and_credentials(mocker, rag
        ],
    )
    query = mocker.Mock()


Unused query mock: this test now patches db.session.scalar directly, so the query variable can be deleted.

Copilot · 2026-04-03T03:47:54Z

@@ -2310,8 +2239,7 @@ def test_get_pipeline_returns_pipeline_when_found(mocker, rag_pipeline_service)
    dataset = SimpleNamespace(pipeline_id="p1")
    pipeline = SimpleNamespace(id="p1")
    query = mocker.Mock()


Unused query mock in this test after migrating to db.session.scalar mocking; please remove it.

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

refactor: select in rag_pipeline

143f095

RenzoMXD requested review from JohnJyong, QuantumGhost and laipz8200 as code owners April 3, 2026 03:18

dosubot Bot added size:L This PR changes 100-499 lines, ignoring generated files. refactor labels Apr 3, 2026

[autofix.ci] apply automated fixes

9598f6c

asukaminato0721 enabled auto-merge April 3, 2026 03:41

asukaminato0721 approved these changes Apr 3, 2026

View reviewed changes

asukaminato0721 added this pull request to the merge queue Apr 3, 2026

dosubot Bot added the lgtm This PR has been approved by a maintainer label Apr 3, 2026

asukaminato0721 requested a review from Copilot April 3, 2026 03:42

Copilot started reviewing on behalf of asukaminato0721 April 3, 2026 03:43 View session

Copilot AI reviewed Apr 3, 2026

View reviewed changes

Merged via the queue into langgenius:main with commit 7eb632e Apr 3, 2026
31 of 33 checks passed

RenzoMXD deleted the refactor/select-rag-pipeline branch April 3, 2026 03:58

HanqingZ pushed a commit to HanqingZ/dify that referenced this pull request Apr 23, 2026

refactor: select in rag_pipeline (langgenius#34495)

ce32f29

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

		@@ -1656,8 +1602,7 @@ def _event_generator():

		document = SimpleNamespace(indexing_status="waiting", error=None)
		query = mocker.Mock()



		def test_get_pipeline_raises_when_dataset_missing(mocker, rag_pipeline_service) -> None:
		query = mocker.Mock()

Conversation

RenzoMXD commented Apr 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Uh oh!

github-actions Bot commented Apr 3, 2026

Pyrefly Diff

Uh oh!

github-actions Bot commented Apr 3, 2026

Pyrefly Diff

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

RenzoMXD commented Apr 3, 2026 •

edited

Loading