streaming: Make streams cancellable (ctrl+c etc...)

Signed-off-by: Ford <[email protected]>

Author: fordN

src/amp/client.py

@@ -323,12 +323,14 @@ def query_and_load_streaming(
             any_command.Pack(command_query)
             cmd = any_command.SerializeToString()
 
+            self.logger.info('Establishing Flight SQL connection...')
             flight_descriptor = flight.FlightDescriptor.for_command(cmd)
             info = self.conn.get_flight_info(flight_descriptor)
             reader = self.conn.do_get(info.endpoints[0].ticket)
 
             # Create streaming iterator
             stream_iterator = StreamingResultIterator(reader)
+            self.logger.info('Stream connection established, waiting for data...')
 
             # Optionally wrap with reorg detection
             if with_reorg_detection:
@@ -339,6 +341,7 @@ def query_and_load_streaming(
             loader_instance = create_loader(loader_type, loader_config)
 
             with loader_instance:
+                self.logger.info(f'Starting continuous load to {destination}. Press Ctrl+C to stop.')
                 yield from loader_instance.load_stream_continuous(stream_iterator, destination, **load_config.__dict__)
 
         except Exception as e:

src/amp/loaders/base.py

@@ -364,18 +364,7 @@ def load_stream_continuous(
 
                     except Exception as e:
                         self.logger.error(f'Failed to handle reorg: {str(e)}')
-                        yield LoadResult(
-                            rows_loaded=0,
-                            duration=duration,
-                            ops_per_second=0,
-                            table_name=table_name,
-                            loader_type=self.__class__.__name__.replace('Loader', '').lower(),
-                            success=False,
-                            error=str(e),
-                            is_reorg=True,
-                            invalidation_ranges=response.invalidation_ranges,
-                        )
-
+                        raise
                 else:
                     # Normal data batch
                     batch_count += 1
@@ -401,6 +390,9 @@ def load_stream_continuous(
 
                     yield result
 
+        except KeyboardInterrupt:
+            self.logger.info(f'Streaming cancelled by user after {batch_count} batches, {rows_loaded} rows loaded')
+            raise
         except Exception as e:
             self.logger.error(f'Streaming failed after {batch_count} batches: {str(e)}')
             duration = time.time() - start_time

src/amp/streaming/iterator.py

@@ -3,6 +3,7 @@
 """
 
 import logging
+import signal
 from typing import Iterator, Optional, Tuple
 
 import pyarrow as pa
@@ -30,10 +31,27 @@ def __init__(self, flight_reader: flight.FlightStreamReader):
         self.logger = logging.getLogger(__name__)
         self._closed = False
 
+        signal.signal(signal.SIGINT, self._handle_interrupt)
+
     def __iter__(self) -> Iterator[ResponseBatch]:
         """Return iterator instance"""
         return self
 
+    def close(self):
+        """Close the stream"""
+        if not self._closed:
+            self.logger.info('Closing stream')
+            self._closed = True
+        try:
+            self.flight_reader.cancel()
+        except Exception as e:
+            self.logger.warning(f'Error cancelling flight reader: {e}')
+
+    def _handle_interrupt(self, signum, frame):
+        """Handle SIGINT (Ctrl+C) signal"""
+        self.logger.info('Interrupt signal received (%s), cancelling stream...', signum)
+        self.close()
+
     def __next__(self) -> ResponseBatch:
         """
         Get the next batch from the stream.
@@ -43,6 +61,7 @@ def __next__(self) -> ResponseBatch:
 
         Raises:
             StopIteration: When stream is exhausted
+            KeyboardInterrupt: When user cancels the stream
         """
         if self._closed:
             raise StopIteration('Stream has been closed')
@@ -57,12 +76,16 @@ def __next__(self) -> ResponseBatch:
 
             return ResponseBatch(data=batch, metadata=metadata)
 
+        except KeyboardInterrupt:
+            self.logger.info('Stream cancelled by user')
+            self.close()
+            raise
         except StopIteration:
-            self._closed = True
+            self.close()
             raise
         except Exception as e:
             self.logger.error(f'Error reading from stream: {e}')
-            self._closed = True
+            self.close()
             raise
 
     def _read_next_batch(self) -> Tuple[Optional[pa.RecordBatch], Optional[BatchMetadata]]:
@@ -89,13 +112,6 @@ def _read_next_batch(self) -> Tuple[Optional[pa.RecordBatch], Optional[BatchMeta
         except StopIteration:
             return None, None
 
-    def close(self) -> None:
-        """Close the underlying stream"""
-        if not self._closed:
-            self._closed = True
-            # PyArrow's FlightStreamReader doesn't have an explicit close method
-            # but we mark it as closed to prevent further reads
-
     def __enter__(self) -> 'StreamingResultIterator':
         """Context manager entry"""
         return self

src/amp/streaming/reorg.py

@@ -45,39 +45,48 @@ def __next__(self) -> ResponseBatchWithReorg:
 
         Raises:
             StopIteration: When stream is exhausted
+            KeyboardInterrupt: When user cancels the stream
         """
-        # Get next batch from underlying stream
-        batch = next(self.stream_iterator)
-
-        # Check if this batch contains only duplicate ranges
-        if self._is_duplicate_batch(batch.metadata.ranges):
-            self.logger.debug(f'Skipping duplicate batch with ranges: {batch.metadata.ranges}')
-            # Recursively call to get the next non-duplicate batch
-            return self.__next__()
-
-        # Detect reorgs by comparing with previous ranges
-        invalidation_ranges = self._detect_reorg(batch.metadata.ranges)
-
-        # Update previous ranges for each network
-        for range in batch.metadata.ranges:
-            self.prev_ranges_by_network[range.network] = range
-
-        # If we detected a reorg, yield the reorg notification first
-        if invalidation_ranges:
-            self.logger.info(f'Reorg detected with {len(invalidation_ranges)} invalidation ranges')
-            # We need to yield the reorg and then the batch
-            # Store the batch to yield after the reorg
-            self._pending_batch = batch
-            return ResponseBatchWithReorg.reorg_batch(invalidation_ranges)
-
-        # Check if we have a pending batch from a previous reorg detection
-        if hasattr(self, '_pending_batch'):
-            pending = self._pending_batch
-            delattr(self, '_pending_batch')
-            return ResponseBatchWithReorg.data_batch(pending)
-
-        # Normal case - just return the data batch
-        return ResponseBatchWithReorg.data_batch(batch)
+        try:
+            # Get next batch from underlying stream
+            batch = next(self.stream_iterator)
+
+            # TODO: look for metadata.ranges_complete to see if it's a batch end. mostly for resuming streams
+            # also document the metadata.  numbers, network, hash, prev_hash (could be null)
+            # Check if this batch contains only duplicate ranges
+            if self._is_duplicate_batch(batch.metadata.ranges):
+                self.logger.debug(f'Skipping duplicate batch with ranges: {batch.metadata.ranges}')
+                # Recursively call to get the next non-duplicate batch
+                return self.__next__()
+
+            # Detect reorgs by comparing with previous ranges
+            invalidation_ranges = self._detect_reorg(batch.metadata.ranges)
+
+            # Update previous ranges for each network
+            for range in batch.metadata.ranges:
+                self.prev_ranges_by_network[range.network] = range
+
+            # If we detected a reorg, yield the reorg notification first
+            if invalidation_ranges:
+                self.logger.info(f'Reorg detected with {len(invalidation_ranges)} invalidation ranges')
+                # We need to yield the reorg and then the batch
+                # Store the batch to yield after the reorg
+                self._pending_batch = batch
+                return ResponseBatchWithReorg.reorg_batch(invalidation_ranges)
+
+            # Check if we have a pending batch from a previous reorg detection
+            if hasattr(self, '_pending_batch'):
+                pending = self._pending_batch
+                delattr(self, '_pending_batch')
+                return ResponseBatchWithReorg.data_batch(pending)
+
+            # Normal case - just return the data batch
+            return ResponseBatchWithReorg.data_batch(batch)
+
+        except KeyboardInterrupt:
+            self.logger.info('Reorg-aware stream cancelled by user')
+            self.stream_iterator.close()
+            raise
 
     def _detect_reorg(self, current_ranges: List[BlockRange]) -> List[BlockRange]:
         """