Improve reorg handling with conservative invalidation and auto-restart

- Use prev_range.start for invalidation range to ensure all potentially
  affected data is invalidated (addresses Ford review comment)
- Remove _pending_batch logic - caller now handles stream restart
- Clear prev_ranges_by_network for affected networks on reorg detection
- Add auto-restart loop in query_and_load_streaming to transparently
  handle reorgs without requiring user intervention
- Close old stream before restarting to prevent resource leaks
- Update tests to reflect new invalidation behavior

Author: incrypto32

src/amp/client.py

@@ -791,42 +791,62 @@ def query_and_load_streaming(
                 self.logger.warning(f'Failed to load checkpoint, starting from beginning: {e}')
 
         try:
-            # Execute streaming query with Flight SQL
-            # Create a CommandStatementQuery message
-            command_query = FlightSql_pb2.CommandStatementQuery()
-            command_query.query = query
-
-            # Add resume watermark if provided
-            if resume_watermark:
-                # TODO: Add watermark to query metadata when Flight SQL supports it
-                self.logger.info(f'Resuming stream from watermark: {resume_watermark}')
-
-            # Wrap the CommandStatementQuery in an Any type
-            any_command = Any()
-            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:
-                stream_iterator = ReorgAwareStream(stream_iterator, resume_watermark=resume_watermark)
-                self.logger.info('Reorg detection enabled for streaming query')
-
-            # Start continuous loading with checkpoint support
             with loader_instance:
-                self.logger.info(f'Starting continuous load to {destination}. Press Ctrl+C to stop.')
-                # Pass connection_name for checkpoint saving
-                yield from loader_instance.load_stream_continuous(
-                    stream_iterator, destination, connection_name=connection_name, **load_config.__dict__
-                )
+                while True:
+                    # Execute streaming query with Flight SQL
+                    # Create a CommandStatementQuery message
+                    command_query = FlightSql_pb2.CommandStatementQuery()
+                    command_query.query = query
+
+                    # Add resume watermark if provided
+                    if resume_watermark:
+                        # TODO: Add watermark to query metadata when Flight SQL supports it
+                        self.logger.info(f'Resuming stream from watermark: {resume_watermark}')
+
+                    # Wrap the CommandStatementQuery in an Any type
+                    any_command = Any()
+                    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:
+                        stream_iterator = ReorgAwareStream(stream_iterator, resume_watermark=resume_watermark)
+                        self.logger.info('Reorg detection enabled for streaming query')
+
+                    # Start continuous loading with checkpoint support
+                    self.logger.info(f'Starting continuous load to {destination}. Press Ctrl+C to stop.')
+
+                    reorg_result = None
+                    # Pass connection_name for checkpoint saving
+                    for result in loader_instance.load_stream_continuous(
+                        stream_iterator, destination, connection_name=connection_name, **load_config.__dict__
+                    ):
+                        yield result
+                        # Break on reorg to restart stream
+                        if result.is_reorg:
+                            reorg_result = result
+                            break
+
+                    # Check if we need to restart due to reorg
+                    if reorg_result:
+                        # Close the old stream before restarting
+                        if hasattr(stream_iterator, 'close'):
+                            stream_iterator.close()
+                        self.logger.info('Reorg detected, restarting stream with new resume position...')
+                        resume_watermark = loader_instance.state_store.get_resume_position(connection_name, destination)
+                        continue
+
+                    # Normal exit - stream completed
+                    break
 
         except Exception as e:
             self.logger.error(f'Streaming query failed: {e}')

src/amp/streaming/reorg.py

@@ -75,20 +75,16 @@ def __next__(self) -> ResponseBatch:
             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 we detected a reorg, return reorg batch
+            # Caller decides whether to stop/restart or continue
             if invalidation_ranges:
                 self.logger.info(f'Reorg detected with {len(invalidation_ranges)} invalidation ranges')
-                # Store the batch to yield after the reorg
-                self._pending_batch = batch
+                # Clear memory for affected networks so restart works correctly
+                for inv_range in invalidation_ranges:
+                    if inv_range.network in self.prev_ranges_by_network:
+                        del self.prev_ranges_by_network[inv_range.network]
                 return ResponseBatch.reorg_batch(invalidation_ranges)
 
-            # Check if we have a pending batch from a previous reorg detection
-            # REVIEW: I think we should remove this
-            if hasattr(self, '_pending_batch'):
-                pending = self._pending_batch
-                delattr(self, '_pending_batch')
-                return pending
-
             # Normal case - just return the data batch
             return batch
 
@@ -144,7 +140,7 @@ def _detect_reorg(self, current_ranges: List[BlockRange]) -> List[BlockRange]:
                 if is_reorg:
                     invalidation = BlockRange(
                         network=current_range.network,
-                        start=current_range.start,
+                        start=prev_range.start,
                         end=max(current_range.end, prev_range.end),
                         hash=prev_range.hash,
                     )

tests/unit/test_streaming_types.py

@@ -474,7 +474,7 @@ class MockIterator:
 
         assert len(invalidations) == 1
         assert invalidations[0].network == 'ethereum'
-        assert invalidations[0].start == 180
+        assert invalidations[0].start == 100  # prev_range.start
         assert invalidations[0].end == 280  # max(280, 200)
 
     def test_detect_reorg_multiple_networks(self):
@@ -504,12 +504,12 @@ class MockIterator:
 
         # Check ethereum reorg
         eth_inv = next(inv for inv in invalidations if inv.network == 'ethereum')
-        assert eth_inv.start == 150
+        assert eth_inv.start == 100  # prev_range.start
         assert eth_inv.end == 250
 
         # Check polygon reorg
         poly_inv = next(inv for inv in invalidations if inv.network == 'polygon')
-        assert poly_inv.start == 140
+        assert poly_inv.start == 50  # prev_range.start
         assert poly_inv.end == 240
 
     def test_detect_reorg_same_range_no_reorg(self):
@@ -546,7 +546,7 @@ class MockIterator:
         invalidations = stream._detect_reorg(current_ranges)
 
         assert len(invalidations) == 1
-        assert invalidations[0].start == 250
+        assert invalidations[0].start == 100  # prev_range.start
         assert invalidations[0].end == 300  # max(280, 300)
 
     def test_is_duplicate_batch_all_same(self):