Fix pagination issue

Author: Peng Ren

pymongosql/result_set.py

@@ -92,7 +92,7 @@ def _ensure_results_available(self, count: int = 1) -> None:
         while len(self._cached_results) < count and self._cursor_id != 0:
             try:
                 # Use getMore to fetch next batch
-                if self._database and self._execution_plan.collection:
+                if self._database is not None and self._execution_plan.collection:
                     getmore_cmd = {
                         "getMore": self._cursor_id,
                         "collection": self._execution_plan.collection,
@@ -256,12 +256,16 @@ def fetchall(self) -> List[Sequence[Any]]:
         all_results = []
 
         try:
-            # Handle command result (db.command)
-            if not self._cache_exhausted:
-                # Results are already processed in constructor, just extend
-                all_results.extend(self._cached_results)
-                self._total_fetched += len(self._cached_results)
-                self._cache_exhausted = True
+            # Ensure all results are available in cache by requesting a very large number
+            # This will trigger getMore calls until all data is exhausted
+            if not self._cache_exhausted and self._cursor_id != 0:
+                self._ensure_results_available(float("inf"))
+
+            # Now get everything from cache
+            all_results.extend(self._cached_results)
+            self._total_fetched += len(self._cached_results)
+            self._cached_results.clear()
+            self._cache_exhausted = True
 
         except PyMongoError as e:
             self._errors.append({"error": str(e), "type": type(e).__name__})

tests/test_result_set_pagination.py

@@ -0,0 +1,182 @@
+# -*- coding: utf-8 -*-
+from pymongosql.result_set import ResultSet
+from pymongosql.sql.builder import ExecutionPlan
+
+
+class TestResultSetPagination:
+    """Test suite for ResultSet pagination with getMore"""
+
+    # Shared projections used by tests
+    PROJECTION_WITH_FIELDS = {"name": 1, "email": 1}
+    PROJECTION_EMPTY = {}
+
+    def test_pagination_cursor_id_zero(self, conn):
+        """Test pagination when cursor_id is 0 (all results in firstBatch)"""
+        db = conn.database
+        # Query with small limit - all results fit in firstBatch
+        command_result = db.command({"find": "users", "limit": 5})
+
+        execution_plan = ExecutionPlan(collection="users", projection_stage=self.PROJECTION_EMPTY)
+        result_set = ResultSet(command_result=command_result, execution_plan=execution_plan, database=db)
+
+        # Check cursor_id - should be 0 when all results fit in firstBatch
+        assert result_set._cursor_id == 0
+        assert result_set._cache_exhausted is False  # Not exhausted yet, but no getMore needed
+
+        # Fetch all results
+        rows = result_set.fetchall()
+        assert len(rows) == 5
+
+        # After fetching all, cache should be exhausted
+        assert result_set._cache_exhausted is True
+
+    def test_pagination_multiple_batches(self, conn):
+        """Test pagination across multiple batches with getMore"""
+        db = conn.database
+        # Use a small batch size (batchSize) to force pagination
+        command_result = db.command({"find": "users", "batchSize": 5})  # Only 5 results per batch
+
+        execution_plan = ExecutionPlan(collection="users", projection_stage=self.PROJECTION_EMPTY)
+        result_set = ResultSet(command_result=command_result, execution_plan=execution_plan, database=db)
+
+        # Initial results should have cursor_id > 0 since we have 22 total users and batchSize is 5
+        initial_cached = len(result_set._cached_results)
+        assert initial_cached <= 5  # Should have at most 5 in cache from firstBatch
+
+        # Fetch multiple results (should trigger getMore)
+        rows = result_set.fetchmany(10)
+        assert len(rows) == 10
+
+        # After fetching, we should have processed multiple batches
+        assert result_set._total_fetched >= 10
+
+    def test_pagination_ensure_results_available(self, conn):
+        """Test _ensure_results_available with pagination"""
+        db = conn.database
+        # Request results with small batch size
+        command_result = db.command({"find": "users", "batchSize": 3})  # Small batch to test pagination
+
+        execution_plan = ExecutionPlan(collection="users", projection_stage=self.PROJECTION_EMPTY)
+        result_set = ResultSet(command_result=command_result, execution_plan=execution_plan, database=db)
+
+        # Initially, cache might have 3 results
+        initial_cache_size = len(result_set._cached_results)
+        assert initial_cache_size <= 3
+
+        # Ensure we have 8 results available - should trigger getMore
+        result_set._ensure_results_available(8)
+        assert len(result_set._cached_results) >= 8
+
+        # Check that cursor_id was updated
+        assert result_set._cursor_id >= 0
+
+    def test_pagination_fetchone_triggers_getmore(self, conn):
+        """Test that fetchone triggers getMore when needed"""
+        db = conn.database
+        # Create result set with small batch size
+        command_result = db.command({"find": "users", "batchSize": 2})  # Very small batch
+
+        execution_plan = ExecutionPlan(collection="users", projection_stage=self.PROJECTION_WITH_FIELDS)
+        result_set = ResultSet(command_result=command_result, execution_plan=execution_plan, database=db)
+
+        _ = result_set._cursor_id
+        rows_fetched = []
+
+        # Fetch many single rows - should trigger getMore multiple times
+        for _ in range(10):
+            row = result_set.fetchone()
+            if row:
+                rows_fetched.append(row)
+
+        assert len(rows_fetched) == 10
+        # rowcount should reflect total fetched
+        assert result_set.rowcount >= 10
+
+    def test_pagination_cache_exhausted_flag(self, conn):
+        """Test cache exhausted flag is set correctly"""
+        db = conn.database
+        command_result = db.command({"find": "users", "limit": 3})
+
+        execution_plan = ExecutionPlan(collection="users", projection_stage=self.PROJECTION_EMPTY)
+        result_set = ResultSet(command_result=command_result, execution_plan=execution_plan, database=db)
+
+        assert result_set._cache_exhausted is False
+
+        # Fetch all results
+        rows = result_set.fetchall()
+        assert len(rows) == 3
+
+        # After exhausting results, flag should be set
+        assert result_set._cache_exhausted is True
+
+        # Subsequent fetches should return empty
+        more_rows = result_set.fetchall()
+        assert more_rows == []
+
+    def test_pagination_rowcount_tracking(self, conn):
+        """Test rowcount is accurately tracked during pagination"""
+        db = conn.database
+        command_result = db.command({"find": "users", "batchSize": 4})
+
+        execution_plan = ExecutionPlan(collection="users", projection_stage=self.PROJECTION_EMPTY)
+        result_set = ResultSet(command_result=command_result, execution_plan=execution_plan, database=db)
+
+        initial_rowcount = result_set.rowcount
+        assert initial_rowcount <= 4  # Initial batch size
+
+        # Fetch multiple batches
+        batch1 = result_set.fetchmany(8)
+        assert result_set.rowcount >= 8
+
+        batch2 = result_set.fetchmany(5)
+        assert result_set.rowcount >= 13
+
+        # Fetch all remaining
+        all_remaining = result_set.fetchall()
+        _ = result_set.rowcount
+
+        # All 22 users should be fetched eventually
+        total_fetched = len(batch1) + len(batch2) + len(all_remaining)
+        assert total_fetched == 22
+
+    def test_pagination_with_projection(self, conn):
+        """Test pagination with field projection applied"""
+        db = conn.database
+        command_result = db.command({"find": "users", "projection": {"name": 1, "email": 1}, "batchSize": 3})
+
+        execution_plan = ExecutionPlan(collection="users", projection_stage=self.PROJECTION_WITH_FIELDS)
+        result_set = ResultSet(command_result=command_result, execution_plan=execution_plan, database=db)
+
+        # Fetch across multiple batches
+        rows = result_set.fetchall()
+
+        # Should have all 22 users
+        assert len(rows) == 22
+
+        # Each row should have exactly 2 projected fields
+        col_names = [desc[0] for desc in result_set.description]
+        for row in rows:
+            assert len(row) == 2
+            assert isinstance(row[col_names.index("name")], (str, type(None)))
+            assert isinstance(row[col_names.index("email")], (str, type(None)))
+
+    def test_pagination_fetchmany_across_batches(self, conn):
+        """Test fetchmany that spans multiple getMore calls"""
+        db = conn.database
+        command_result = db.command({"find": "users", "batchSize": 3})
+
+        execution_plan = ExecutionPlan(collection="users", projection_stage=self.PROJECTION_EMPTY)
+        result_set = ResultSet(command_result=command_result, execution_plan=execution_plan, database=db)
+
+        # Fetch 10 rows - should span multiple batches
+        batch1 = result_set.fetchmany(10)
+        assert len(batch1) == 10
+
+        # Fetch next 10 - should get more users
+        batch2 = result_set.fetchmany(10)
+        assert len(batch2) == 10
+
+        # Fetch remaining results
+        batch3 = result_set.fetchmany(5)
+        # Should get remaining users (total > 20, depends on actual data size)
+        assert len(batch3) > 0