opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-08 11:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(models): remove the last standalone legacy db.session properties on Dataset (#41911)

Browse Source
This commit is contained in:
Keith 2026-09-07 09:27:04 +00:00 committed by GitHub
parent b93101459d
commit d5be01b253
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 61 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
api/events/event_handlers/clean_when_dataset_deleted.py
View File

@ -1,4 +1,5 @@
from events.dataset_event import dataset_was_deleted
from extensions.ext_database import db
from models import Dataset
from tasks.clean_dataset_task import clean_dataset_task
@ -6,7 +7,8 @@ from tasks.clean_dataset_task import clean_dataset_task
@dataset_was_deleted.connect
def handle(sender: Dataset, **kwargs):
dataset = sender
if not dataset.doc_form or not dataset.indexing_technique:
doc_form = dataset.get_doc_form(session=db.session())
if not doc_form or not dataset.indexing_technique:
return
clean_dataset_task.delay(
dataset.id,
@ -14,6 +16,6 @@ def handle(sender: Dataset, **kwargs):
dataset.indexing_technique,
dataset.index_struct,
dataset.collection_binding_id,
dataset.doc_form,
doc_form,
dataset.pipeline_id,
)

12
api/models/dataset.py
View File

@ -299,10 +299,6 @@ class Dataset(Base):
or 0
)
@property
def word_count(self) -> int:
return self.get_word_count(session=db.session())
def get_word_count(self, *, session: Session) -> int:
return (
session.scalar(
@ -311,10 +307,6 @@ class Dataset(Base):
or 0
)
@property
def doc_form(self) -> str | None:
return self.get_doc_form(session=db.session())
def get_doc_form(self, *, session: Session) -> str | None:
if self.chunk_structure:
return self.chunk_structure
@ -344,10 +336,6 @@ class Dataset(Base):
return {**default_retrieval_model, **self.retrieval_model}
@property
def tags(self) -> Sequence[Tag]:
return self.get_tags(session=db.session())
def get_tags(self, *, session: Session) -> Sequence[Tag]:
tags = session.scalars(
select(Tag)

2
api/tests/test_containers_integration_tests/models/test_dataset_models.py
View File

@ -132,7 +132,7 @@ class TestDatasetDocumentProperties:
db_session_with_containers.add(doc)
db_session_with_containers.flush()
assert dataset.word_count == 5000
assert dataset.get_word_count(session=db_session_with_containers) == 5000
def test_dataset_available_segment_count(self, db_session_with_containers: Session) -> None:
"""Test Dataset.available_segment_count counts completed and enabled segments."""

4
api/tests/test_containers_integration_tests/services/test_dataset_service_delete_dataset.py
View File

@ -83,7 +83,7 @@ class DatasetDeleteIntegrationDataFactory:
created_by: str,
doc_form: str = IndexStructureType.PARAGRAPH_INDEX,
) -> Document:
"""Persist a document so dataset.doc_form resolves through the real document path."""
"""Persist a document so dataset.get_doc_form resolves through the real document path."""
document = Document(
tenant_id=tenant_id,
dataset_id=dataset_id,
@ -142,7 +142,7 @@ class TestDatasetServiceDeleteDataset:
dataset.indexing_technique,
dataset.index_struct,
dataset.collection_binding_id,
dataset.doc_form,
dataset.get_doc_form(session=db_session_with_containers),
dataset.pipeline_id,
)

9
api/tests/test_containers_integration_tests/services/test_dataset_service_document.py
View File

@ -624,7 +624,12 @@ def test_delete_documents_ignores_empty_input(db_session_with_containers: Sessio
dataset_ref = DatasetRefService.create_dataset_ref(dataset)
with patch("services.dataset_service.batch_clean_document_task.delay") as delay:
DocumentService.delete_documents(dataset_ref, [], dataset.doc_form, session=db_session_with_containers)
DocumentService.delete_documents(
dataset_ref,
[],
dataset.get_doc_form(session=db_session_with_containers),
session=db_session_with_containers,
)
delay.assert_not_called()
@ -662,7 +667,7 @@ def test_delete_documents_deletes_rows_and_dispatches_cleanup_task(db_session_wi
DocumentService.delete_documents(
dataset_ref,
[document_a.id, document_b.id],
dataset.doc_form,
dataset.get_doc_form(session=db_session_with_containers),
session=db_session_with_containers,
)

35
api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py
View File

@ -255,7 +255,10 @@ class TestBatchCleanDocumentTask:
# Execute the task
batch_clean_document_task(
document_ids=[document_id], dataset_id=dataset.id, doc_form=dataset.doc_form, file_ids=[file_id]
document_ids=[document_id],
dataset_id=dataset.id,
doc_form=dataset.get_doc_form(session=db_session_with_containers),
file_ids=[file_id],
)
# Verify that the task completed successfully
@ -311,7 +314,10 @@ class TestBatchCleanDocumentTask:
# Execute the task
batch_clean_document_task(
document_ids=[document_id], dataset_id=dataset.id, doc_form=dataset.doc_form, file_ids=[]
document_ids=[document_id],
dataset_id=dataset.id,
doc_form=dataset.get_doc_form(session=db_session_with_containers),
file_ids=[],
)
# Verify database cleanup
@ -351,7 +357,10 @@ class TestBatchCleanDocumentTask:
# Execute the task
batch_clean_document_task(
document_ids=[document_id], dataset_id=dataset.id, doc_form=dataset.doc_form, file_ids=[file_id]
document_ids=[document_id],
dataset_id=dataset.id,
doc_form=dataset.get_doc_form(session=db_session_with_containers),
file_ids=[file_id],
)
# Verify that the task completed successfully
@ -446,7 +455,10 @@ class TestBatchCleanDocumentTask:
# Execute the task
batch_clean_document_task(
document_ids=[document_id], dataset_id=dataset.id, doc_form=dataset.doc_form, file_ids=[file_id]
document_ids=[document_id],
dataset_id=dataset.id,
doc_form=dataset.get_doc_form(session=db_session_with_containers),
file_ids=[file_id],
)
# Verify that the task completed successfully despite storage failure
@ -504,7 +516,10 @@ class TestBatchCleanDocumentTask:
# Execute the task with multiple documents
batch_clean_document_task(
document_ids=document_ids, dataset_id=dataset.id, doc_form=dataset.doc_form, file_ids=file_ids
document_ids=document_ids,
dataset_id=dataset.id,
doc_form=dataset.get_doc_form(session=db_session_with_containers),
file_ids=file_ids,
)
# Verify that the task completed successfully for all documents
@ -641,7 +656,10 @@ class TestBatchCleanDocumentTask:
# Execute the task with large batch
batch_clean_document_task(
document_ids=document_ids, dataset_id=dataset.id, doc_form=dataset.doc_form, file_ids=file_ids
document_ids=document_ids,
dataset_id=dataset.id,
doc_form=dataset.get_doc_form(session=db_session_with_containers),
file_ids=file_ids,
)
end_time = time.perf_counter()
@ -734,7 +752,10 @@ class TestBatchCleanDocumentTask:
# Execute the task
batch_clean_document_task(
document_ids=[document_id], dataset_id=dataset.id, doc_form=dataset.doc_form, file_ids=[file_id]
document_ids=[document_id],
dataset_id=dataset.id,
doc_form=dataset.get_doc_form(session=db_session_with_containers),
file_ids=[file_id],
)
# Verify that the task completed successfully

14
api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py
View File

@ -295,7 +295,7 @@ class TestCleanDatasetTask:
indexing_technique=dataset.indexing_technique,
index_struct=dataset.index_struct,
collection_binding_id=dataset.collection_binding_id,
doc_form=dataset.doc_form,
doc_form=dataset.get_doc_form(session=db_session_with_containers),
)
# Verify results
@ -419,7 +419,7 @@ class TestCleanDatasetTask:
indexing_technique=dataset.indexing_technique,
index_struct=dataset.index_struct,
collection_binding_id=dataset.collection_binding_id,
doc_form=dataset.doc_form,
doc_form=dataset.get_doc_form(session=db_session_with_containers),
)
# Verify results
@ -552,7 +552,7 @@ class TestCleanDatasetTask:
indexing_technique=dataset.indexing_technique,
index_struct=dataset.index_struct,
collection_binding_id=dataset.collection_binding_id,
doc_form=dataset.doc_form,
doc_form=dataset.get_doc_form(session=db_session_with_containers),
)
# Verify results - even with vector cleanup failure, documents and segments should be deleted
@ -638,7 +638,7 @@ class TestCleanDatasetTask:
indexing_technique=dataset.indexing_technique,
index_struct=dataset.index_struct,
collection_binding_id=dataset.collection_binding_id,
doc_form=dataset.doc_form,
doc_form=dataset.get_doc_form(session=db_session_with_containers),
)
# Verify results
@ -751,7 +751,7 @@ class TestCleanDatasetTask:
indexing_technique=dataset.indexing_technique,
index_struct=dataset.index_struct,
collection_binding_id=dataset.collection_binding_id,
doc_form=dataset.doc_form,
doc_form=dataset.get_doc_form(session=db_session_with_containers),
)
end_time = time.time()
@ -845,7 +845,7 @@ class TestCleanDatasetTask:
indexing_technique=dataset.indexing_technique,
index_struct=dataset.index_struct,
collection_binding_id=dataset.collection_binding_id,
doc_form=dataset.doc_form,
doc_form=dataset.get_doc_form(session=db_session_with_containers),
)
# Verify results
@ -999,7 +999,7 @@ class TestCleanDatasetTask:
indexing_technique=dataset.indexing_technique,
index_struct=dataset.index_struct,
collection_binding_id=dataset.collection_binding_id,
doc_form=dataset.doc_form,
doc_form=dataset.get_doc_form(session=db_session_with_containers),
)
# Verify results

16
api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py
View File

@ -226,7 +226,9 @@ class TestCreateSegmentToIndexTask:
assert segment.error is None
# Verify index processor was called
mock_external_service_dependencies["index_processor_factory"].assert_called_once_with(dataset.doc_form)
mock_external_service_dependencies["index_processor_factory"].assert_called_once_with(
dataset.get_doc_form(session=db_session_with_containers)
)
mock_external_service_dependencies["index_processor"].load.assert_called_once()
# Verify Redis cache cleanup
@ -552,7 +554,9 @@ class TestCreateSegmentToIndexTask:
assert segment.completed_at is not None
# Verify index processor was called
mock_external_service_dependencies["index_processor_factory"].assert_called_once_with(dataset.doc_form)
mock_external_service_dependencies["index_processor_factory"].assert_called_once_with(
dataset.get_doc_form(session=db_session_with_containers)
)
mock_external_service_dependencies["index_processor"].load.assert_called_once()
def test_create_segment_to_index_different_doc_forms(
@ -983,7 +987,9 @@ class TestCreateSegmentToIndexTask:
assert segment.completed_at is not None
# Verify index processor was called
mock_external_service_dependencies["index_processor_factory"].assert_called_once_with(dataset.doc_form)
mock_external_service_dependencies["index_processor_factory"].assert_called_once_with(
dataset.get_doc_form(session=db_session_with_containers)
)
mock_external_service_dependencies["index_processor"].load.assert_called_once()
def test_create_segment_to_index_tenant_isolation(
@ -1057,7 +1063,9 @@ class TestCreateSegmentToIndexTask:
assert segment.completed_at is not None
# Verify index processor was called
mock_external_service_dependencies["index_processor_factory"].assert_called_once_with(dataset.doc_form)
mock_external_service_dependencies["index_processor_factory"].assert_called_once_with(
dataset.get_doc_form(session=db_session_with_containers)
)
mock_external_service_dependencies["index_processor"].load.assert_called_once()
def test_create_segment_to_index_comprehensive_integration(