diff --git a/flow/record/adapter/jsonfile.py b/flow/record/adapter/jsonfile.py index b32fb2a0..bfbe51fa 100644 --- a/flow/record/adapter/jsonfile.py +++ b/flow/record/adapter/jsonfile.py @@ -1,7 +1,7 @@ from __future__ import annotations -import json -from typing import TYPE_CHECKING, BinaryIO +import logging +from typing import TYPE_CHECKING, Any, BinaryIO from flow import record from flow.record import JsonRecordPacker @@ -28,6 +28,9 @@ """ +log = logging.getLogger(__name__) + + class JsonfileWriter(AbstractWriter): fp = None @@ -65,33 +68,56 @@ def close(self) -> None: class JsonfileReader(AbstractReader): fp = None - def __init__(self, path: str | Path | BinaryIO, selector: str | None = None, **kwargs): + def __init__( + self, path: str | Path | BinaryIO, selector: str | None = None, record_name: str = "json/record", **kwargs + ): self.selector = make_selector(selector) self.fp = record.open_path_or_stream(path, "r") self.packer = JsonRecordPacker() + self.record_name = record_name def close(self) -> None: if self.fp: self.fp.close() self.fp = None + def obj_hook(self, obj: record.Record | dict) -> Any: + return obj + def __iter__(self) -> Iterator[Record]: ctx = get_app_context() selector = self.selector + + if not self.fp: + return + for line in self.fp: - obj = self.packer.unpack(line) - if isinstance(obj, record.Record): - if match_record_with_context(obj, selector, ctx): - yield obj - elif isinstance(obj, record.RecordDescriptor): - pass - else: - # fallback for plain jsonlines (non flow.record format) - jd = json.loads(line) + if not line or line == "\n": + continue + + try: + obj = self.packer.unpack(line) + except Exception as e: + log.warning("Failed unpacking line '%s'", line) + log.debug("", exc_info=e) + continue + + if isinstance(obj, record.RecordDescriptor): + continue + + elif isinstance(obj, record.Record) and match_record_with_context(obj, selector, ctx): + yield self.obj_hook(obj) + + elif isinstance(obj, dict): + obj = self.obj_hook(obj) + fields = [ - (fieldtype_for_value(val, "string"), key) for key, val in jd.items() if not key.startswith("_") + (fieldtype_for_value(val, "string"), key) for key, val in obj.items() if not key.startswith("_") ] - desc = record.RecordDescriptor("json/record", fields) - obj = desc(**jd) + desc = record.RecordDescriptor(self.record_name, fields) + obj = desc(**obj) if match_record_with_context(obj, selector, ctx): yield obj + + else: + log.warning("Failed handling unpacked line '%s'", line) diff --git a/tests/adapter/test_json.py b/tests/adapter/test_jsonfile.py similarity index 78% rename from tests/adapter/test_json.py rename to tests/adapter/test_jsonfile.py index 074cfd51..2a8a937c 100644 --- a/tests/adapter/test_json.py +++ b/tests/adapter/test_jsonfile.py @@ -1,11 +1,12 @@ from __future__ import annotations import json -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import pytest from flow.record import RecordReader, RecordWriter +from flow.record.adapter.jsonfile import JsonfileReader from flow.record.base import RecordDescriptor from tests._utils import generate_records @@ -13,7 +14,7 @@ from pathlib import Path -def test_json_adapter(tmp_path: Path) -> None: +def test_jsonfile_adapter(tmp_path: Path) -> None: json_file = tmp_path.joinpath("records.json") record_adapter_path = f"jsonfile://{json_file}" writer = RecordWriter(record_adapter_path) @@ -31,7 +32,7 @@ def test_json_adapter(tmp_path: Path) -> None: assert nr_records == nr_received_records -def test_json_adapter_contextmanager(tmp_path: Path) -> None: +def test_jsonfile_adapter_contextmanager(tmp_path: Path) -> None: json_file = tmp_path.joinpath("records.json") record_adapter_path = f"jsonfile://{json_file}" with RecordWriter(record_adapter_path) as writer: @@ -47,7 +48,7 @@ def test_json_adapter_contextmanager(tmp_path: Path) -> None: assert nr_records == nr_received_records -def test_json_adapter_jsonlines(tmp_path: Path) -> None: +def test_jsonfile_adapter_jsonlines(tmp_path: Path) -> None: json_file = tmp_path.joinpath("data.jsonl") items = [ @@ -75,7 +76,7 @@ def test_json_adapter_jsonlines(tmp_path: Path) -> None: "jsonfile://{json_file}?descriptors=0", ], ) -def test_json_adapter_no_record_descriptors(tmp_path: Path, record_adapter_path: str) -> None: +def test_jsonfile_adapter_no_record_descriptors(tmp_path: Path, record_adapter_path: str) -> None: json_file = tmp_path.joinpath("records.jsonl") record_adapter_path = record_adapter_path.format(json_file=json_file) @@ -99,7 +100,7 @@ def test_json_adapter_no_record_descriptors(tmp_path: Path, record_adapter_path: "jsonfile://{json_file}?descriptors=1", ], ) -def test_json_adapter_with_record_descriptors(tmp_path: Path, record_adapter_path: str) -> None: +def test_jsonfile_adapter_with_record_descriptors(tmp_path: Path, record_adapter_path: str) -> None: json_file = tmp_path.joinpath("records.jsonl") record_adapter_path = record_adapter_path.format(json_file=json_file) @@ -120,7 +121,7 @@ def test_json_adapter_with_record_descriptors(tmp_path: Path, record_adapter_pat assert descriptor_seen == 2 -def test_json_command_fieldtype(tmp_path: Path) -> None: +def test_jsonfile_command_fieldtype(tmp_path: Path) -> None: json_file = tmp_path.joinpath("records.json") record_adapter_path = f"jsonfile://{json_file}" writer = RecordWriter(record_adapter_path) @@ -155,3 +156,23 @@ def test_json_command_fieldtype(tmp_path: Path) -> None: assert records[1].commando.args == ("bash",) assert len(records) == 3 + + +def test_jsonfile_object_hook(tmp_path: Path) -> None: + """Test if the :class:`JsonFileReader` object hook method works as intended.""" + json_file = tmp_path.joinpath("records.json") + json_file.write_text('{"foo": "bar"}\n{"foo": "baz"}') + + class CustomJsonfileReader(JsonfileReader): + def obj_hook(self, obj: dict) -> Any: + obj["hello"] = "world" + return obj + + with CustomJsonfileReader(json_file) as reader: + record = next(iter(reader)) + assert record.foo == "bar" + assert record.hello == "world" + + record = next(iter(reader)) + assert record.foo == "baz" + assert record.hello == "world"