loaders: Integrate resilience, checkpointing, and idempotency

Author: fordN

src/amp/client.py

@@ -346,6 +346,29 @@ def query_and_load_streaming(
 
         self.logger.info(f'Starting streaming query to {loader_type}:{destination}')
 
+        # Create loader instance early to access checkpoint store
+        loader_instance = create_loader(loader_type, loader_config)
+
+        # Load checkpoint and create resume watermark if enabled (default: enabled)
+        if resume_watermark is None and kwargs.get('resume', True):
+            try:
+                checkpoint = loader_instance.checkpoint_store.load(connection_name, destination)
+
+                if checkpoint:
+                    resume_watermark = checkpoint.to_resume_watermark()
+                    checkpoint_type = 'reorg checkpoint' if checkpoint.is_reorg else 'checkpoint'
+                    self.logger.info(
+                        f'Resuming from {checkpoint_type}: {len(checkpoint.ranges)} ranges, '
+                        f'timestamp {checkpoint.timestamp}'
+                    )
+                    if checkpoint.is_reorg:
+                        resume_points = ', '.join(
+                            f'{r.network}:{r.start}' for r in checkpoint.ranges
+                        )
+                        self.logger.info(f'Reorg resume points: {resume_points}')
+            except Exception as e:
+                self.logger.warning(f'Failed to load checkpoint, starting from beginning: {e}')
+
         try:
             # Execute streaming query with Flight SQL
             # Create a CommandStatementQuery message
@@ -376,12 +399,13 @@ def query_and_load_streaming(
                 stream_iterator = ReorgAwareStream(stream_iterator)
                 self.logger.info('Reorg detection enabled for streaming query')
 
-            # Create loader instance and start continuous loading
-            loader_instance = create_loader(loader_type, loader_config)
-
+            # Start continuous loading with checkpoint support
             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__)
+                # Pass connection_name for checkpoint saving
+                yield from loader_instance.load_stream_continuous(
+                    stream_iterator, destination, connection_name=connection_name, **load_config.__dict__
+                )
 
         except Exception as e:
             self.logger.error(f'Streaming query failed: {e}')

src/amp/loaders/base.py

@@ -4,6 +4,7 @@
 import pyarrow as pa
 from psycopg2.pool import ThreadedConnectionPool
 
+from ...streaming.idempotency import DatabaseProcessedRangesStore, IdempotencyConfig
 from ...streaming.types import BlockRange
 from ..base import DataLoader, LoadMode
 from ._postgres_helpers import has_binary_columns, prepare_csv_data, prepare_insert_data
@@ -84,6 +85,22 @@ def connect(self) -> None:
                 finally:
                     self.pool.putconn(conn)
 
+            # Replace NullStores with database-backed implementations
+            # This enables persistent checkpointing and idempotency
+            conn = self.pool.getconn()
+            try:
+                if self.checkpoint_config.enabled:
+                    from ...streaming.checkpoint import DatabaseCheckpointStore
+
+                    self.checkpoint_store = DatabaseCheckpointStore(self.checkpoint_config, conn)
+                    self.logger.info('Enabled database-backed checkpoint store')
+
+                if self.idempotency_config.enabled:
+                    self.processed_ranges_store = DatabaseProcessedRangesStore(self.idempotency_config, conn)
+                    self.logger.info('Enabled database-backed idempotency store')
+            finally:
+                self.pool.putconn(conn)
+
             self._is_connected = True
 
         except Exception as e:
@@ -109,6 +126,90 @@ def _load_batch_impl(self, batch: pa.RecordBatch, table_name: str, **kwargs) ->
         finally:
             self.pool.putconn(conn)
 
+    def load_batch_transactional(
+        self,
+        batch: pa.RecordBatch,
+        table_name: str,
+        connection_name: str,
+        ranges: List[BlockRange],
+        batch_hash: Optional[str] = None,
+    ) -> int:
+        """
+        Load a batch with transactional exactly-once semantics.
+
+        This method wraps the duplicate check, data loading, and processed marking
+        in a single PostgreSQL transaction, ensuring atomic exactly-once processing.
+
+        The transaction flow:
+        1. BEGIN TRANSACTION
+        2. Check if batch already processed (with SELECT FOR UPDATE lock)
+        3. If not processed:
+           - Load data into target table
+           - Mark ranges as processed in processed_ranges table
+        4. COMMIT (or ROLLBACK on error)
+
+        This guarantees that either both operations succeed or both fail,
+        preventing duplicate data even in case of crashes between operations.
+
+        Args:
+            batch: PyArrow RecordBatch to load
+            table_name: Target table name
+            connection_name: Connection identifier for tracking
+            ranges: Block ranges covered by this batch
+            batch_hash: Optional hash for additional validation
+
+        Returns:
+            Number of rows loaded (0 if duplicate)
+        """
+        if not self.idempotency_config.enabled:
+            raise ValueError('Transactional loading requires idempotency to be enabled')
+
+        conn = self.pool.getconn()
+        try:
+            # Create processed ranges store with this connection for transactional operations
+            store = DatabaseProcessedRangesStore(self.idempotency_config, conn)
+
+            # Disable autocommit to manage transaction manually
+            original_autocommit = conn.autocommit
+            conn.autocommit = False
+
+            try:
+                # Check if already processed (within transaction)
+                if store.is_processed(connection_name, table_name, ranges):
+                    self.logger.info(
+                        f'Batch already processed (ranges: {[f"{r.network}:{r.start}-{r.end}" for r in ranges]}), '
+                        f'skipping (transactional check)'
+                    )
+                    conn.rollback()
+                    return 0
+
+                # Load data within transaction
+                with conn.cursor() as cur:
+                    self._copy_arrow_data(cur, batch, table_name)
+
+                # Mark as processed within same transaction
+                store.mark_processed(connection_name, table_name, ranges, batch_hash)
+
+                # Commit transaction - both data load and processed marking succeed atomically
+                conn.commit()
+                self.logger.debug(
+                    f'Transactional batch load committed: {batch.num_rows} rows, '
+                    f'ranges: {[f"{r.network}:{r.start}-{r.end}" for r in ranges]}'
+                )
+                return batch.num_rows
+
+            except Exception as e:
+                # Rollback on any error - ensures no partial state
+                conn.rollback()
+                self.logger.error(f'Transactional batch load failed, rolled back: {e}')
+                raise
+            finally:
+                # Restore original autocommit setting
+                conn.autocommit = original_autocommit
+
+        finally:
+            self.pool.putconn(conn)
+
     def _clear_table(self, table_name: str) -> None:
         """Clear table for overwrite mode"""
         conn = self.pool.getconn()
@@ -208,11 +309,9 @@ def _create_table_from_schema(self, schema: pa.Schema, table_name: str) -> None:
 
                 # Build CREATE TABLE statement
                 columns = []
-                # Check if this is streaming data with metadata columns
-                has_metadata = any(field.name.startswith('_meta_') for field in schema)
 
                 for field in schema:
-                    # Skip generic metadata columns - we'll use _meta_block_range instead
+                    # Skip generic metadata columns - we'll use _meta_block_ranges instead
                     if field.name in ('_meta_range_start', '_meta_range_end'):
                         continue
                     # Special handling for JSONB metadata column
@@ -258,13 +357,14 @@ def _create_table_from_schema(self, schema: pa.Schema, table_name: str) -> None:
                     # Quote column name for safety (important for blockchain field names)
                     columns.append(f'"{field.name}" {pg_type}{nullable}')
 
-                # Add metadata columns for streaming/reorg support if this is streaming data
-                # but only if they don't already exist in the schema
-                if has_metadata:
-                    schema_field_names = [field.name for field in schema]
-                    if '_meta_block_ranges' not in schema_field_names:
-                        # Use JSONB for multi-network block ranges with GIN index support
-                        columns.append('"_meta_block_ranges" JSONB')
+                # Always add metadata column for streaming/reorg support
+                # This supports hybrid streaming (parallel catch-up → continuous streaming)
+                # where initial batches don't have metadata but later ones do
+                schema_field_names = [field.name for field in schema]
+                if '_meta_block_ranges' not in schema_field_names:
+                    # Use JSONB for multi-network block ranges with GIN index support
+                    # This column is optional and can be NULL for non-streaming loads
+                    columns.append('"_meta_block_ranges" JSONB')
 
                 # Create the table - Fixed: use proper identifier quoting
                 create_sql = f"""

src/amp/loaders/implementations/postgresql_loader.py

@@ -1,4 +1,19 @@
 # Streaming module for continuous data loading
+from .checkpoint import (
+    CheckpointConfig,
+    CheckpointState,
+    CheckpointStore,
+    DatabaseCheckpointStore,
+    NullCheckpointStore,
+)
+from .idempotency import (
+    DatabaseProcessedRangesStore,
+    IdempotencyConfig,
+    NullProcessedRangesStore,
+    ProcessedRange,
+    ProcessedRangesStore,
+    compute_batch_hash,
+)
 from .iterator import StreamingResultIterator
 from .parallel import (
     BlockRangePartitionStrategy,
@@ -27,4 +42,15 @@
     'ParallelStreamExecutor',
     'QueryPartition',
     'BlockRangePartitionStrategy',
+    'CheckpointConfig',
+    'CheckpointState',
+    'CheckpointStore',
+    'DatabaseCheckpointStore',
+    'NullCheckpointStore',
+    'IdempotencyConfig',
+    'ProcessedRange',
+    'ProcessedRangesStore',
+    'DatabaseProcessedRangesStore',
+    'NullProcessedRangesStore',
+    'compute_batch_hash',
 ]

src/amp/streaming/__init__.py

@@ -18,6 +18,7 @@
 from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional
 
 from ..loaders.types import LoadResult
+from .resilience import BackPressureConfig, RetryConfig
 
 if TYPE_CHECKING:
     from ..client import Client
@@ -53,7 +54,7 @@ def metadata(self) -> Dict[str, Any]:
 
 @dataclass
 class ParallelConfig:
-    """Configuration for parallel streaming execution"""
+    """Configuration for parallel streaming execution with resilience support"""
 
     num_workers: int
     table_name: str  # Name of the table to partition (e.g., 'blocks', 'transactions')
@@ -64,6 +65,11 @@ class ParallelConfig:
     stop_on_error: bool = False  # Stop all workers on first error
     reorg_buffer: int = 200  # Block overlap when transitioning to continuous streaming (for reorg detection)
 
+    # Resilience configuration (applied to all workers)
+    # If not specified, uses sensible defaults from resilience module
+    retry_config: Optional[RetryConfig] = None
+    back_pressure_config: Optional[BackPressureConfig] = None
+
     def __post_init__(self):
         if self.num_workers < 1:
             raise ValueError(f'num_workers must be >= 1, got {self.num_workers}')
@@ -74,6 +80,37 @@ def __post_init__(self):
         if not self.table_name:
             raise ValueError('table_name is required')
 
+    def get_resilience_config(self) -> Dict[str, Any]:
+        """
+        Get resilience configuration as a dict suitable for loader config.
+
+        Returns:
+            Dict with resilience settings, or empty dict if all None (use defaults)
+        """
+        resilience_dict = {}
+
+        if self.retry_config is not None:
+            resilience_dict['retry'] = {
+                'enabled': self.retry_config.enabled,
+                'max_retries': self.retry_config.max_retries,
+                'initial_backoff_ms': self.retry_config.initial_backoff_ms,
+                'max_backoff_ms': self.retry_config.max_backoff_ms,
+                'backoff_multiplier': self.retry_config.backoff_multiplier,
+                'jitter': self.retry_config.jitter,
+            }
+
+        if self.back_pressure_config is not None:
+            resilience_dict['back_pressure'] = {
+                'enabled': self.back_pressure_config.enabled,
+                'initial_delay_ms': self.back_pressure_config.initial_delay_ms,
+                'max_delay_ms': self.back_pressure_config.max_delay_ms,
+                'adapt_on_429': self.back_pressure_config.adapt_on_429,
+                'adapt_on_timeout': self.back_pressure_config.adapt_on_timeout,
+                'recovery_factor': self.back_pressure_config.recovery_factor,
+            }
+
+        return {'resilience': resilience_dict} if resilience_dict else {}
+
 
 class BlockRangePartitionStrategy:
     """
@@ -317,6 +354,13 @@ def execute_parallel_stream(
         """
         load_config = load_config or {}
 
+        # Merge resilience configuration into load_config
+        # This ensures all workers inherit the resilience behavior
+        resilience_config = self.config.get_resilience_config()
+        if resilience_config:
+            load_config.update(resilience_config)
+            self.logger.info('Applied resilience configuration to parallel workers')
+
         # Detect if we should continue with live streaming after parallel phase
         continue_streaming = self.config.max_block is None