diff --git a/api/models/dataset.py b/api/models/dataset.py
index b15980ae4507ea..796e3d50ab6418 100644
--- a/api/models/dataset.py
+++ b/api/models/dataset.py
@@ -29,7 +29,6 @@
 
 from .account import Account
 from .base import Base, TypeBase
-from .engine import db
 from .enums import (
     CollectionBindingType,
     CreatorUserRole,
@@ -224,10 +223,9 @@ def get_available_document_count(self, *, session: Session) -> int:
             or 0
         )
 
-    @property
-    def available_segment_count(self):
+    def get_available_segment_count(self, *, session: Session) -> int:
         return (
-            db.session.scalar(
+            session.scalar(
                 select(func.count(DocumentSegment.id)).where(
                     DocumentSegment.dataset_id == self.id,
                     DocumentSegment.status == "completed",
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 73097625552394..52d9b42c3563d6 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
@@ -188,7 +188,7 @@ def test_dataset_available_segment_count(self, db_session_with_containers: Sessi
         db_session_with_containers.add(seg_waiting)
         db_session_with_containers.flush()
 
-        assert dataset.available_segment_count == 2
+        assert dataset.get_available_segment_count(session=db_session_with_containers) == 2
 
     def test_document_segment_count_property(self, db_session_with_containers: Session) -> None:
         """Test document can count its segments."""