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

api/events/event_handlers/clean_when_dataset_deleted.py

@@ -1,19 +1,21 @@
 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
 
 
 @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,
         dataset.tenant_id,
         dataset.indexing_technique,
         dataset.index_struct,
         dataset.collection_binding_id,
-        dataset.doc_form,
+        doc_form,
         dataset.pipeline_id,
     )

api/models/dataset.py

@@ -299,10 +299,6 @@ def available_segment_count(self):
             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 @@ def get_word_count(self, *, session: Session) -> int:
             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 @@ def retrieval_model_dict(self):
 
         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)

api/tests/test_containers_integration_tests/models/test_dataset_models.py

@@ -132,7 +132,7 @@ def test_dataset_word_count_aggregation(self, db_session_with_containers: Sessio
             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."""

api/tests/test_containers_integration_tests/services/test_dataset_service_delete_dataset.py

@@ -83,7 +83,7 @@ def create_document(
         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 @@ def test_delete_dataset_with_documents_success(self, db_session_with_containers:
             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,
         )
 

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,
         )
 

api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py

@@ -255,7 +255,10 @@ def test_batch_clean_document_task_successful_cleanup(
 
         # 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 @@ def test_batch_clean_document_task_with_image_files(
 
         # 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 @@ def test_batch_clean_document_task_no_segments(
 
         # 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 @@ def test_batch_clean_document_task_storage_cleanup_failure(
 
         # 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 @@ def test_batch_clean_document_task_multiple_documents(
 
         # 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 @@ def test_batch_clean_document_task_large_batch_performance(
 
         # 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 @@ def test_batch_clean_document_task_integration_with_real_database(
 
         # 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

api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py

@@ -295,7 +295,7 @@ def test_clean_dataset_task_success_basic_cleanup(
             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 @@ def test_clean_dataset_task_success_with_documents_and_segments(
             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 @@ def test_clean_dataset_task_error_handling_and_rollback(
             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 @@ def test_clean_dataset_task_with_image_file_references(
                 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 @@ def test_clean_dataset_task_performance_with_large_dataset(
             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 @@ def test_clean_dataset_task_storage_exception_handling(
             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 @@ def test_clean_dataset_task_edge_cases_and_boundary_conditions(
             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

api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py

@@ -226,7 +226,9 @@ def test_create_segment_to_index_success(
         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 @@ def test_create_segment_to_index_with_keywords(
         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 @@ def test_create_segment_to_index_with_long_keywords(
         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 @@ def test_create_segment_to_index_with_none_keywords(
         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(