Add --skip flag to rdump (#76)

Adds the ability to skip a number of records when reading or writing records using rdump.

Co-authored-by: Yun Zheng Hu <hu@fox-it.com>

flow/record/tools/rdump.py

@@ -4,6 +4,7 @@
 import logging
 import sys
 from importlib import import_module
+from itertools import islice
 from pathlib import Path
 from textwrap import indent
 from urllib.parse import parse_qsl, urlencode, urlparse
@@ -95,6 +96,7 @@ def main(argv=None):
     output = parser.add_argument_group("output control")
     output.add_argument("-f", "--format", metavar="FORMAT", help="Format string")
     output.add_argument("-c", "--count", type=int, help="Exit after COUNT records")
+    output.add_argument("--skip", metavar="COUNT", type=int, default=0, help="Skip the first COUNT records")
     output.add_argument("-w", "--writer", metavar="OUTPUT", default=None, help="Write records to output")
     output.add_argument("-m", "--mode", default=None, choices=("csv", "json", "jsonlines", "line"), help="Output mode")
     output.add_argument(
@@ -201,9 +203,11 @@ def main(argv=None):
 
     selector = make_selector(args.selector, not args.no_compile)
     seen_desc = set()
+    islice_stop = (args.count + args.skip) if args.count else None
+    record_iterator = islice(record_stream(args.src, selector), args.skip, islice_stop)
     count = 0
     with RecordWriter(uri) as record_writer:
-        for count, rec in enumerate(record_stream(args.src, selector), start=1):
+        for count, rec in enumerate(record_iterator, start=1):
             if args.record_source is not None:
                 rec._source = args.record_source
             if args.record_classification is not None:
@@ -227,9 +231,6 @@ def main(argv=None):
                 else:
                     record_writer.write(rec)
 
-            if args.count and count >= args.count:
-                break
-
     if args.list:
         print("Processed {} records".format(count))
 

tests/test_rdump.py

@@ -455,3 +455,57 @@ def test_rdump_headerless_csv(tmp_path, capsysbinary):
         b"<csv/reader count='2' text='world'>",
         b"<csv/reader count='3' text='bar'>",
     ]
+
+
+@pytest.mark.parametrize(
+    ("total_records", "count", "skip", "expected_numbers"),
+    [
+        (10, None, 2, [2, 3, 4, 5, 6, 7, 8, 9]),
+        (10, 3, None, [0, 1, 2]),
+        (10, 2, 3, [3, 4]),
+        (10, None, 9, [9]),
+        (10, None, 10, []),
+    ],
+)
+def test_rdump_count_and_skip(tmp_path, capsysbinary, total_records, count, skip, expected_numbers):
+    TestRecord = RecordDescriptor(
+        "test/record",
+        [
+            ("varint", "number"),
+            ("string", "foo"),
+        ],
+    )
+
+    # Write test records to a file
+    full_set_path = tmp_path / "test_full_set.records"
+    with RecordWriter(full_set_path) as writer:
+        for i in range(total_records):
+            record = TestRecord(number=i, foo="bar" + "baz" * i)
+            writer.write(record)
+
+    rdump_parameters = []
+    if count is not None:
+        rdump_parameters.append(f"--count={count}")
+    if skip is not None:
+        rdump_parameters.append(f"--skip={skip}")
+
+    rdump.main([str(full_set_path), "--csv", "-F", "number"] + rdump_parameters)
+    captured = capsysbinary.readouterr()
+    assert captured.err == b""
+
+    # Skip csv header
+    record_lines = captured.out.splitlines()[1:]
+
+    # Convert numbers to integers and validate
+    numbers = list(map(int, record_lines))
+    assert numbers == expected_numbers
+
+    # Write records using --skip and --count to a new file
+    subset_path = tmp_path / "test_subset.records"
+    rdump.main([str(full_set_path), "-w", str(subset_path)] + rdump_parameters)
+
+    # Read records from new file and validate
+    numbers = None
+    with RecordReader(subset_path) as reader:
+        numbers = [rec.number for rec in reader]
+    assert numbers == expected_numbers