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

This commit is contained in:
Keith 2026-09-07 05:25:51 +00:00 committed by GitHub
parent 62a286332a
commit 11a9627cbb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 30 additions and 23 deletions

View File

@ -243,10 +243,6 @@ class Dataset(Base):
}
return self.retrieval_model or default_retrieval_model
@property
def created_by_account(self):
return self.get_created_by_account(session=db.session())
def get_created_by_account(self, *, session: Session) -> Account | None:
return session.get(Account, self.created_by)
@ -264,10 +260,6 @@ class Dataset(Base):
.limit(1)
)
@property
def app_count(self) -> int:
return self.get_app_count(session=db.session())
def get_app_count(self, *, session: Session) -> int:
return (
session.scalar(
@ -278,17 +270,12 @@ class Dataset(Base):
or 0
)
@property
def document_count(self) -> int:
return self.get_document_count(session=db.session())
def get_document_count(self, *, session: Session) -> int:
return session.scalar(select(func.count(Document.id)).where(Document.dataset_id == self.id)) or 0
@property
def available_document_count(self):
def get_available_document_count(self, *, session: Session) -> int:
return (
db.session.scalar(
session.scalar(
select(func.count(Document.id)).where(
Document.dataset_id == self.id,
Document.indexing_status == "completed",
@ -401,10 +388,6 @@ class Dataset(Base):
"external_knowledge_api_endpoint": json.loads(external_knowledge_api.settings).get("endpoint", ""),
}
@property
def is_published(self) -> bool:
return self.get_is_published(session=db.session())
def get_is_published(self, *, session: Session) -> bool:
if self.pipeline_id:
pipeline = session.scalar(select(Pipeline).where(Pipeline.id == self.pipeline_id))
@ -412,10 +395,6 @@ class Dataset(Base):
return pipeline.is_published
return False
@property
def doc_metadata(self) -> list[dict[str, str]]:
return self.get_doc_metadata(session=db.session())
def get_doc_metadata(self, *, session: Session) -> list[dict[str, str]]:
dataset_metadatas = session.scalars(select(DatasetMetadata).where(DatasetMetadata.dataset_id == self.id)).all()

View File

@ -1662,6 +1662,34 @@ class TestChildChunkSessionAccessors:
assert child_chunk.segment(session=sqlite_session) is None
class TestDatasetAvailableDocumentCount:
"""Regression coverage for ``Dataset.get_available_document_count``.
The accessor was a ``@property`` reaching for the global ``db.session``; it now takes a
caller-provided session, matching the other ``get_*`` accessors on ``Dataset``.
"""
def test_counts_only_completed_enabled_documents(self, sqlite_session: Session):
dataset = _make_dataset(dataset_id=str(uuid4()), tenant_id=str(uuid4()))
completed = _make_document(
document_id=str(uuid4()),
dataset_id=dataset.id,
tenant_id=dataset.tenant_id,
indexing_status=IndexingStatus.COMPLETED,
)
waiting = _make_document(
document_id=str(uuid4()),
dataset_id=dataset.id,
tenant_id=dataset.tenant_id,
position=2,
indexing_status=IndexingStatus.WAITING,
)
sqlite_session.add_all([dataset, completed, waiting])
sqlite_session.flush()
assert dataset.get_available_document_count(session=sqlite_session) == 1
class TestDocumentSegmentNeighborAccessors:
"""Regression coverage for ``DocumentSegment.previous_segment`` and ``next_segment`` refactored
to take a caller-provided session.