Adding support for limit offset

Author: brunovcosta

json_sql/apply_test.py

@@ -25,7 +25,7 @@
 )
 
 
-class TestEvalExpression(TestCase):
+class TestApplyExpression(TestCase):
     def test_addition(self):
         expression = PlusExpression(
             left=NameExpression(name="a"),
@@ -87,7 +87,7 @@ def test_boolean(self):
         self.assertEqual(result, True)
 
 
-class TestEvalWhere(TestCase):
+class TestApplyWhere(TestCase):
     def test_where_no_context(self):
         where = Where(
             expression=GreaterThanExpression(
@@ -149,7 +149,7 @@ def test_with_context(self):
         )
 
 
-class TestEvalOrderBy(TestCase):
+class TestApplyOrderBy(TestCase):
     def test_order_by(self):
         order_by = OrderBy(
             fields=[
@@ -182,7 +182,7 @@ def test_order_by(self):
         )
 
 
-class TestEvalGroupBy(TestCase):
+class TestApplyGroupBy(TestCase):
     def test_group_by(self):
         group_by = GroupBy(
             fields=[
@@ -216,7 +216,7 @@ def test_group_by(self):
         )
 
 
-class TestEvalLimit(TestCase):
+class TestApplyLimit(TestCase):
     def test_limit(self):
         data = [
             {"name": "Alice", "age": 20},

json_sql/parser.py

@@ -4,6 +4,7 @@
     Ast,
     Select,
     SelectField,
+    Limit,
     IntExpression,
     SelectWildcard,
     From,
@@ -186,6 +187,29 @@ def parse_fields(tokens: List[Token]) -> Tuple[List[SelectField], List[Token]]:
     return fields, tokens
 
 
+def parse_limit(tokens: List[Token]) -> Tuple[Limit, List[Token]]:
+    if not tokens or tokens[0].type != "keyword" or tokens[0].value.upper() != "LIMIT":
+        return None, tokens
+
+    tokens = tokens[1:]
+    limit_token = tokens[0]
+    assert limit_token.type == "int", f"Expected limit value, got {limit_token}"
+    limit_value = limit_token.value
+    limit_int = int(limit_value)
+    tokens = tokens[1:]
+
+    if tokens and tokens[0].type == "keyword" and tokens[0].value.upper() == "OFFSET":
+        tokens = tokens[1:]
+        offset_token = tokens[0]
+        assert offset_token.type == "int", f"Expected offset value, got {offset_token}"
+        offset_value = offset_token.value
+        offset_int = int(offset_value)
+        tokens = tokens[1:]
+        return Limit(limit=limit_int, offset=offset_int), tokens
+    else:
+        return Limit(limit=limit_int), tokens
+
+
 def parse_select(tokens: List[Token]) -> Tuple[Select, List[Token]]:
     if not tokens or tokens[0].type != "keyword" or tokens[0].value.upper() != "SELECT":
         raise ValueError("Expected SELECT statement")
@@ -194,12 +218,14 @@ def parse_select(tokens: List[Token]) -> Tuple[Select, List[Token]]:
     from_part, tokens = parse_from(tokens)
     where_part, tokens = parse_where(tokens)
     order_part, tokens = parse_order(tokens)
+    limit_part, tokens = parse_limit(tokens)
 
     return Select(
         field_parts=field_parts,
         from_part=from_part,
         where_part=where_part,
         order_part=order_part,
+        limit_part=limit_part,
     ), tokens
 
 

json_sql/parser_test.py

@@ -1,5 +1,5 @@
 from unittest import TestCase
-from .parser import parse, parse_expression, parse_from
+from .parser import parse, parse_expression, parse_from, parse_limit
 from .lexer import scan
 from .ast import (
     Select,
@@ -13,6 +13,7 @@
     StringExpression,
     Where,
     OrderBy,
+    Limit,
     OrderField,
 )
 
@@ -159,3 +160,28 @@ def test_with_alias(self):
             ),
         )
         self.assertEqual(tokens, [])
+
+
+class LimitTest(TestCase):
+    def test_limit(self):
+        tokens = scan("LIMIT 10")
+        ast, tokens = parse_limit(tokens)
+        self.assertEqual(
+            ast,
+            Limit(
+                limit=10,
+            ),
+        )
+        self.assertEqual(tokens, [])
+
+    def test_limit_with_offset(self):
+        tokens = scan("LIMIT 10 OFFSET 5")
+        ast, tokens = parse_limit(tokens)
+        self.assertEqual(
+            ast,
+            Limit(
+                limit=10,
+                offset=5,
+            ),
+        )
+        self.assertEqual(tokens, [])

json_sql/tokens.py

@@ -30,6 +30,7 @@
     "IS",
     "NULL",
     "EXISTS",
+    "OFFSET",
     "DISTINCT",
     "ORDER BY",
     "GROUP BY",
@@ -43,6 +44,7 @@
     "CROSS JOIN",
     "NATURAL JOIN",
     "JOIN",
+    "LIMIT",
 ]