Support frame clauses

These are used in window functions and would likely only need parameters
for offsets, as implemented.

Author: pgjones

src/sql_tstring/__init__.py

@@ -131,12 +131,15 @@ def _check_valid(
     *,
     case_sensitive: set[str] | None = None,
     case_insensitive: set[str] | None = None,
+    value_type: type = str,
 ) -> None:
     if case_sensitive is None:
         case_sensitive = set()
     if case_insensitive is None:
         case_insensitive = set()
-    if not isinstance(value, str) or (
+    if not isinstance(value, value_type):
+        raise ValueError(f"{value} is not valid, must be {value_type}")
+    if isinstance(value, str) and (
         value not in case_sensitive and value.lower() not in case_insensitive
     ):
         raise ValueError(
@@ -274,36 +277,40 @@ def _replace_placeholder(
                     )
                     result.append(value)
     else:
-        if clause is not None and clause.text.lower() == "order by":
-            _check_valid(
-                value,
-                case_sensitive=ctx.columns,
-                case_insensitive={"asc", "ascending", "desc", "descending"},
-            )
-            new_node = Part(text=typing.cast(str, value), parent=node.parent)
-        elif placeholder_type == PlaceholderType.COLUMN:
-            _check_valid(value, case_sensitive=ctx.columns)
-            new_node = Part(text=typing.cast(str, value), parent=node.parent)
-        elif placeholder_type == PlaceholderType.TABLE:
-            _check_valid(value, case_sensitive=ctx.tables)
-            new_node = Part(text=typing.cast(str, value), parent=node.parent)
-        elif placeholder_type == PlaceholderType.LOCK:
-            _check_valid(value, case_insensitive={"", "nowait", "skip locked"})
-            new_node = Part(text=typing.cast(str, value), parent=node.parent)
-        else:
-            if (
-                value is RewritingValue.IS_NULL or value is RewritingValue.IS_NOT_NULL
-            ) and placeholder_type == PlaceholderType.VARIABLE_CONDITION:
-                for part in node.parent.parts:
-                    if isinstance(part, Operator):
-                        if value is RewritingValue.IS_NULL:
-                            part.text = "IS"
-                        else:
-                            part.text = "IS NOT"
-                new_node = Part(text="NULL", parent=node.parent)
-            else:
-                new_node = node
-                result.append(value)
+        match placeholder_type:
+            case PlaceholderType.COLUMN:
+                _check_valid(value, case_sensitive=ctx.columns)
+                new_node = Part(text=typing.cast(str, value), parent=node.parent)
+            case PlaceholderType.FRAME:
+                _check_valid(value, value_type=int)
+                new_node = Part(text=str(value), parent=node.parent)
+            case PlaceholderType.LOCK:
+                _check_valid(value, case_insensitive={"", "nowait", "skip locked"})
+                new_node = Part(text=typing.cast(str, value), parent=node.parent)
+            case PlaceholderType.SORT:
+                _check_valid(
+                    value,
+                    case_sensitive=ctx.columns,
+                    case_insensitive={"asc", "ascending", "desc", "descending"},
+                )
+                new_node = Part(text=typing.cast(str, value), parent=node.parent)
+            case PlaceholderType.TABLE:
+                _check_valid(value, case_sensitive=ctx.tables)
+                new_node = Part(text=typing.cast(str, value), parent=node.parent)
+            case _:
+                if (
+                    value is RewritingValue.IS_NULL or value is RewritingValue.IS_NOT_NULL
+                ) and placeholder_type == PlaceholderType.VARIABLE_CONDITION:
+                    for part in node.parent.parts:
+                        if isinstance(part, Operator):
+                            if value is RewritingValue.IS_NULL:
+                                part.text = "IS"
+                            else:
+                                part.text = "IS NOT"
+                    new_node = Part(text="NULL", parent=node.parent)
+                else:
+                    new_node = node
+                    result.append(value)
 
         if isinstance(node.parent, (Expression, ExpressionGroup, Function, Group)):
             node.parent.parts[index] = new_node

src/sql_tstring/parser.py

@@ -19,7 +19,9 @@
 class PlaceholderType(Enum):
     COLUMN = auto()
     DISALLOWED = auto()
+    FRAME = auto()
     LOCK = auto()
+    SORT = auto()
     TABLE = auto()
     VARIABLE = auto()
     VARIABLE_CONDITION = auto()
@@ -117,6 +119,13 @@ class ClauseProperties:
         ),
     },
     "order": {
+        "by": {
+            "": ClauseProperties(
+                allow_empty=False, placeholder_type=PlaceholderType.SORT, separators={","}
+            )
+        },
+    },
+    "partition": {
         "by": {
             "": ClauseProperties(
                 allow_empty=False, placeholder_type=PlaceholderType.COLUMN, separators={","}
@@ -152,6 +161,11 @@ class ClauseProperties:
             allow_empty=False, placeholder_type=PlaceholderType.TABLE, separators=set()
         )
     },
+    "groups": {
+        "": ClauseProperties(
+            allow_empty=False, placeholder_type=PlaceholderType.FRAME, separators=set()
+        )
+    },
     "having": {
         "": ClauseProperties(
             allow_empty=False,
@@ -172,11 +186,21 @@ class ClauseProperties:
             allow_empty=False, placeholder_type=PlaceholderType.VARIABLE, separators=set()
         )
     },
+    "range": {
+        "": ClauseProperties(
+            allow_empty=False, placeholder_type=PlaceholderType.FRAME, separators=set()
+        )
+    },
     "returning": {
         "": ClauseProperties(
             allow_empty=False, placeholder_type=PlaceholderType.DISALLOWED, separators={","}
         )
     },
+    "rows": {
+        "": ClauseProperties(
+            allow_empty=False, placeholder_type=PlaceholderType.FRAME, separators=set()
+        )
+    },
     "select": {
         "": ClauseProperties(
             allow_empty=False, placeholder_type=PlaceholderType.COLUMN, separators={","}

tests/test_identifiers.py

@@ -31,6 +31,45 @@ def test_order_by_invalid_column() -> None:
         sql("SELECT x FROM y ORDER BY {a}, {b}", locals())
 
 
+def test_partition_by() -> None:
+    a = RewritingValue.ABSENT
+    b = "x"
+    with sql_context(columns={"x"}):
+        assert ("SELECT x OVER(PARTITION BY x) FROM y", []) == sql(
+            "SELECT x OVER(PARTITION BY {b}) FROM y", locals()
+        )
+
+
+def test_partition_by_invalid_column() -> None:
+    a = RewritingValue.ABSENT
+    b = "x"
+    with pytest.raises(ValueError):
+        sql("SELECT x OVER(PARTITION BY {b}) FROM y", locals())
+
+
+@pytest.mark.parametrize(
+    "frame_clause",
+    ["GROUPS", "RANGE", "ROWS"],
+)
+def test_frame_clause_int(frame_clause: str) -> None:
+    a = RewritingValue.ABSENT
+    b = 2
+    assert (f"SELECT x OVER(PARTITION BY x {frame_clause} 2 PRECEDING) FROM y", []) == sql(
+        f"SELECT x OVER(PARTITION BY x {frame_clause} {{b}} PRECEDING) FROM y", locals()
+    )
+
+
+@pytest.mark.parametrize(
+    "frame_clause",
+    ["GROUPS", "RANGE", "ROWS"],
+)
+def test_frame_clause_invalid(frame_clause: str) -> None:
+    a = RewritingValue.ABSENT
+    b = "INVALID"
+    with pytest.raises(ValueError):
+        sql(f"SELECT x OVER(PARTITION BY x {frame_clause} {{b}} PRECEDING) FROM y", locals())
+
+
 @pytest.mark.parametrize(
     "lock_type, expected",
     (