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 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: def get_created_by_account(self, *, session: Session) -> Account | None:
return session.get(Account, self.created_by) return session.get(Account, self.created_by)
@ -264,10 +260,6 @@ class Dataset(Base):
.limit(1) .limit(1)
) )
@property
def app_count(self) -> int:
return self.get_app_count(session=db.session())
def get_app_count(self, *, session: Session) -> int: def get_app_count(self, *, session: Session) -> int:
return ( return (
session.scalar( session.scalar(
@ -278,17 +270,12 @@ class Dataset(Base):
or 0 or 0
) )
@property
def document_count(self) -> int:
return self.get_document_count(session=db.session())
def get_document_count(self, *, session: Session) -> int: def get_document_count(self, *, session: Session) -> int:
return session.scalar(select(func.count(Document.id)).where(Document.dataset_id == self.id)) or 0 return session.scalar(select(func.count(Document.id)).where(Document.dataset_id == self.id)) or 0
@property def get_available_document_count(self, *, session: Session) -> int:
def available_document_count(self):
return ( return (
db.session.scalar( session.scalar(
select(func.count(Document.id)).where( select(func.count(Document.id)).where(
Document.dataset_id == self.id, Document.dataset_id == self.id,
Document.indexing_status == "completed", Document.indexing_status == "completed",
@ -401,10 +388,6 @@ class Dataset(Base):
"external_knowledge_api_endpoint": json.loads(external_knowledge_api.settings).get("endpoint", ""), "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: def get_is_published(self, *, session: Session) -> bool:
if self.pipeline_id: if self.pipeline_id:
pipeline = session.scalar(select(Pipeline).where(Pipeline.id == 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 pipeline.is_published
return False 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]]: def get_doc_metadata(self, *, session: Session) -> list[dict[str, str]]:
dataset_metadatas = session.scalars(select(DatasetMetadata).where(DatasetMetadata.dataset_id == self.id)).all() 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 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: class TestDocumentSegmentNeighborAccessors:
"""Regression coverage for ``DocumentSegment.previous_segment`` and ``next_segment`` refactored """Regression coverage for ``DocumentSegment.previous_segment`` and ``next_segment`` refactored
to take a caller-provided session. to take a caller-provided session.