Prototype of stopping after N rows loaded from taps

src/meltano/core/block/extract_load.py

@@ -4,6 +4,7 @@
 
 import asyncio
 import logging
+import os
 import typing as t
 from contextlib import asynccontextmanager, closing
 
@@ -38,6 +39,29 @@
 logger = structlog.getLogger(__name__)
 
 
+TAP_ROW_LIMITED = False
+try:
+    TAP_ROW_LIMIT = int(os.environ.get("TAP_ROW_LIMIT"))
+    TAP_ROW_LIMITED = True
+except (ValueError, TypeError):
+    TAP_ROW_LIMITED = False
+
+
+class GotEnoughRowsException(Exception):
+    pass
+
+
+class RowLimitHandler:
+    def __init__(self, row_limit):
+        self.linecount = 0
+        self.row_limit = row_limit
+
+    def writeline(self, line):
+        self.linecount += 1
+        if self.linecount > self.row_limit:
+            raise GotEnoughRowsException
+
+
 class BlockSetHasNoStateError(Exception):
     """Block has no state."""
 
@@ -223,6 +247,10 @@
             plugin_invoker=self.invoker_for(ctx),
             plugin_args=plugin_args,
         )
+
+        if TAP_ROW_LIMITED:
+            pi.add_output_handler("stdout", RowLimitHandler(TAP_ROW_LIMIT))
+
         self._blocks.append(block)
         self._env.update(ctx.env)
         return block
@@ -734,7 +762,10 @@
                     line_length_limit=self.line_length_limit,
                     stream_buffer_size=self.stream_buffer_size,
                 )
-            raise output_futures_failed.exception()  # noqa: RSE102
+            if not isinstance(ex, GotEnoughRowsException):
+                raise output_futures_failed.exception()  # noqa: RSE102
+            else:
+                logger.debug("Target Rows Reached")
         else:
             # If all the output handlers completed without raising an
             # exception, we still need to wait for all the underlying block
@@ -754,16 +785,20 @@
         else:
             logger.warning("Intermediate block in sequence failed.")
             await self._stop_all_blocks(start_idx)
-            raise RunnerError(
-                (
-                    "Unexpected completion sequence in ExtractLoadBlock set. "
-                    "Intermediate block (likely a mapper) failed."
-                ),
-                {
-                    PluginType.EXTRACTORS: 1,
-                    PluginType.LOADERS: 1,
-                },
-            )
+
+            if not isinstance(
+                output_futures_failed.exception(), GotEnoughRowsException
+            ):
+                raise RunnerError(
+                    (
+                        "Unexpected completion sequence in ExtractLoadBlock set. "
+                        "Intermediate block (likely a mapper) failed."
+                    ),
+                    {
+                        PluginType.EXTRACTORS: 1,
+                        PluginType.LOADERS: 1,
+                    },
+                )
 
     async def _handle_head_completed(
         self,