estuary-cdk: add support for NDJSON processing in IncrementalJsonProcessor

Author: JustinASmith

estuary-cdk/estuary_cdk/incremental_json_processor.py

@@ -15,17 +15,16 @@ class _NoopRemainder(BaseModel):
 
 class IncrementalJsonProcessor(Generic[StreamedItem, Remainder]):
     """
-    Processes a stream of JSON bytes incrementally, yielding objects of type
-    `streamed_item_cls` along the way. Once iteration is complete, the
+    Processes a stream of JSON or NDJSON bytes incrementally, yielding objects of
+    type `streamed_item_cls` along the way. Once iteration is complete, the
     "remainder" of the input that was not present under the specific prefix can
     be obtained using the `get_remainder` method.
 
     The prefix is a path within the JSON document where objects to parse reside.
     Usually it will end with the ".item" suffix, which allows iteration through
-    objects in an array. More nuanced scenarios are also possible.
-
-    Does not currently support parsing NDJSON since that can be more efficiently
-    processed via other means.
+    objects in an array. Prefixes can be used with NDJSON to extract specific
+    nested objects from each newline-delimited JSON object. More nuanced
+    scenarios are also possible.
 
     Example usage if you only want the items and don't care about the remainder:
     ```python
@@ -37,6 +36,16 @@ class IncrementalJsonProcessor(Generic[StreamedItem, Remainder]):
         do_something_with(item)
     ```
 
+    Example usage for NDJSON with no prefix:
+    ```python
+    async for item in IncrementalJsonProcessor(
+        input,
+        prefix="",  # Use "" to yield the whole object per line
+        streamed_item_cls=MyItem,
+    ):
+        do_something_with(item)
+    ```
+
     For using the remainder after iteration is complete, keep a reference to the
     processor:
     ```python
@@ -47,7 +56,7 @@ class IncrementalJsonProcessor(Generic[StreamedItem, Remainder]):
         remainder_cls=MyRemainder,
     )
 
-    async for item processor:
+    async for item in processor:
         do_something_with(item)
 
     do_something_with(processor.get_remainder())
@@ -69,7 +78,10 @@ def __init__(
         self.prefix = prefix
         self.streamed_item_cls = streamed_item_cls
         self.remainder_cls = remainder_cls
-        self.parser = ijson.parse_async(_AsyncStreamWrapper(input))
+        self.parser = ijson.parse_async(
+            _AsyncStreamWrapper(input),
+            multiple_values=True,
+        )
         self.remainder_builder = _ObjectBuilder()
         self.done = False
 
@@ -121,18 +133,31 @@ class _AsyncStreamWrapper:
     def __init__(self, gen: AsyncGenerator[bytes, None]):
         self.gen = gen
         self.buf = b""
+        self.had_data = False
+        self.eof = False
 
     async def read(self, size: int = -1) -> bytes:
         if size == -1:
-            return self.buf + b"".join([chunk async for chunk in self.gen])
-
-        while len(self.buf) < size:
-            try:
-                self.buf += await anext(self.gen)
-            except StopAsyncIteration:
-                break
+            data = self.buf + b"".join([chunk async for chunk in self.gen])
+            self.eof = True
+        else:
+            while len(self.buf) < size:
+                try:
+                    self.buf += await anext(self.gen)
+                except StopAsyncIteration:
+                    self.eof = True
+                    break
 
         data, self.buf = self.buf[:size], self.buf[size:]
+        if not self.had_data:
+            data = data.lstrip(b" \t\n\r")
+
+        if self.eof and not data and not self.had_data:
+            raise StopAsyncIteration
+
+        if not self.had_data and data:
+            self.had_data = True
+
         return data
 
 

estuary-cdk/tests/test_incremental_json_processor.py

@@ -1,7 +1,8 @@
 from typing import AsyncGenerator
 
 import pytest
-from pydantic import BaseModel, TypeAdapter
+from pydantic import BaseModel
+from ijson.common import IncompleteJSONError
 
 from estuary_cdk.incremental_json_processor import IncrementalJsonProcessor
 
@@ -238,17 +239,106 @@ class Timestamp(BaseModel):
             ],
             None,
         ),
+        # NDJSON (Newline Delimited JSON)
+        (
+            b"""{"id": 1, "value": "test1"}
+            {"id": 2, "value": "test2"}
+            {"id": 3, "value": "test3"}""",
+            "",  # Note: No prefix typically for NDJSON if you want the whole record.
+            SimpleRecord,
+            None,
+            [
+                SimpleRecord(id=1, value="test1"),
+                SimpleRecord(id=2, value="test2"),
+                SimpleRecord(id=3, value="test3"),
+            ],
+            None,
+        ),
+        (
+            b"""{"record": {"id": 1, "value": "test1"}}
+            {"record": {"id": 2, "value": "test2"}}
+            {"not_a_record": {"something": "else"}}
+            {"record": {"id": 3, "value": "test3"}}""",
+            "record",  # Note: Prefix here is the key to the object in NDJSON.
+            SimpleRecord,
+            None,
+            [
+                SimpleRecord(id=1, value="test1"),
+                SimpleRecord(id=2, value="test2"),
+                SimpleRecord(id=3, value="test3"),
+            ],
+            None,
+        ),
+        # emptry NDJSON
+        (
+            b"",
+            "",
+            SimpleRecord,
+            None,
+            [],
+            None,
+        ),
+        # whitespace only NDJSON
+        (
+            b"   \n\t\n   ",
+            "",
+            SimpleRecord,
+            None,
+            [],
+            None,
+        ),
+        # NDJSON with trailing newline
+        (
+            b"""{"id": 1, "value": "test1"}
+            {"id": 2, "value": "test2"}
+            
+            """,
+            "",
+            SimpleRecord,
+            None,
+            [
+                SimpleRecord(id=1, value="test1"),
+                SimpleRecord(id=2, value="test2"),
+            ],
+            None,
+        ),
+        # NDJSON with multiple trailing newlines
+        (
+            b"""{"id": 1, "value": "test1"}
+            
+
+            """,
+            "",
+            SimpleRecord,
+            None,
+            [
+                SimpleRecord(id=1, value="test1"),
+            ],
+            None,
+        ),
     ],
 )
 async def test_incremental_json_processor(
-    input_data, prefix, record_cls, remainder_cls, want_records, want_remainder
+    input_data,
+    prefix,
+    record_cls,
+    remainder_cls,
+    want_records,
+    want_remainder,
 ):
     if remainder_cls:
         processor = IncrementalJsonProcessor(
-            bytes_gen(input_data), prefix, record_cls, remainder_cls
+            bytes_gen(input_data),
+            prefix,
+            record_cls,
+            remainder_cls,
         )
     else:
-        processor = IncrementalJsonProcessor(bytes_gen(input_data), prefix, record_cls)
+        processor = IncrementalJsonProcessor(
+            bytes_gen(input_data),
+            prefix,
+            record_cls,
+        )
 
     got = []
     async for record in processor:
@@ -257,3 +347,23 @@ async def test_incremental_json_processor(
     assert want_records == got
     if remainder_cls:
         assert want_remainder == processor.get_remainder()
+
+
+@pytest.mark.asyncio
+async def test_malformed_ndjson_fails():
+    # When ijson encounters malformed JSON in NDJSON mode, it fails immediately
+    # This is the expected behavior - we want to fail fast on malformed data
+    malformed_data = b"""
+    {"id": 1, "value": "test1"}
+    {"id": 2, "value": "test2"
+    {"id": 3, "value": "test3"}"""
+
+    processor = IncrementalJsonProcessor(
+        bytes_gen(malformed_data),
+        "",
+        SimpleRecord,
+    )
+
+    with pytest.raises(IncompleteJSONError):
+        async for record in processor:
+            pass  # We expect the error before any records are yielded due to ijson's buffering