Add optional merge_events parameter to stream and scheduled decorators

Author: corva-alex-s

src/corva/handlers.py

@@ -1,5 +1,6 @@
 import contextlib
 import functools
+import itertools
 import logging
 import sys
 import warnings
@@ -27,7 +28,9 @@
 from corva.models.merge.merge import PartialRerunMergeEvent
 from corva.models.merge.raw import RawPartialRerunMergeEvent
 from corva.models.scheduled.raw import RawScheduledEvent
-from corva.models.scheduled.scheduled import ScheduledEvent, ScheduledNaturalTimeEvent
+from corva.models.scheduled.scheduled import ScheduledEvent, ScheduledNaturalTimeEvent, ScheduledDataTimeEvent, \
+    ScheduledDepthEvent
+from corva.models.scheduled.scheduler_type import SchedulerType
 from corva.models.stream.raw import RawStreamEvent
 from corva.models.stream.stream import StreamEvent
 from corva.models.task import RawTaskEvent, TaskEvent, TaskStatus
@@ -58,6 +61,7 @@ def base_handler(
     func: Callable,
     raw_event_type: Type[RawBaseEvent],
     handler: Optional[logging.Handler],
+    merge_events: Optional[bool] = False
 ) -> Callable[[Any, Any], List[Any]]:
     @functools.wraps(func)
     def wrapper(aws_event: Any, aws_context: Any) -> List[Any]:
@@ -84,6 +88,8 @@ def wrapper(aws_event: Any, aws_context: Any) -> List[Any]:
                     url=SETTINGS.CACHE_URL, decode_responses=True, max_connections=1
                 )
                 data_transformation_type = raw_custom_event_type or raw_event_type
+                if merge_events:
+                    aws_event = _merge_events(aws_event, data_transformation_type)
                 raw_events = data_transformation_type.from_raw_event(event=aws_event)
 
                 if (
@@ -120,18 +126,20 @@ def stream(
     func: Optional[Callable[[StreamEventT, Api, UserRedisSdk], Any]] = None,
     *,
     handler: Optional[logging.Handler] = None,
+    merge_events: Optional[bool] = False
 ) -> Callable:
     """Runs stream app.
 
     Arguments:
         handler: logging handler to include in Corva logger.
+        merge_events: if True - merge all incoming events into one before passing them to func
     """
 
     if func is None:
-        return functools.partial(stream, handler=handler)
+        return functools.partial(stream, handler=handler, merge_events=merge_events)
 
     @functools.wraps(func)
-    @functools.partial(base_handler, raw_event_type=RawStreamEvent, handler=handler)
+    @functools.partial(base_handler, raw_event_type=RawStreamEvent, handler=handler, merge_events=merge_events)
     def wrapper(
         event: RawStreamEvent,
         api_key: str,
@@ -221,6 +229,7 @@ def scheduled(
     func: Optional[Callable[[ScheduledEventT, Api, UserRedisSdk], Any]] = None,
     *,
     handler: Optional[logging.Handler] = None,
+    merge_events: Optional[bool] = False
 ) -> Callable:
     """Runs scheduled app.
 
@@ -229,10 +238,10 @@ def scheduled(
     """
 
     if func is None:
-        return functools.partial(scheduled, handler=handler)
+        return functools.partial(scheduled, handler=handler, merge_events=merge_events)
 
     @functools.wraps(func)
-    @functools.partial(base_handler, raw_event_type=RawScheduledEvent, handler=handler)
+    @functools.partial(base_handler, raw_event_type=RawScheduledEvent, handler=handler, merge_events=merge_events)
     def wrapper(
         event: RawScheduledEvent,
         api_key: str,
@@ -544,3 +553,33 @@ def _get_custom_event_type_by_raw_aws_event(
         if events:
             return event_type, handler
     return None, None
+
+
+def _merge_events(aws_event: Any, data_transformation_type: RawBaseEvent) -> Any:
+    """
+    Merges incoming aws_events into one.
+    Merge happens differently, depending on app type.
+    """
+    if data_transformation_type == RawScheduledEvent:
+        if not isinstance(aws_event[0], dict):
+            aws_event: List[dict] = list(itertools.chain(*aws_event))
+        is_depth = aws_event[0]["scheduler_type"] == SchedulerType.data_depth_milestone
+        event_start, event_end = ("top_depth", "bottom_depth") if is_depth else ("schedule_start", "schedule_end")
+        min_event_start, max_event_end = aws_event[0][event_start], aws_event[0].get(event_end)
+        for event in aws_event[1:]:
+            if event[event_start] < min_event_start:
+                min_event_start = event[event_start]
+            if max_event_end and event[event_end] > max_event_end:
+                max_event_end = event[event_end]
+
+        aws_event[0][event_start] = min_event_start
+        if max_event_end:
+            aws_event[0][event_end] = max_event_end
+        aws_event = aws_event[0]
+        return aws_event
+
+    # stream event
+    for event in aws_event[1:]:
+        aws_event[0]["records"].extend(event["records"])
+    aws_event = [aws_event[0]]
+    return aws_event

tests/unit/test_scheduled_app.py

@@ -532,3 +532,66 @@ def scheduled_app(event, api, cache):
 
     with pytest.raises(redis.exceptions.ConnectionError):
         scheduled_app(event, context)
+
+
+def test_merge_events_scheduled_event(context, requests_mock):
+    @scheduled(merge_events=True)
+    def scheduled_app(event, api, state):
+        return event
+
+    event = [
+        [
+            RawScheduledDataTimeEvent(
+                asset_id=int(),
+                interval=60,
+                schedule=int(),
+                schedule_start=60,
+                schedule_end=120,
+                app_connection=int(),
+                app_stream=int(),
+                company=int(),
+                scheduler_type=SchedulerType.data_time,
+            ).dict(
+                by_alias=True,
+                exclude_unset=True,
+            ),
+            RawScheduledDataTimeEvent(
+                asset_id=int(),
+                interval=60,
+                schedule=int(),
+                schedule_start=61,
+                schedule_end=121,
+                app_connection=int(),
+                app_stream=int(),
+                company=int(),
+                scheduler_type=SchedulerType.data_time,
+            ).dict(
+                by_alias=True,
+                exclude_unset=True,
+            ),
+            RawScheduledDataTimeEvent(
+                asset_id=int(),
+                interval=60,
+                schedule=int(),
+                schedule_start=62,
+                schedule_end=122,
+                app_connection=int(),
+                app_stream=int(),
+                company=int(),
+                scheduler_type=SchedulerType.data_time,
+            ).dict(
+                by_alias=True,
+                exclude_unset=True,
+            )
+        ]
+    ]
+
+    # patch post request, that sets scheduled task as completed
+    # looks for url path like /scheduler/123/completed
+    post_mock = requests_mock.post(re.compile(r'/scheduler/\d+/completed'))
+
+    result_event: ScheduledDataTimeEvent = scheduled_app(event, context)[0]
+
+    assert result_event.start_time == 1
+    assert result_event.end_time == 60
+    assert result_event.schedule_end == 122

tests/unit/test_stream_app.py

@@ -615,3 +615,78 @@ def stream_app(event, api, cache):
     result_event: StreamDepthEvent = stream_app(event, context)[0]
 
     assert result_event.log_identifier == 'log_identifier'
+
+
+def test_merge_events_stream_event(context):
+    """
+    merge_events parameter is merging records for "stream" apps
+
+    When 3 events with 2 records each are sent and @stream decorator has optional "merge_events" param set to True -
+    we're supposed to merge incoming events into one event with all records(6 in our case) combined
+    """
+    @stream(merge_events=True)
+    def stream_app(event, api, cache):
+        return event
+
+    event = [
+        RawStreamTimeEvent(
+            records=[
+                RawTimeRecord(
+                    collection=str(),
+                    timestamp=1,
+                    **{"asset_id": 1, "company_id": 1},
+                ),
+                RawTimeRecord(
+                    collection=str(),
+                    timestamp=2,
+                    **{"asset_id": 1, "company_id": 1},
+                ),
+            ],
+            metadata=RawMetadata(
+                app_stream_id=1,
+                apps={SETTINGS.APP_KEY: RawAppMetadata(app_connection_id=1)},
+                log_type=LogType.time,
+            ),
+        ).dict(),
+        RawStreamTimeEvent(
+            records=[
+                RawTimeRecord(
+                    collection=str(),
+                    timestamp=3,
+                    **{"asset_id": 1, "company_id": 1},
+                ),
+                RawTimeRecord(
+                    collection=str(),
+                    timestamp=4,
+                    **{"asset_id": 1, "company_id": 1},
+                )
+            ],
+            metadata=RawMetadata(
+                app_stream_id=1,
+                apps={SETTINGS.APP_KEY: RawAppMetadata(app_connection_id=1)},
+                log_type=LogType.time,
+            ),
+        ).dict(),
+        RawStreamTimeEvent(
+            records=[
+                RawTimeRecord(
+                    collection=str(),
+                    timestamp=5,
+                    **{"asset_id": 1, "company_id": 1},
+                ),
+                RawTimeRecord(
+                    collection=str(),
+                    timestamp=6,
+                    **{"asset_id": 1, "company_id": 1},
+                )
+            ],
+            metadata=RawMetadata(
+                app_stream_id=1,
+                apps={SETTINGS.APP_KEY: RawAppMetadata(app_connection_id=1)},
+                log_type=LogType.time,
+            ),
+        ).dict()
+    ]
+
+    result_event: StreamEvent = stream_app(event, context)[0]
+    assert len(result_event.records) == 6, "records were not merged into a single event"