Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
25 changes: 2 additions & 23 deletions api/models/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -243,10 +243,6 @@ def external_retrieval_model(self):
}
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)

Expand All @@ -264,10 +260,6 @@ def get_latest_process_rule(self, *, session: Session) -> "DatasetProcessRule |
.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(
Expand All @@ -278,17 +270,12 @@ def get_app_count(self, *, session: Session) -> int:
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",
Expand Down Expand Up @@ -401,21 +388,13 @@ def get_external_knowledge_info(self, *, session: Session) -> dict[str, Any] | N
"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))
if pipeline:
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()

Expand Down
28 changes: 28 additions & 0 deletions api/tests/unit_tests/models/test_dataset_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -1662,6 +1662,34 @@ def test_accessors_return_none_when_related_rows_are_absent(self, sqlite_session
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.
Expand Down
Loading