From d5be01b253f0c927f4a8e36149d64763c395faba Mon Sep 17 00:00:00 2001 From: Keith <148296621+keith991001@users.noreply.github.com> Date: Mon, 7 Sep 2026 09:27:04 +0000 Subject: [PATCH] refactor(models): remove the last standalone legacy db.session properties on Dataset (#41911) --- .../clean_when_dataset_deleted.py | 6 ++-- api/models/dataset.py | 12 ------- .../models/test_dataset_models.py | 2 +- .../test_dataset_service_delete_dataset.py | 4 +-- .../services/test_dataset_service_document.py | 9 +++-- .../tasks/test_batch_clean_document_task.py | 35 +++++++++++++++---- .../tasks/test_clean_dataset_task.py | 14 ++++---- .../test_create_segment_to_index_task.py | 16 ++++++--- 8 files changed, 61 insertions(+), 37 deletions(-) diff --git a/api/events/event_handlers/clean_when_dataset_deleted.py b/api/events/event_handlers/clean_when_dataset_deleted.py index d6007662d8d..ce0cf66f8f1 100644 --- a/api/events/event_handlers/clean_when_dataset_deleted.py +++ b/api/events/event_handlers/clean_when_dataset_deleted.py @@ -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, ) diff --git a/api/models/dataset.py b/api/models/dataset.py index db87aa71fa3..87e6d3e5635 100644 --- a/api/models/dataset.py +++ b/api/models/dataset.py @@ -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) diff --git a/api/tests/test_containers_integration_tests/models/test_dataset_models.py b/api/tests/test_containers_integration_tests/models/test_dataset_models.py index 939e87d8695..0f25a543c9c 100644 --- a/api/tests/test_containers_integration_tests/models/test_dataset_models.py +++ b/api/tests/test_containers_integration_tests/models/test_dataset_models.py @@ -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.""" diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_delete_dataset.py b/api/tests/test_containers_integration_tests/services/test_dataset_service_delete_dataset.py index 1c5ec6835bd..7881cc00733 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service_delete_dataset.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_service_delete_dataset.py @@ -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, ) diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_document.py b/api/tests/test_containers_integration_tests/services/test_dataset_service_document.py index f9a176c1f57..bf53483f8b3 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service_document.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_service_document.py @@ -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, ) diff --git a/api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py b/api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py index 4193223f311..770a145c6cb 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py @@ -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 diff --git a/api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py b/api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py index a31552a09ea..e801597e434 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py @@ -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 diff --git a/api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py b/api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py index a8d295e6a90..0410fb5ebeb 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py @@ -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(