functions

Author: brunovcosta

abstra_json_sql/apply.py

@@ -29,6 +29,11 @@
 )
 
 
+def is_aggregate_function(name: str) -> bool:
+    # Placeholder for aggregate function check
+    return name.lower() in ["sum", "avg", "count", "min", "max"]
+
+
 def apply_expression(expression: Expression, ctx: dict):
     if isinstance(expression, StringExpression):
         return expression.value
@@ -151,9 +156,18 @@ def apply_expression(expression: Expression, ctx: dict):
                 f"Unsupported types for less than or equal: {type(left_value)}, {type(right_value)}"
             )
     elif isinstance(expression, FunctionCallExpression):
-        raise NotImplementedError(
-            f"Function call expressions are not implemented: {expression.name}"
-        )
+        if is_aggregate_function(expression.name):
+            raise NotImplementedError(
+                f"Function call expressions are not implemented: {expression.name}"
+            )
+        else:
+            args = [apply_expression(arg, ctx) for arg in expression.args]
+            if expression.name == "lower":
+                return args[0].lower()
+            elif expression.name == "upper":
+                return args[0].upper()
+            else:
+                raise ValueError(f"Unknown function: {expression.name}")
     else:
         raise ValueError(f"Unsupported expression type: {type(expression)}")
 

abstra_json_sql/eval_test.py

@@ -1,10 +1,10 @@
 from unittest import TestCase
 from .eval import eval_sql
-from .tables import InMemoryTables, Table
+from .tables import InMemoryTables, Table, Column
 
 
 class TestEvalSQL(TestCase):
-    def test_eval_sql(self):
+    def test_sql(self):
         code = "select 1+1"
         tables = InMemoryTables(
             tables=[],
@@ -13,7 +13,7 @@ def test_eval_sql(self):
         result = eval_sql(code=code, tables=tables, ctx=ctx)
         self.assertEqual(result, [{"?column?": 2}])
 
-    def test_eval_select_alias(self):
+    def test_select_alias(self):
         code = "select 1+1 as a"
         tables = InMemoryTables(
             tables=[],
@@ -22,21 +22,54 @@ def test_eval_select_alias(self):
         result = eval_sql(code=code, tables=tables, ctx=ctx)
         self.assertEqual(result, [{"a": 2}])
 
-    def test_eval_aggregate_sum(self):
-        code = "select sum(foo) from bar"
+    def test_lower(self):
+        code = "select lower(foo) from bar"
         tables = InMemoryTables(
             tables=[
                 Table(
                     name="bar",
-                    columns=["foo"],
-                    rows=[
-                        {"foo": 1},
-                        {"foo": 2},
-                        {"foo": 3},
+                    columns=[Column(name="foo", type="text")],
+                    data=[
+                        {"foo": "AAA"},
+                        {"foo": "BBB"},
+                        {"foo": "CCC"},
                     ],
                 )
             ],
         )
         ctx = {}
         result = eval_sql(code=code, tables=tables, ctx=ctx)
-        self.assertEqual(result, [{"a": 6}])
+        self.assertEqual(
+            result,
+            [
+                {"lower": "aaa"},
+                {"lower": "bbb"},
+                {"lower": "ccc"},
+            ],
+        )
+
+    def test_upper(self):
+        code = "select upper(foo) from bar"
+        tables = InMemoryTables(
+            tables=[
+                Table(
+                    name="bar",
+                    columns=[Column(name="foo", type="text")],
+                    data=[
+                        {"foo": "aaa"},
+                        {"foo": "bbb"},
+                        {"foo": "ccc"},
+                    ],
+                )
+            ],
+        )
+        ctx = {}
+        result = eval_sql(code=code, tables=tables, ctx=ctx)
+        self.assertEqual(
+            result,
+            [
+                {"upper": "AAA"},
+                {"upper": "BBB"},
+                {"upper": "CCC"},
+            ],
+        )

abstra_json_sql/field_name.py

@@ -1,4 +1,4 @@
-from .ast import SelectField, NameExpression
+from .ast import SelectField, NameExpression, FunctionCallExpression
 
 
 def field_name(field: SelectField) -> str:
@@ -9,5 +9,7 @@ def field_name(field: SelectField) -> str:
         return field.alias
     elif isinstance(field.expression, NameExpression):
         return field.expression.name
+    elif isinstance(field.expression, FunctionCallExpression):
+        return field.expression.name
     else:
         return "?column?"

abstra_json_sql/parser.py

@@ -79,14 +79,14 @@ def parse_expression(tokens: List[Token]) -> Tuple[Expression, List[Token]]:
                 args = []
                 while True:
                     param_expression, tokens = parse_expression(tokens)
+                    args.append(param_expression)
                     if tokens and tokens[0].type == "comma":
                         tokens = tokens[1:]
                     elif tokens and tokens[0].type == "paren_right":
                         tokens = tokens[1:]
                         break
                     else:
                         raise ValueError("Expected comma or closing parenthesis")
-                    args.append(param_expression)
 
                 stack.append(FunctionCallExpression(name=name_value, args=args))
             else:

abstra_json_sql/parser_test.py

@@ -6,6 +6,7 @@
     From,
     SelectWildcard,
     PlusExpression,
+    FunctionCallExpression,
     SelectField,
     IntExpression,
     NameExpression,
@@ -136,6 +137,18 @@ def test_equal_expression(self):
         )
         self.assertEqual(tokens, [])
 
+    def test_function_call_expression(self):
+        tokens = scan("SUM(foo)")
+        ast, tokens = parse_expression(tokens)
+        self.assertEqual(
+            ast,
+            FunctionCallExpression(
+                name="SUM",
+                args=[NameExpression(name="foo")],
+            ),
+        )
+        self.assertEqual(tokens, [])
+
 
 class FromTest(TestCase):
     def test_simple(self):