count

Author: brunovcosta

abstra_json_sql/apply.py

@@ -1,6 +1,6 @@
 from typing import List, Dict, Optional
 from .tables import ITablesSnapshot
-from .field_name import field_name
+from .field_name import field_name, expression_name
 from .ast import (
     Expression,
     StringExpression,
@@ -25,6 +25,7 @@
     GreaterThanOrEqualExpression,
     LessThanExpression,
     LessThanOrEqualExpression,
+    WildcardExpression,
     Limit,
 )
 
@@ -157,14 +158,32 @@ def apply_expression(expression: Expression, ctx: dict):
             )
     elif isinstance(expression, FunctionCallExpression):
         if is_aggregate_function(expression.name):
-            raise NotImplementedError(
-                f"Function call expressions are not implemented: {expression.name}"
-            )
+            if expression.name.lower() == "count":
+                assert len(expression.args) == 1, "Count function requires one argument"
+                if isinstance(expression.args[0], WildcardExpression):
+                    return len(ctx["__grouped_rows"])
+                elif isinstance(expression.args[0], NameExpression):
+                    return len(
+                        [
+                            row
+                            for row in ctx["__grouped_rows"]
+                            if expression.args[0].name in row
+                            and row[expression.args[0].name] is not None
+                        ]
+                    )
+            else:
+                raise ValueError(f"Unknown aggregate function: {expression.name}")
         else:
             args = [apply_expression(arg, ctx) for arg in expression.args]
             if expression.name == "lower":
+                assert isinstance(
+                    args[0], str
+                ), "lower function requires a string argument"
                 return args[0].lower()
             elif expression.name == "upper":
+                assert isinstance(
+                    args[0], str
+                ), "upper function requires a string argument"
                 return args[0].upper()
             else:
                 raise ValueError(f"Unknown function: {expression.name}")
@@ -196,14 +215,27 @@ def apply_order_by(order_by: OrderBy, data: List[dict], ctx: dict):
 
 def apply_group_by(group_by: GroupBy, data: List[dict], ctx: dict):
     groups: Dict[tuple, list] = {}
-    for row in data:
+    for idx, row in enumerate(data):
         key = tuple(
             apply_expression(field, {**ctx, **row}) for field in group_by.fields
         )
         if key not in groups:
             groups[key] = []
         groups[key].append(row)
-    return groups
+    if not groups:
+        return [
+            {
+                "__grouped_rows": data,
+                **{expression_name(field): None for field in group_by.fields},
+            }
+        ]
+    return [
+        {
+            "__grouped_rows": rows,
+            **{expression_name(field): key for field, key in zip(group_by.fields, key)},
+        }
+        for key, rows in groups.items()
+    ]
 
 
 def apply_limit(limit: Limit, data: List[dict], ctx: dict):
@@ -212,11 +244,19 @@ def apply_limit(limit: Limit, data: List[dict], ctx: dict):
     return data[start:end]
 
 
+def has_aggregation_fields(fields: List[SelectField]) -> bool:
+    for field in fields:
+        if isinstance(field.expression, FunctionCallExpression):
+            if is_aggregate_function(field.expression.name):
+                return True
+    return False
+
+
 def apply_select_fields(fields: List[SelectField], data: List[dict], ctx: dict):
     return [
         {
             field_name(field) or field.expression: apply_expression(
-                field.expression, {**ctx, **row}
+                field.expression, {**ctx, "__grouped_rows": data, **row}
             )
             for field in fields
         }
@@ -247,12 +287,22 @@ def apply_from(
     return data
 
 
+def has_implicit_aggregation(fields: List[SelectField]) -> bool:
+    for field in fields:
+        if isinstance(field.expression, FunctionCallExpression):
+            if is_aggregate_function(field.expression.name):
+                return True
+    return False
+
+
 def apply_select(select: Select, tables: ITablesSnapshot, ctx: dict):
     data = apply_from(select.from_part, tables, ctx)
     if select.where_part:
         data = apply_where(select.where_part, data, ctx)
     if select.group_part:
         data = apply_group_by(select.group_part, data, ctx)
+    elif has_implicit_aggregation(select.field_parts):
+        data = apply_group_by(GroupBy(fields=[]), data, ctx)
     if select.order_part:
         data = apply_order_by(select.order_part, data, ctx)
     if select.limit_part:

abstra_json_sql/apply_test.py

@@ -202,17 +202,23 @@ def test_group_by(self):
 
         self.assertEqual(
             result,
-            {
-                ("foo",): [
-                    {"name": "Alice", "team": "foo"},
-                    {"name": "Charlie", "team": "foo"},
-                    {"name": "Eve", "team": "foo"},
-                ],
-                ("bar",): [
-                    {"name": "Bob", "team": "bar"},
-                    {"name": "David", "team": "bar"},
-                ],
-            },
+            [
+                {
+                    "team": "foo",
+                    "__grouped_rows": [
+                        {"name": "Alice", "team": "foo"},
+                        {"name": "Charlie", "team": "foo"},
+                        {"name": "Eve", "team": "foo"},
+                    ],
+                },
+                {
+                    "team": "bar",
+                    "__grouped_rows": [
+                        {"name": "Bob", "team": "bar"},
+                        {"name": "David", "team": "bar"},
+                    ],
+                },
+            ],
         )
 
 

abstra_json_sql/ast.py

@@ -101,6 +101,11 @@ class NotEqualExpression(Expression):
     right: Expression
 
 
+@dataclass
+class WildcardExpression(Expression):
+    pass
+
+
 @dataclass
 class Join(Ast):
     table: str
@@ -138,11 +143,6 @@ class OrderBy(Ast):
     fields: List[OrderField]
 
 
-@dataclass
-class SelectWildcard(Ast):
-    pass
-
-
 @dataclass
 class GroupBy(Ast):
     fields: List[Expression]
@@ -156,7 +156,7 @@ class Limit(Ast):
 
 @dataclass
 class Select(Command):
-    field_parts: List[Union[SelectField, SelectWildcard]]
+    field_parts: List[Union[SelectField, WildcardExpression]]
     from_part: Optional[From] = None
     where_part: Optional[Where] = None
     order_part: Optional[OrderBy] = None

abstra_json_sql/eval_test.py

@@ -73,3 +73,22 @@ def test_upper(self):
                 {"upper": "CCC"},
             ],
         )
+
+    def test_count(self):
+        code = "select count(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, [{"count": 3}])

abstra_json_sql/field_name.py

@@ -1,4 +1,4 @@
-from .ast import SelectField, NameExpression, FunctionCallExpression
+from .ast import SelectField, NameExpression, FunctionCallExpression, Expression
 
 
 def field_name(field: SelectField) -> str:
@@ -7,9 +7,17 @@ def field_name(field: SelectField) -> str:
     """
     if field.alias:
         return field.alias
-    elif isinstance(field.expression, NameExpression):
-        return field.expression.name
-    elif isinstance(field.expression, FunctionCallExpression):
-        return field.expression.name
+    else:
+        return expression_name(field.expression)
+
+
+def expression_name(expression: Expression):
+    """
+    Get the field name from an Expression object.
+    """
+    if isinstance(expression, NameExpression):
+        return expression.name
+    elif isinstance(expression, FunctionCallExpression):
+        return expression.name
     else:
         return "?column?"

abstra_json_sql/parser.py

@@ -6,7 +6,7 @@
     SelectField,
     Limit,
     IntExpression,
-    SelectWildcard,
+    WildcardExpression,
     FunctionCallExpression,
     From,
     NameExpression,
@@ -184,7 +184,7 @@ def parse_fields(tokens: List[Token]) -> Tuple[List[SelectField], List[Token]]:
         if tokens[0].type == "keyword" and tokens[0].value.upper() == "FROM":
             break
         if tokens[0].type == "wildcard":
-            fields.append(SelectWildcard())
+            fields.append(WildcardExpression())
             tokens = tokens[1:]
         else:
             exp, tokens = parse_expression(tokens)

abstra_json_sql/parser_test.py

@@ -4,7 +4,7 @@
 from .ast import (
     Select,
     From,
-    SelectWildcard,
+    WildcardExpression,
     PlusExpression,
     FunctionCallExpression,
     SelectField,
@@ -37,7 +37,7 @@ def test_select_wildcard(self):
         self.assertEqual(
             ast,
             Select(
-                field_parts=[SelectWildcard()],
+                field_parts=[WildcardExpression()],
                 from_part=From(
                     table="users",
                 ),