Merge pull request #432 from opencybersecurityalliance/limit

kestrel2 Limit instruction

packages-nextgen/kestrel_core/src/kestrel/frontend/compile.py

@@ -26,6 +26,7 @@
 from kestrel.ir.instructions import (
     Filter,
     Source,
+    Limit,
     ProjectEntity,
     Variable,
 )
@@ -90,9 +91,14 @@ def get(self, args):
         source_node = graph.add_node(args[1])
         filter_node = graph.add_node(args[2], source_node)
         projection_node = graph.add_node(ProjectEntity(args[0].value), filter_node)
+        root = projection_node
         if len(args) > 3:
-            filter_node.timerange = args[3]
-        return graph, projection_node
+            for arg in args[3:]:
+                if isinstance(arg, TimeRange):
+                    filter_node.timerange = args[3]
+                elif isinstance(arg, Limit):
+                    root = graph.add_node(arg, projection_node)
+        return graph, root
 
     def where_clause(self, args):
         exp = args[0]
@@ -180,3 +186,8 @@ def second(self, _args):
 
     def timestamp(self, args):
         return args[0]
+
+    # Limit
+    def limit_clause(self, args):
+        n = int(args[0])
+        return Limit(n)

packages-nextgen/kestrel_core/src/kestrel/ir/instructions.py

@@ -165,6 +165,11 @@ class Reference(SourceInstruction):
     name: str
 
 
+@dataclass(eq=False)
+class Limit(TransformingInstruction):
+    num: int
+
+
 @typechecked
 def get_instruction_class(name: str) -> Type[Instruction]:
     classes = inspect.getmembers(sys.modules[__name__], inspect.isclass)

packages-nextgen/kestrel_core/tests/test_parser.py

@@ -2,7 +2,7 @@
 
 from kestrel.frontend.parser import parse_kestrel
 from kestrel.ir.graph import IRGraph
-from kestrel.ir.instructions import Filter, ProjectEntity, Source, Variable
+from kestrel.ir.instructions import Filter, ProjectEntity, Source, Variable, Limit
 
 import pytest
 
@@ -57,3 +57,19 @@ def test_parser_get_timespan_absolute():
     assert delta == timedelta(hours=5)
     assert filt.timerange.start == datetime(2023, 11, 29, 0, 0, tzinfo=timezone.utc)
     assert filt.timerange.stop == datetime(2023, 11, 29, 5, 0, tzinfo=timezone.utc)
+
+
+@pytest.mark.parametrize(
+    "stmt, expected", [
+        ("x = GET url FROM if://ds WHERE url = 'http://example.com/' LIMIT 1", 1),
+        ("x = GET url FROM if://ds WHERE url = 'http://example.com/' LAST 3d LIMIT 2", 2),
+        (("x = GET url FROM if://ds WHERE url = 'http://example.com/'"
+          " START '2023-11-29T00:00:00Z' STOP '2023-11-29T05:00:00Z' LIMIT 3"), 3),
+    ]
+)
+def test_parser_get_with_limit(stmt, expected):
+    graph = parse_kestrel(stmt)
+    limits = graph.get_nodes_by_type(Limit)
+    assert len(limits) == 1
+    limit = limits[0]
+    assert limit.num == expected