fix: resolve aggregate ORDER BY parsing, HAVING context, and ORDER BY scope issues

Author: yeongseon

docs/SQL_SPEC.md

@@ -147,10 +147,14 @@ See [Section 7: WHERE Clause](#7-where-clause).
 |---------|-----------|---------|
 | Single column | ✅ | `ORDER BY name ASC` |
 | Direction | ✅ | `ASC` (default) or `DESC` |
-| Multi-column | ❌ | `ORDER BY name, age` — only first column used |
+| Multi-column | ❌ | `ORDER BY name, age` — not supported |
 | Expressions | ❌ | `ORDER BY UPPER(name)` — not supported |
 
-- Column must exist in the worksheet headers.
+- Non-aggregate queries: ORDER BY column must exist in worksheet headers.
+- Aggregate queries (`GROUP BY` and/or aggregate SELECT): ORDER BY may reference:
+  - projected output columns,
+  - GROUP BY columns (even if not projected),
+  - aggregate expressions (e.g., `COUNT(*)`, `SUM(score)`).
 - NULL values sort last (after all non-NULL values).
 - Numeric strings are compared numerically when both operands parse as numbers.
 

src/excel_dbapi/executor.py

@@ -326,14 +326,17 @@ def _execute_aggregate_select(
                 )
 
         output_columns: list[str] = []
+        required_aggregates: dict[str, tuple[str, str]] = {}
         for column in columns:
             if self._is_aggregate_column(column):
                 arg = str(column.get("arg"))
                 if arg != "*" and arg not in headers:
                     raise ValueError(
                         f"Unknown column: {arg}. Available columns: {headers}"
                     )
-                output_columns.append(self._aggregate_label(column))
+                label = self._aggregate_label(column)
+                required_aggregates[label] = (str(column.get("func")), arg)
+                output_columns.append(label)
                 continue
 
             column_name = str(column)
@@ -351,6 +354,24 @@ def _execute_aggregate_select(
                 )
             output_columns.append(column_name)
 
+        for clause in (having, parsed.get("order_by")):
+            if clause is None:
+                continue
+            if "conditions" in clause:
+                refs = [str(condition["column"]) for condition in clause["conditions"]]
+            else:
+                refs = [str(clause["column"])]
+            for ref in refs:
+                aggregate_spec = self._aggregate_spec_from_label(ref)
+                if aggregate_spec is None:
+                    continue
+                func, arg = aggregate_spec
+                if arg != "*" and arg not in headers:
+                    raise ValueError(
+                        f"Unknown column: {arg}. Available columns: {headers}"
+                    )
+                required_aggregates[ref] = (func, arg)
+
         grouped_rows: list[dict[str, Any]] = []
         if group_by:
             groups: dict[tuple[Any, ...], list[dict[str, Any]]] = {}
@@ -360,33 +381,18 @@ def _execute_aggregate_select(
 
             for group_key, group_values in groups.items():
                 group_map = dict(zip(group_by, group_key))
-                output_row: dict[str, Any] = {}
-                for column in columns:
-                    if self._is_aggregate_column(column):
-                        label = self._aggregate_label(column)
-                        output_row[label] = self._compute_aggregate(
-                            str(column.get("func")),
-                            str(column.get("arg")),
-                            group_values,
-                        )
-                    else:
-                        col_name = str(column)
-                        output_row[col_name] = group_map.get(col_name)
-                if having and not self._matches_where(output_row, having):
+                context_row: dict[str, Any] = dict(group_map)
+                for label, (func, arg) in required_aggregates.items():
+                    context_row[label] = self._compute_aggregate(func, arg, group_values)
+                if having and not self._matches_where(context_row, having):
                     continue
-                grouped_rows.append(output_row)
+                grouped_rows.append(context_row)
         else:
-            output_row_single: dict[str, Any] = {}
-            for column in columns:
-                if self._is_aggregate_column(column):
-                    label = self._aggregate_label(column)
-                    output_row_single[label] = self._compute_aggregate(
-                        str(column.get("func")),
-                        str(column.get("arg")),
-                        rows,
-                    )
-            if not having or self._matches_where(output_row_single, having):
-                grouped_rows.append(output_row_single)
+            context_row_single: dict[str, Any] = {}
+            for label, (func, arg) in required_aggregates.items():
+                context_row_single[label] = self._compute_aggregate(func, arg, rows)
+            if not having or self._matches_where(context_row_single, having):
+                grouped_rows.append(context_row_single)
 
         distinct = parsed.get("distinct", False)
         if distinct:
@@ -402,9 +408,13 @@ def _execute_aggregate_select(
         order_by = parsed.get("order_by")
         if order_by:
             order_column = str(order_by["column"])
-            if order_column not in output_columns:
+            available_order_columns = set(output_columns)
+            if group_by:
+                available_order_columns.update(group_by)
+            available_order_columns.update(required_aggregates.keys())
+            if order_column not in available_order_columns:
                 raise ValueError(
-                    f"Unknown column: {order_column}. Available columns: {output_columns}"
+                    f"Unknown column: {order_column}. Available columns: {sorted(available_order_columns)}"
                 )
             reverse = order_by["direction"] == "DESC"
             grouped_rows = sorted(
@@ -445,6 +455,16 @@ def _is_aggregate_column(self, column: Any) -> bool:
     def _aggregate_label(self, aggregate: dict[str, Any]) -> str:
         return f"{str(aggregate['func']).upper()}({aggregate['arg']})"
 
+    def _aggregate_spec_from_label(self, expression: str) -> tuple[str, str] | None:
+        match = re.fullmatch(r"(?i)(COUNT|SUM|AVG|MIN|MAX)\(([^\)]+)\)", expression.strip())
+        if not match:
+            return None
+        func = match.group(1).upper()
+        arg = match.group(2).strip()
+        if not arg:
+            return None
+        return func, arg
+
     def _compute_aggregate(
         self, func: str, arg: str, rows: list[dict[str, Any]]
     ) -> Any:

src/excel_dbapi/parser.py

@@ -469,6 +469,7 @@ def _parse_select(query: str, params: Optional[tuple[Any, ...]]) -> Dict[str, An
         ]
         order_end = min(order_end_candidates) if order_end_candidates else len(remainder)
         order_part = remainder[order_start:order_end].strip()
+        order_part = _normalize_aggregate_expressions(order_part)
         order_tokens = order_part.split()
         if not order_tokens:
             raise ValueError("Invalid ORDER BY clause format")

tests/test_executor.py

@@ -48,6 +48,19 @@ def _create_empty_select_workbook(path: Path) -> None:
     workbook.save(path)
 
 
+def _create_users_workbook(path: Path) -> None:
+    workbook = Workbook()
+    sheet = workbook.active
+    assert sheet is not None
+    sheet.title = "users"
+    sheet.append(["id", "name", "age"])
+    sheet.append([1, "Alice", 30])
+    sheet.append([2, "Bob", 25])
+    sheet.append([3, "Alice", 35])
+    sheet.append([4, "Charlie", 40])
+    workbook.save(path)
+
+
 def test_executor_distinct_removes_duplicates_and_preserves_order(tmp_path: Path):
     file_path = tmp_path / "distinct_order.xlsx"
     _create_select_workbook(file_path)
@@ -222,3 +235,36 @@ def test_executor_distinct_with_group_by(tmp_path: Path):
     results = SharedExecutor(engine).execute(parsed)
 
     assert results.rows == [("A",), ("B",), ("C",)]
+
+
+def test_executor_group_by_having_count_not_in_select(tmp_path: Path):
+    file_path = tmp_path / "users_having_count_not_selected.xlsx"
+    _create_users_workbook(file_path)
+
+    engine = OpenpyxlBackend(str(file_path))
+    parsed = parse_sql("SELECT name FROM users GROUP BY name HAVING COUNT(*) > 1")
+    results = SharedExecutor(engine).execute(parsed)
+
+    assert results.rows == [("Alice",)]
+
+
+def test_executor_group_by_having_group_column_not_in_select(tmp_path: Path):
+    file_path = tmp_path / "users_having_group_column_not_selected.xlsx"
+    _create_users_workbook(file_path)
+
+    engine = OpenpyxlBackend(str(file_path))
+    parsed = parse_sql("SELECT COUNT(*) FROM users GROUP BY name HAVING name = 'Alice'")
+    results = SharedExecutor(engine).execute(parsed)
+
+    assert results.rows == [(2,)]
+
+
+def test_executor_group_by_order_by_group_key_not_in_select(tmp_path: Path):
+    file_path = tmp_path / "users_order_by_group_key_not_selected.xlsx"
+    _create_users_workbook(file_path)
+
+    engine = OpenpyxlBackend(str(file_path))
+    parsed = parse_sql("SELECT COUNT(*) FROM users GROUP BY name ORDER BY name")
+    results = SharedExecutor(engine).execute(parsed)
+
+    assert results.rows == [(2,), (1,), (1,)]

tests/test_parser.py

@@ -139,6 +139,13 @@ def test_parse_select_with_having_aggregate_expression():
     }
 
 
+def test_parse_select_order_by_aggregate_expression():
+    parsed = parse_sql(
+        "SELECT name, COUNT(*) FROM Sheet1 GROUP BY name ORDER BY COUNT(*) DESC"
+    )
+    assert parsed["order_by"] == {"column": "COUNT(*)", "direction": "DESC"}
+
+
 def test_parse_select_all_aggregate_functions():
     parsed = parse_sql(
         "SELECT COUNT(*), COUNT(score), SUM(score), AVG(score), MIN(score), MAX(score) FROM Sheet1"