Skip to content

refactor(models): remove legacy db.session property wrappers on Dataset #41569

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
asukaminato0721 merged 2 commits into from
Sep 1, 2026
Merged

refactor(models): remove legacy db.session property wrappers on Dataset #41569

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
78ad674
refactor(models): remove legacy db.session property wrappers on Dataset
keith991001 Sep 1, 2026
81ffc00
fix: route empty RAG pipeline dataset response through dataset_detail…
keith991001 Sep 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions api/controllers/console/datasets/rag_pipeline/rag_pipeline_datasets.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
with_current_user,
)
from extensions.ext_database import db
from fields.dataset_fields import DatasetDetailResponse
from fields.dataset_fields import DatasetDetailResponse, dataset_detail_response_source
from libs.helper import dump_response
from libs.login import login_required
from models import Account
Expand Down Expand Up @@ -116,6 +116,7 @@ def post(self, current_tenant_id: str, current_user: Account) -> JsonResponseWit
# The role of the current user in the ta table must be admin, owner, or editor, or dataset_operator
if not current_user.is_dataset_editor:
raise Forbidden()
session = db.session()
dataset = DatasetService.create_empty_rag_pipeline_dataset(
tenant_id=current_tenant_id,
rag_pipeline_dataset_create_entity=RagPipelineDatasetCreateEntity(
Expand All @@ -129,6 +130,6 @@ def post(self, current_tenant_id: str, current_user: Account) -> JsonResponseWit
permission=DatasetPermissionEnum.ONLY_ME,
partial_member_list=None,
),
session=db.session(),
session=session,
)
return dump_response(DatasetDetailResponse, dataset), 201
return dump_response(DatasetDetailResponse, dataset_detail_response_source(dataset, session=session)), 201
20 changes: 0 additions & 20 deletions api/models/dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -212,17 +212,9 @@ class Dataset(Base):
enable_api = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("true"))
is_multimodal = mapped_column(sa.Boolean, default=False, nullable=False, server_default=sa.text("false"))

@property
def total_documents(self) -> int:
return self.get_total_documents(session=db.session())

def get_total_documents(self, *, session: Session) -> int:
return self.get_document_count(session=session)

@property
def total_available_documents(self) -> int:
return self.get_total_available_documents(session=db.session())

def get_total_available_documents(self, *, session: Session) -> int:
return (
session.scalar(
Expand Down Expand Up @@ -258,20 +250,12 @@ def created_by_account(self):
def get_created_by_account(self, *, session: Session) -> Account | None:
return session.get(Account, self.created_by)

@property
def author_name(self) -> str | None:
return self.get_author_name(session=db.session())

def get_author_name(self, *, session: Session) -> str | None:
account = self.get_created_by_account(session=session)
if account:
return account.name
return None

@property
def latest_process_rule(self):
return self.get_latest_process_rule(session=db.session())

def get_latest_process_rule(self, *, session: Session) -> "DatasetProcessRule | None":
return session.scalar(
select(DatasetProcessRule)
Expand Down Expand Up @@ -391,10 +375,6 @@ def get_tags(self, *, session: Session) -> Sequence[Tag]:

return tags or []

@property
def external_knowledge_info(self) -> dict[str, Any] | None:
return self.get_external_knowledge_info(session=db.session())

def get_external_knowledge_info(self, *, session: Session) -> dict[str, Any] | None:
if self.provider != "external":
return None
Expand Down
4 changes: 2 additions & 2 deletions api/tests/test_containers_integration_tests/models/test_dataset_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ def test_dataset_with_documents_relationship(self, db_session_with_containers: S
db_session_with_containers.add(doc)
db_session_with_containers.flush()

assert dataset.total_documents == 3
assert dataset.get_total_documents(session=db_session_with_containers) == 3

def test_dataset_available_documents_count(self, db_session_with_containers: Session) -> None:
"""Test dataset can count available documents."""
Expand Down Expand Up @@ -104,7 +104,7 @@ def test_dataset_available_documents_count(self, db_session_with_containers: Ses
db_session_with_containers.add_all([doc_available, doc_pending, doc_disabled])
db_session_with_containers.flush()

assert dataset.total_available_documents == 1
assert dataset.get_total_available_documents(session=db_session_with_containers) == 1

def test_dataset_word_count_aggregation(self, db_session_with_containers: Session) -> None:
"""Test dataset can aggregate word count from documents."""
Expand Down