Merge pull request #82 from corva-ai/DEVC-496-be-stream-spike-contact-devcenter-team-to-update-python-sdk-to-increase-data-batch-size

corva-alex-s · web-flow · commit e510d57540e6 · 2023-12-07T16:37:57.000+02:00
New optional merge_events parameter for stream and scheduled apps
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,7 +5,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
--  Added functionality to make retries for the failed HTTP requests.
+
+### Added
+- Optional merge_events parameter to @stream and @scheduled decorators. Default is False, if True - merge all incoming events into one.
+- Added functionality to make retries for the failed HTTP requests.
 
 
 # [1.10.0] - 2023-11-08
diff --git a/docs/modules/ROOT/examples/merging/tutorial001.py b/docs/modules/ROOT/examples/merging/tutorial001.py
@@ -0,0 +1,10 @@
+from corva import Api, Cache, StreamTimeEvent, stream
+
+
+# imagine we actually have 3 incoming events with 3 records each
+@stream(merge_events=True)
+def app(event: StreamTimeEvent, api: Api, cache: Cache):
+    # since we passed merge_events=True all 3 incoming events
+    # and their records will be merged into a single event with 9 records
+    assert len(event.records) == 9  # this will not fail
+    return event
diff --git a/docs/modules/ROOT/examples/merging/tutorial002.py b/docs/modules/ROOT/examples/merging/tutorial002.py
@@ -0,0 +1,6 @@
+from corva import Api, Cache, ScheduledNaturalTimeEvent, scheduled
+
+
+@scheduled(merge_events=True)
+def app(event: ScheduledNaturalTimeEvent, api: Api, cache: Cache):
+    return event
diff --git a/docs/modules/ROOT/pages/index.adoc b/docs/modules/ROOT/pages/index.adoc
@@ -541,6 +541,34 @@ and pass corresponding logging handler
 as a keyword argument to the app decorator.
 Use code samples above as the examples.
 
+== Merging incoming events
+
+[TIP]
+====
+Only <<stream,`stream`>>
+and <<scheduled,`scheduled`>>
+apps can use this feature.
+====
+
+Sometimes Corva can send more than one event to <<scheduled,`scheduled`>> and <<stream,`stream`>> apps.
+Optionally we can ask to merge them into one event by providing `merge_events=True` parameter.
+[source,python]
+----
+include::example$merging/tutorial001.py[]
+----
+[source,python]
+----
+include::example$merging/tutorial002.py[]
+----
+Usually this is needed to save some IO operations by processing data in bigger batches.
+Use this parameter with care, in pessimistic scenario you can receive too much data, try to
+process it in "one go" and fail with timeout. In that case your app will be automatically
+restarted and you will start from the beginning and fail again.
+Without this parameter after each processed event corva-sdk will "remember" that event was processed.
+So, for example, if you will fail at event #5 and your app will be restarted - app will start processing
+from event #5(and not #1 like in case of `merge_events=True`)
+
+
 == Followable apps
 
 [TIP]
diff --git a/src/corva/handlers.py b/src/corva/handlers.py
@@ -1,5 +1,6 @@
 import contextlib
 import functools
+import itertools
 import logging
 import sys
 import warnings
@@ -28,6 +29,7 @@
 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.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 +60,7 @@ def base_handler(
     func: Callable,
     raw_event_type: Type[RawBaseEvent],
     handler: Optional[logging.Handler],
+    merge_events: bool = False,
 ) -> Callable[[Any, Any], List[Any]]:
     @functools.wraps(func)
     def wrapper(aws_event: Any, aws_context: Any) -> List[Any]:
@@ -84,6 +87,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 +125,26 @@ def stream(
     func: Optional[Callable[[StreamEventT, Api, UserRedisSdk], Any]] = None,
     *,
     handler: Optional[logging.Handler] = None,
+    merge_events: 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,18 +234,26 @@ def scheduled(
     func: Optional[Callable[[ScheduledEventT, Api, UserRedisSdk], Any]] = None,
     *,
     handler: Optional[logging.Handler] = None,
+    merge_events: bool = False,
 ) -> Callable:
     """Runs scheduled 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(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,
@@ -441,7 +462,6 @@ def wrapper(
         logging_ctx: LoggingContext,
         redis_client: redis.Redis,
     ) -> Any:
-
         logging_ctx.asset_id = event.data.asset_id
         logging_ctx.app_connection_id = event.data.app_connection_id
 
@@ -544,3 +564,51 @@ 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: Type[RawBaseEvent],
+) -> Any:
+    """
+    Merges incoming aws_events into one.
+    Merge happens differently, depending on app type.
+    Only "scheduled" and "stream" type of apps can be processed here.
+    If somehow any other type is passed - raise an exception
+    """
+    if data_transformation_type is RawScheduledEvent:
+        # scheduled event
+        if not isinstance(aws_event[0], dict):
+            aws_event = list(itertools.chain(*aws_event))
+        scheduler_type = aws_event[0]["scheduler_type"]
+        if isinstance(scheduler_type, SchedulerType):
+            scheduler_type = scheduler_type.value
+        is_depth = scheduler_type == SchedulerType.data_depth_milestone.value
+        event_start, event_end = (
+            ("top_depth", "bottom_depth")
+            if is_depth
+            else ("schedule_start", "schedule_end")
+        )
+        min_event_start = min(e[event_start] for e in aws_event)
+        max_event_end = max(
+            (e[event_end] for e in aws_event if e.get(event_end) is not None),
+            default=None,
+        )
+        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]
+
+    elif data_transformation_type is RawStreamEvent:
+        # stream event
+        for event in aws_event[1:]:
+            aws_event[0]["records"].extend(event["records"])
+        aws_event = [aws_event[0]]
+
+    else:
+        CORVA_LOGGER.warning(
+            f"{data_transformation_type.__name__} does not support `merge event` "
+            "parameter."
+        )
+
+    return aws_event
diff --git a/tests/unit/test_docs/test_merging.py b/tests/unit/test_docs/test_merging.py
@@ -0,0 +1,102 @@
+from typing import Dict, List, Union
+
+import pytest
+
+from corva import ScheduledDataTimeEvent, StreamTimeEvent
+from corva.configuration import SETTINGS
+from corva.models.scheduled.raw import RawScheduledDataTimeEvent
+from corva.models.scheduled.scheduler_type import SchedulerType
+from corva.models.stream.log_type import LogType
+from corva.models.stream.raw import (
+    RawAppMetadata,
+    RawMetadata,
+    RawStreamTimeEvent,
+    RawTimeRecord,
+)
+from docs.modules.ROOT.examples.merging import tutorial001, tutorial002
+
+
+def test_tutorial001(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
+    """
+
+    event = []
+    timestamp = 1
+    # generate 3 events with 2 records each
+    for _ in range(3):
+        event.extend(
+            [
+                RawStreamTimeEvent(
+                    records=[
+                        RawTimeRecord(
+                            collection=str(),
+                            timestamp=timestamp,
+                            asset_id=1,
+                            company_id=1,
+                        ),
+                        RawTimeRecord(
+                            collection=str(),
+                            timestamp=timestamp + 1,
+                            asset_id=1,
+                            company_id=1,
+                        ),
+                        RawTimeRecord(
+                            collection=str(),
+                            timestamp=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()
+            ]
+        )
+        timestamp += 3
+
+    result_event: StreamTimeEvent = tutorial001.app(event, context)[0]
+    assert len(result_event.records) == 9, "records were not merged into a single event"
+
+
+@pytest.mark.parametrize(
+    "time_ranges, flat",
+    (
+        [((60, 120), (61, None), (62, 122)), True],
+        [((61, None), (60, None), (62, None)), False],
+    ),
+)
+def test_tutorial002(context, time_ranges, flat):
+
+    event: List[Union[List, Dict]] = []
+    for schedule_start, schedule_end in time_ranges:
+        event.append(
+            RawScheduledDataTimeEvent(
+                asset_id=int(),
+                interval=60,
+                schedule=int(),
+                schedule_start=schedule_start,
+                schedule_end=schedule_end,
+                app_connection=int(),
+                app_stream=int(),
+                company=int(),
+                scheduler_type=SchedulerType.data_time,
+            ).dict(by_alias=True, exclude_unset=True)
+        )
+    if not flat:
+        event = [event]
+
+    result_event: ScheduledDataTimeEvent = tutorial002.app(event, context)[0]
+
+    assert result_event.start_time == 1
+    assert result_event.end_time == 60
+    max_schedule_value = time_ranges[-1][-1]
+    assert result_event.schedule_end == max_schedule_value  # type: ignore[attr-defined]
