Clean up the code

Author: Peng Ren

pymongosql/cursor.py

@@ -2,7 +2,6 @@
 import logging
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, TypeVar
 
-from pymongo.cursor import Cursor as MongoCursor
 from pymongo.errors import PyMongoError
 
 from .common import BaseCursor, CursorIterator
@@ -32,16 +31,15 @@ def __init__(self, connection: "Connection", **kwargs) -> None:
         self._result_set: Optional[ResultSet] = None
         self._result_set_class = ResultSet
         self._current_execution_plan: Optional[ExecutionPlan] = None
-        self._mongo_cursor: Optional[MongoCursor] = None
         self._is_closed = False
 
     @property
     def result_set(self) -> Optional[ResultSet]:
         return self._result_set
 
     @result_set.setter
-    def result_set(self, val: ResultSet) -> None:
-        self._result_set = val
+    def result_set(self, rs: ResultSet) -> None:
+        self._result_set = rs
 
     @property
     def has_result_set(self) -> bool:
@@ -52,8 +50,8 @@ def result_set_class(self) -> Optional[type]:
         return self._result_set_class
 
     @result_set_class.setter
-    def result_set_class(self, val: type) -> None:
-        self._result_set_class = val
+    def result_set_class(self, rs_cls: type) -> None:
+        self._result_set_class = rs_cls
 
     @property
     def rowcount(self) -> int:
@@ -107,7 +105,7 @@ def _execute_execution_plan(self, execution_plan: ExecutionPlan) -> None:
             # Build MongoDB find command
             find_command = {"find": execution_plan.collection, "filter": execution_plan.filter_stage or {}}
 
-            # Apply projection if specified (already in MongoDB format)
+            # Apply projection if specified
             if execution_plan.projection_stage:
                 find_command["projection"] = execution_plan.projection_stage
 
@@ -236,15 +234,6 @@ def fetchall(self) -> List[Sequence[Any]]:
     def close(self) -> None:
         """Close the cursor and free resources"""
         try:
-            if self._mongo_cursor:
-                # Close MongoDB cursor
-                try:
-                    self._mongo_cursor.close()
-                except Exception as e:
-                    _logger.warning(f"Error closing MongoDB cursor: {e}")
-                finally:
-                    self._mongo_cursor = None
-
             if self._result_set:
                 # Close result set
                 try:

pymongosql/sql/ast.py

@@ -154,3 +154,20 @@ def visitLimitClause(self, ctx: PartiQLParser.LimitClauseContext) -> Any:
         except Exception as e:
             _logger.warning(f"Error processing LIMIT clause: {e}")
             return self.visitChildren(ctx)
+
+    def visitOffsetByClause(self, ctx: PartiQLParser.OffsetByClauseContext) -> Any:
+        """Handle OFFSET clause for result skipping"""
+        _logger.debug("Processing OFFSET clause")
+        try:
+            if hasattr(ctx, "exprSelect") and ctx.exprSelect():
+                offset_text = ctx.exprSelect().getText()
+                try:
+                    offset_value = int(offset_text)
+                    self._parse_result.offset_value = offset_value
+                    _logger.debug(f"Extracted offset value: {offset_value}")
+                except ValueError as e:
+                    _logger.warning(f"Invalid OFFSET value '{offset_text}': {e}")
+            return self.visitChildren(ctx)
+        except Exception as e:
+            _logger.warning(f"Error processing OFFSET clause: {e}")
+            return self.visitChildren(ctx)

tests/test_cursor.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 import pytest
 
-from pymongosql.error import ProgrammingError
+from pymongosql.error import DatabaseError, ProgrammingError, SqlSyntaxError
 from pymongosql.result_set import ResultSet
 
 
@@ -64,9 +64,7 @@ def test_execute_with_limit(self, conn):
         assert isinstance(cursor.result_set, ResultSet)
         rows = cursor.result_set.fetchall()
 
-        # Should return results from 22 users in dataset (LIMIT parsing may not be implemented yet)
-        # TODO: Fix LIMIT parsing in SQL grammar
-        assert len(rows) >= 1  # At least we get some results
+        assert len(rows) == 2  # At least we get some results
 
         # Check that names are present using DB API 2.0
         if len(rows) > 0:
@@ -85,7 +83,7 @@ def test_execute_with_skip(self, conn):
         rows = cursor.result_set.fetchall()
 
         # Should return users after skipping 1 (from 22 users in dataset)
-        assert len(rows) >= 0  # Could be 0-21 depending on implementation
+        assert len(rows) == 21  # 22 - 1 = 21 users after skipping the first one
 
         # Check that results have name field if any results using DB API 2.0
         if len(rows) > 0:
@@ -111,11 +109,8 @@ def test_execute_with_sort(self, conn):
         assert "name" in col_names
         assert all(len(row) >= 1 for row in rows)  # All rows should have data
 
-        # Verify that we have actual user names from the dataset using DB API 2.0
-        if "name" in col_names:
-            name_idx = col_names.index("name")
-            names = [row[name_idx] for row in rows]
-            assert "John Doe" in names  # First user from dataset
+        # Verify that the first name in the result
+        assert "Patricia Johnson" == rows[0][0]
 
     def test_execute_complex_query(self, conn):
         """Test executing complex query with multiple clauses"""
@@ -140,7 +135,7 @@ def test_execute_complex_query(self, conn):
 
     def test_execute_nested_fields_query(self, conn):
         """Test executing query with nested field access"""
-        sql = "SELECT profile.bio, address.city, address.coordinates FROM users WHERE salary >= 100000"
+        sql = "SELECT name, profile.bio, address.city FROM users WHERE salary >= 100000 ORDER BY salary DESC"
 
         cursor = conn.cursor()
         result = cursor.execute(sql)
@@ -156,17 +151,20 @@ def test_execute_nested_fields_query(self, conn):
         if cursor.result_set.description:
             col_names = [desc[0] for desc in cursor.result_set.description]
             # Should include nested field names in projection
+            assert "name" in col_names
             assert "profile.bio" in col_names
             assert "address.city" in col_names
-            assert "address.coordinates" in col_names
+
+        # Verify the first record matched the highest salary
+        assert "Patricia Johnson" == rows[0][0]
 
     def test_execute_parser_error(self, conn):
         """Test executing query with parser errors"""
         sql = "INVALID SQL SYNTAX"
 
         # This should raise an exception due to invalid SQL
         cursor = conn.cursor()
-        with pytest.raises(Exception):  # Could be SqlSyntaxError or other parsing error
+        with pytest.raises(SqlSyntaxError):  # Could be SqlSyntaxError or other parsing error
             cursor.execute(sql)
 
     def test_execute_database_error(self, conn, make_connection):
@@ -178,7 +176,7 @@ def test_execute_database_error(self, conn, make_connection):
 
         # This should raise an exception due to closed connection
         cursor = conn.cursor()
-        with pytest.raises(Exception):  # Could be DatabaseError or OperationalError
+        with pytest.raises(DatabaseError):
             cursor.execute(sql)
 
         # Reconnect for other tests
@@ -188,31 +186,6 @@ def test_execute_database_error(self, conn, make_connection):
         finally:
             new_conn.close()
 
-    def test_execute_with_aliases(self, conn):
-        """Test executing query with field aliases"""
-        sql = "SELECT name AS full_name, email AS user_email FROM users"
-        cursor = conn.cursor()
-        result = cursor.execute(sql)
-
-        assert result == cursor  # execute returns self
-        assert isinstance(cursor.result_set, ResultSet)
-        rows = cursor.result_set.fetchall()
-
-        # Should return users with aliased field names
-        assert len(rows) == 22
-
-        # Check that alias fields are present if aliasing works using DB API 2.0
-        col_names = [desc[0] for desc in cursor.result_set.description]
-        # Aliases might not work yet, so check for either original or alias names
-        assert "name" in col_names or "full_name" in col_names
-        # Check for email columns in description
-        has_email = "email" in col_names or "user_email" in col_names
-        for row in rows:
-            assert len(row) >= 2  # Should have at least 2 columns
-        # Verify we have email data if expected
-        if has_email:
-            assert True  # Email column exists in description
-
     def test_fetchone_without_execute(self, conn):
         """Test fetchone without previous execute"""
         fresh_cursor = conn.cursor()
@@ -238,11 +211,10 @@ def test_fetchone_with_result(self, conn):
         # Execute query first
         cursor = conn.cursor()
         _ = cursor.execute(sql)
-
-        # Test fetchone - DB API 2.0 returns sequences, not dicts
         row = cursor.fetchone()
+
         assert row is not None
-        assert isinstance(row, (tuple, list))  # Should be sequence, not dict
+        assert isinstance(row, (tuple, list))
         # Verify we have data using DB API 2.0 approach
         col_names = [desc[0] for desc in cursor.result_set.description] if cursor.result_set.description else []
         if "name" in col_names:

tests/test_sql_parser_general.py

@@ -239,6 +239,29 @@ def test_select_with_limit(self):
         assert execution_plan.limit_stage == 10
         assert execution_plan.projection_stage == {"name": 1}
 
+    def test_select_with_offset(self):
+        """Test SELECT with OFFSET clause"""
+        sql = "SELECT name FROM users OFFSET 5"
+        parser = SQLParser(sql)
+
+        assert not parser.has_errors, f"Parser errors: {parser.errors}"
+        execution_plan = parser.get_execution_plan()
+        assert execution_plan.collection == "users"
+        assert execution_plan.skip_stage == 5
+        assert execution_plan.projection_stage == {"name": 1}
+
+    def test_select_with_limit_and_offset(self):
+        """Test SELECT with both LIMIT and OFFSET clauses"""
+        sql = "SELECT name, email FROM users LIMIT 10 OFFSET 5"
+        parser = SQLParser(sql)
+
+        assert not parser.has_errors, f"Parser errors: {parser.errors}"
+        execution_plan = parser.get_execution_plan()
+        assert execution_plan.collection == "users"
+        assert execution_plan.limit_stage == 10
+        assert execution_plan.skip_stage == 5
+        assert execution_plan.projection_stage == {"name": 1, "email": 1}
+
     def test_complex_query_combination(self):
         """Test complex query with multiple clauses"""
         sql = """
@@ -277,83 +300,6 @@ def test_parser_error_handling(self):
             parser = SQLParser("INVALID SQL SYNTAX")
             parser.get_execution_plan()
 
-    def test_select_with_as_aliases(self):
-        """Test SELECT with AS aliases"""
-        sql = "SELECT name AS username, email AS user_email FROM customers"
-        parser = SQLParser(sql)
-
-        assert not parser.has_errors, f"Parser errors: {parser.errors}"
-
-        execution_plan = parser.get_execution_plan()
-        assert execution_plan.collection == "customers"
-        assert execution_plan.filter_stage == {}
-        assert execution_plan.projection_stage == {
-            "name": 1,
-            "email": 1,
-        }
-
-    def test_select_with_mixed_aliases(self):
-        """Test SELECT with mixed alias formats"""
-        sql = "SELECT name AS username, age user_age, status FROM users"
-        parser = SQLParser(sql)
-
-        assert not parser.has_errors, f"Parser errors: {parser.errors}"
-
-        execution_plan = parser.get_execution_plan()
-        assert execution_plan.collection == "users"
-        assert execution_plan.filter_stage == {}
-        assert execution_plan.projection_stage == {
-            "name": 1,  # AS alias
-            "age": 1,  # Space-separated alias
-            "status": 1,  # No alias (field included)
-        }
-
-    def test_select_with_space_separated_aliases(self):
-        """Test SELECT with space-separated aliases"""
-        sql = "SELECT first_name fname, last_name lname, created_at creation_date FROM users"
-        parser = SQLParser(sql)
-
-        assert not parser.has_errors, f"Parser errors: {parser.errors}"
-
-        execution_plan = parser.get_execution_plan()
-        assert execution_plan.collection == "users"
-        assert execution_plan.filter_stage == {}
-        assert execution_plan.projection_stage == {
-            "first_name": 1,
-            "last_name": 1,
-            "created_at": 1,
-        }
-
-    def test_select_with_complex_field_names_and_aliases(self):
-        """Test SELECT with complex field names and aliases"""
-        sql = "SELECT user_profile.name AS display_name, account_settings.theme user_theme FROM users"
-        parser = SQLParser(sql)
-
-        assert not parser.has_errors, f"Parser errors: {parser.errors}"
-
-        execution_plan = parser.get_execution_plan()
-        assert execution_plan.collection == "users"
-        assert execution_plan.filter_stage == {}
-        assert execution_plan.projection_stage == {
-            "user_profile.name": 1,
-            "account_settings.theme": 1,
-        }
-
-    def test_select_function_with_aliases(self):
-        """Test SELECT with functions and aliases"""
-        sql = "SELECT COUNT(*) AS total_count, MAX(age) max_age FROM users"
-        parser = SQLParser(sql)
-
-        assert not parser.has_errors, f"Parser errors: {parser.errors}"
-
-        execution_plan = parser.get_execution_plan()
-        assert execution_plan.collection == "users"
-        assert execution_plan.filter_stage == {}
-        assert execution_plan.projection_stage == {
-            "COUNT(*)": 1,
-            "MAX(age)": 1,
-        }
-
     def test_superset_wrapped_subquery(self):
         """Support Superset wrapping subquery with alias 'virtual'"""
         sql = "SELECT virtual.a, virtual.b FROM (SELECT a, b FROM users WHERE c = 1) virtual"
@@ -369,49 +315,6 @@ def test_superset_wrapped_subquery(self):
         # Inner filter should be preserved
         assert execution_plan.filter_stage == {"c": 1}
 
-    def test_select_single_field_with_alias(self):
-        """Test SELECT with single field and alias"""
-        sql = "SELECT email AS contact_email FROM customers"
-        parser = SQLParser(sql)
-
-        assert not parser.has_errors, f"Parser errors: {parser.errors}"
-
-        execution_plan = parser.get_execution_plan()
-        assert execution_plan.collection == "customers"
-        assert execution_plan.filter_stage == {}
-        assert execution_plan.projection_stage == {"email": 1}
-
-    def test_select_aliases_with_where_clause(self):
-        """Test SELECT with aliases and WHERE clause"""
-        sql = "SELECT name AS username, status AS account_status FROM users WHERE age > 18"
-        parser = SQLParser(sql)
-
-        assert not parser.has_errors, f"Parser errors: {parser.errors}"
-
-        execution_plan = parser.get_execution_plan()
-        assert execution_plan.collection == "users"
-        assert execution_plan.filter_stage == {"age": {"$gt": 18}}
-        assert execution_plan.projection_stage == {
-            "name": 1,
-            "status": 1,
-        }
-
-    def test_select_case_insensitive_as_alias(self):
-        """Test SELECT with case insensitive AS keyword"""
-        sql = "SELECT name as username, email As user_email, status AS account_status FROM users"
-        parser = SQLParser(sql)
-
-        assert not parser.has_errors, f"Parser errors: {parser.errors}"
-
-        execution_plan = parser.get_execution_plan()
-        assert execution_plan.collection == "users"
-        assert execution_plan.filter_stage == {}
-        assert execution_plan.projection_stage == {
-            "name": 1,
-            "email": 1,
-            "status": 1,
-        }
-
     def test_different_collection_names(self):
         """Test parsing with different collection names"""
         test_cases = [