estuary-cdk: Support multi-subtask capture with independent states

Enhance the capture common module to support running multiple subtasks for incremental and backfill captures with independent states. This allows more flexible capture configurations where different subtasks can track their own progress and state within a single binding.

Author: JustinASmith

estuary-cdk/estuary_cdk/capture/common.py

@@ -1,5 +1,6 @@
 import abc
 import asyncio
+import functools
 from enum import Enum
 from dataclasses import dataclass
 from datetime import UTC, datetime, timedelta
@@ -54,7 +55,7 @@
 
 
 class Triggers(Enum):
-    BACKFILL = 'BACKFILL'
+    BACKFILL = "BACKFILL"
 
 
 class BaseDocument(BaseModel):
@@ -114,13 +115,14 @@ def path(self) -> list[str]:
 _ResourceConfig = TypeVar("_ResourceConfig", bound=ResourceConfig)
 
 
-CRON_REGEX = (r"^"
+CRON_REGEX = (
+    r"^"
     r"((?:[0-5]?\d(?:-[0-5]?\d)?|\*(?:/[0-5]?\d)?)(?:,(?:[0-5]?\d(?:-[0-5]?\d)?|\*(?:/[0-5]?\d)?))*)\s+"  # minute
     r"((?:[01]?\d|2[0-3]|(?:[01]?\d|2[0-3])-(?:[01]?\d|2[0-3])|\*(?:/[01]?\d|/2[0-3])?)(?:,(?:[01]?\d|2[0-3]|(?:[01]?\d|2[0-3])-(?:[01]?\d|2[0-3])|\*(?:/[01]?\d|/2[0-3])?))*)\s+"  # hour
     r"((?:0?[1-9]|[12]\d|3[01]|(?:0?[1-9]|[12]\d|3[01])-(?:0?[1-9]|[12]\d|3[01])|\*(?:/[0-9]|/1[0-9]|/2[0-9]|/3[01])?)(?:,(?:0?[1-9]|[12]\d|3[01]|(?:0?[1-9]|[12]\d|3[01])-(?:0?[1-9]|[12]\d|3[01])|\*(?:/[0-9]|/1[0-9]|/2[0-9]|/3[01])?))*)\s+"  # day of month
     r"((?:[1-9]|1[0-2]|(?:[1-9]|1[0-2])-(?:[1-9]|1[0-2])|\*(?:/[1-9]|/1[0-2])?)(?:,(?:[1-9]|1[0-2]|(?:[1-9]|1[0-2])-(?:[1-9]|1[0-2])|\*(?:/[1-9]|/1[0-2])?))*)\s+"  # month
     r"((?:[0-6]|(?:[0-6])-(?:[0-6])|\*(?:/[0-6])?)(?:,(?:[0-6]|(?:[0-6])-(?:[0-6])|\*(?:/[0-6])?))*)"  # day of week
-    r"$|^$" # Empty string to signify no schedule
+    r"$|^$"  # Empty string to signify no schedule
 )
 
 
@@ -129,7 +131,7 @@ class ResourceConfigWithSchedule(ResourceConfig):
         default="",
         title="Schedule",
         description="Schedule to automatically rebackfill this binding. Accepts a cron expression.",
-        pattern=CRON_REGEX
+        pattern=CRON_REGEX,
     )
 
 
@@ -164,8 +166,7 @@ class Backfill(BaseModel, extra="forbid"):
             description="LogCursor at which incremental replication began"
         )
         next_page: PageCursor = Field(
-            description="PageCursor of the next page to fetch",
-            default=None
+            description="PageCursor of the next page to fetch", default=None
         )
 
     class Snapshot(BaseModel, extra="forbid"):
@@ -179,18 +180,20 @@ class Snapshot(BaseModel, extra="forbid"):
             description="The xxh3_128 hex digest of documents of this resource in the last snapshot"
         )
 
-    inc: Incremental | None = Field(
+    inc: Incremental | dict[str, Incremental | None] | None = Field(
         default=None, description="Incremental capture progress"
     )
 
-    backfill: Backfill | None = Field(
+    backfill: Backfill | dict[str, Backfill | None] | None = Field(
         default=None,
         description="Backfill progress, or None if no backfill is occurring",
     )
 
     snapshot: Snapshot | None = Field(default=None, description="Snapshot progress")
 
-    last_initialized: datetime | None = Field(default=None, description="The last time this state was initialized.")
+    last_initialized: datetime | None = Field(
+        default=None, description="The last time this state was initialized."
+    )
 
 
 _ResourceState = TypeVar("_ResourceState", bound=ResourceState)
@@ -213,6 +216,7 @@ class AssociatedDocument(Generic[_BaseDocument]):
     You might use this if your data model requires you to load "child" documents when capturing a "parent" document,
     instead of independently loading the child data stream.
     """
+
     doc: _BaseDocument
     binding: int
 
@@ -317,7 +321,7 @@ class FixedSchema:
                     CaptureBinding[_ResourceConfig],
                     "Resource[_BaseDocument, _ResourceConfig, _ResourceState]",
                 ]
-            ]
+            ],
         ],
         None,
     ]
@@ -363,7 +367,6 @@ def resolve_bindings(
     resources: list[Resource[Any, _BaseResourceConfig, Any]],
     resource_term="Resource",
 ) -> list[tuple[_ResolvableBinding, Resource[Any, _BaseResourceConfig, Any]]]:
-
     resolved: list[
         tuple[_ResolvableBinding, Resource[Any, _BaseResourceConfig, Any]]
     ] = []
@@ -397,7 +400,6 @@ def validated(
         ]
     ],
 ) -> response.Validated:
-
     return response.Validated(
         bindings=[
             response.ValidatedBinding(resourcePath=b[0].resourceConfig.path())
@@ -415,7 +417,6 @@ def open(
         ]
     ],
 ) -> tuple[response.Opened, Callable[[Task], Awaitable[None]]]:
-
     async def _run(task: Task):
         backfill_requests = []
         if open.state.backfillRequests is not None:
@@ -445,17 +446,20 @@ async def _run(task: Task):
                 if state.last_initialized is None:
                     state.last_initialized = datetime.now(tz=UTC)
                     task.checkpoint(
-                        ConnectorState(
-                            bindingStateV1={binding.stateKey: state}
-                        )
+                        ConnectorState(bindingStateV1={binding.stateKey: state})
                     )
 
                 if isinstance(binding.resourceConfig, ResourceConfigWithSchedule):
                     cron_schedule = binding.resourceConfig.schedule
-                    next_scheduled_initialization = next_fire(cron_schedule, state.last_initialized)
+                    next_scheduled_initialization = next_fire(
+                        cron_schedule, state.last_initialized
+                    )
 
-                    if next_scheduled_initialization and next_scheduled_initialization < datetime.now(tz=UTC):
-                    # Re-initialize the binding if we missed a scheduled re-initialization.
+                    if (
+                        next_scheduled_initialization
+                        and next_scheduled_initialization < datetime.now(tz=UTC)
+                    ):
+                        # Re-initialize the binding if we missed a scheduled re-initialization.
                         should_initialize = True
                         if state.backfill:
                             task.log.warning(
@@ -464,12 +468,22 @@ async def _run(task: Task):
                                 " complete before the next scheduled backfill starts."
                             )
 
-                        next_scheduled_initialization = next_fire(cron_schedule, datetime.now(tz=UTC))
+                        next_scheduled_initialization = next_fire(
+                            cron_schedule, datetime.now(tz=UTC)
+                        )
 
-                    if next_scheduled_initialization and soonest_future_scheduled_initialization:
-                        soonest_future_scheduled_initialization = min(soonest_future_scheduled_initialization, next_scheduled_initialization)
+                    if (
+                        next_scheduled_initialization
+                        and soonest_future_scheduled_initialization
+                    ):
+                        soonest_future_scheduled_initialization = min(
+                            soonest_future_scheduled_initialization,
+                            next_scheduled_initialization,
+                        )
                     elif next_scheduled_initialization:
-                        soonest_future_scheduled_initialization = next_scheduled_initialization
+                        soonest_future_scheduled_initialization = (
+                            next_scheduled_initialization
+                        )
 
             if should_initialize:
                 # Checkpoint the binding's initialized state prior to any processing.
@@ -478,7 +492,7 @@ async def _run(task: Task):
 
                 task.checkpoint(
                     ConnectorState(
-                        bindingStateV1={binding.stateKey: state}
+                        bindingStateV1={binding.stateKey: state},
                     )
                 )
 
@@ -487,7 +501,7 @@ async def _run(task: Task):
                 index,
                 state,
                 task,
-                resolved_bindings
+                resolved_bindings,
             )
 
         async def scheduled_stop(future_dt: datetime | None) -> None:
@@ -510,8 +524,12 @@ def open_binding(
     binding_index: int,
     state: _ResourceState,
     task: Task,
-    fetch_changes: FetchChangesFn[_BaseDocument] | None = None,
-    fetch_page: FetchPageFn[_BaseDocument] | None = None,
+    fetch_changes: FetchChangesFn[_BaseDocument]
+    | dict[str, FetchChangesFn[_BaseDocument]]
+    | None = None,
+    fetch_page: FetchPageFn[_BaseDocument]
+    | dict[str, FetchPageFn[_BaseDocument]]
+    | None = None,
     fetch_snapshot: FetchSnapshotFn[_BaseDocument] | None = None,
     tombstone: _BaseDocument | None = None,
 ):
@@ -520,30 +538,96 @@ def open_binding(
 
     It does 'heavy lifting' to actually capture a binding.
 
-    TODO(johnny): Separate into snapshot vs incremental tasks?
+    When fetch_changes, fetch_page, or fetch_snapshot are provided as dictionaries,
+    each function will be run as a separate subtask with its own independent state.
+    The dictionary keys are used as subtask IDs and are used to store and retrieve
+    the state for each subtask in state.inc, state.backfill, or state.snapshot.
     """
 
     prefix = ".".join(binding.resourceConfig.path())
 
     if fetch_changes:
 
-        async def closure(task: Task):
-            assert state.inc
+        async def incremental_closure(
+            task: Task,
+            fetch_changes: FetchChangesFn[_BaseDocument],
+            state: ResourceState.Incremental,
+        ):
+            assert state and not isinstance(state, dict)
             await _binding_incremental_task(
-                binding, binding_index, fetch_changes, state.inc, task,
+                binding,
+                binding_index,
+                fetch_changes,
+                state,
+                task,
             )
 
-        task.spawn_child(f"{prefix}.incremental", closure)
+        if isinstance(fetch_changes, dict):
+            assert state.inc and isinstance(state.inc, dict)
+            for subtask_id, subtask_fetch_changes in fetch_changes.items():
+                inc_state = state.inc.get(subtask_id)
+                assert inc_state
+
+                task.spawn_child(
+                    f"{prefix}.incremental.{subtask_id}",
+                    functools.partial(
+                        incremental_closure,
+                        fetch_changes=subtask_fetch_changes,
+                        state=inc_state,
+                    ),
+                )
+        else:
+            assert state.inc and not isinstance(state.inc, dict)
+            task.spawn_child(
+                f"{prefix}.incremental",
+                functools.partial(
+                    incremental_closure,
+                    fetch_changes=fetch_changes,
+                    state=state.inc,
+                ),
+            )
 
     if fetch_page and state.backfill:
 
-        async def closure(task: Task):
-            assert state.backfill
+        async def backfill_closure(
+            task: Task,
+            fetch_page: FetchPageFn[_BaseDocument],
+            state: ResourceState.Backfill,
+        ):
+            assert state and not isinstance(state, dict)
             await _binding_backfill_task(
-                binding, binding_index, fetch_page, state.backfill, task,
+                binding,
+                binding_index,
+                fetch_page,
+                state,
+                task,
             )
 
-        task.spawn_child(f"{prefix}.backfill", closure)
+        if isinstance(fetch_page, dict):
+            assert state.backfill and isinstance(state.backfill, dict)
+            for subtask_id, subtask_fetch_page in fetch_page.items():
+                backfill_state = state.backfill.get(subtask_id)
+                assert backfill_state
+
+                task.spawn_child(
+                    f"{prefix}.backfill.{subtask_id}",
+                    functools.partial(
+                        backfill_closure,
+                        fetch_page=subtask_fetch_page,
+                        state=backfill_state,
+                    ),
+                )
+
+        else:
+            assert state.backfill and not isinstance(state.backfill, dict)
+            task.spawn_child(
+                f"{prefix}.backfill",
+                functools.partial(
+                    backfill_closure,
+                    fetch_page=fetch_page,
+                    state=state.backfill,
+                ),
+            )
 
     if fetch_snapshot:
 
@@ -612,7 +696,7 @@ async def _binding_snapshot_task(
             if isinstance(doc, dict):
                 doc["meta_"] = {
                     "op": "u" if count < state.last_count else "c",
-                    "row_id": count
+                    "row_id": count,
                 }
             else:
                 doc.meta_ = BaseDocument.Meta(
@@ -719,7 +803,10 @@ async def _binding_incremental_task(
 
         if lag < binding.resourceConfig.interval:
             sleep_for = binding.resourceConfig.interval - lag
-            task.log.info("incremental task ran recently, sleeping until `interval` has fully elapsed", {"sleep_for": sleep_for, "interval": binding.resourceConfig.interval})
+            task.log.info(
+                "incremental task ran recently, sleeping until `interval` has fully elapsed",
+                {"sleep_for": sleep_for, "interval": binding.resourceConfig.interval},
+            )
 
     while True:
         try:
@@ -747,9 +834,7 @@ async def _binding_incremental_task(
                 task.log.info("incremental task triggered backfill")
                 task.stopping.event.set()
                 task.checkpoint(
-                    ConnectorState(
-                        backfillRequests={binding.stateKey: True}
-                    )
+                    ConnectorState(backfillRequests={binding.stateKey: True})
                 )
                 return
             else:
@@ -759,7 +844,12 @@ async def _binding_incremental_task(
                     is_larger = item > state.cursor
                 elif isinstance(item, datetime) and isinstance(state.cursor, datetime):
                     is_larger = item > state.cursor
-                elif isinstance(item, tuple) and isinstance(state.cursor, tuple) and isinstance(item[0], str) and isinstance(state.cursor[0], str):
+                elif (
+                    isinstance(item, tuple)
+                    and isinstance(state.cursor, tuple)
+                    and isinstance(item[0], str)
+                    and isinstance(state.cursor[0], str)
+                ):
                     is_larger = item[0] > state.cursor[0]
                 else:
                     raise RuntimeError(
@@ -786,7 +876,7 @@ async def _binding_incremental_task(
             sleep_for = binding.resourceConfig.interval
 
         elif isinstance(state.cursor, datetime):
-            lag = (datetime.now(tz=UTC) - state.cursor)
+            lag = datetime.now(tz=UTC) - state.cursor
 
             if lag > binding.resourceConfig.interval:
                 # We're not idle. Attempt to fetch the next changes.
@@ -800,4 +890,6 @@ async def _binding_incremental_task(
             sleep_for = timedelta()
             continue
 
-        task.log.debug("incremental task is idle", {"sleep_for": sleep_for, "cursor": state.cursor})
+        task.log.debug(
+            "incremental task is idle", {"sleep_for": sleep_for, "cursor": state.cursor}
+        )