snowflake loader: Add reorg aware streaming support

Author: fordN

src/amp/loaders/implementations/snowflake_loader.py

@@ -1,13 +1,14 @@
 import io
 import time
 from dataclasses import dataclass
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
 import pyarrow as pa
 import pyarrow.csv as pa_csv
 import snowflake.connector
 from snowflake.connector import DictCursor, SnowflakeConnection
 
+from ...streaming.types import BlockRange
 from ..base import DataLoader, LoadMode
 
 
@@ -390,7 +391,7 @@ def get_table_info(self, table_name: str) -> Optional[Dict[str, Any]]:
             # Get table metadata
             self.cursor.execute(
                 """
-                SELECT 
+                SELECT
                     TABLE_NAME,
                     TABLE_SCHEMA,
                     TABLE_CATALOG,
@@ -412,7 +413,7 @@ def get_table_info(self, table_name: str) -> Optional[Dict[str, Any]]:
             # Get column information
             self.cursor.execute(
                 """
-                SELECT 
+                SELECT
                     COLUMN_NAME,
                     DATA_TYPE,
                     IS_NULLABLE,
@@ -456,3 +457,83 @@ def get_table_info(self, table_name: str) -> Optional[Dict[str, Any]]:
         except Exception as e:
             self.logger.error(f"Failed to get table info for '{table_name}': {str(e)}")
             return None
+
+    def _handle_reorg(self, invalidation_ranges: List[BlockRange], table_name: str) -> None:
+        """
+        Handle blockchain reorganization by deleting affected rows from Snowflake.
+
+        Snowflake's SQL capabilities allow for efficient deletion using JSON functions
+        to parse the _meta_block_ranges column and identify affected rows.
+
+        Args:
+            invalidation_ranges: List of block ranges to invalidate (reorg points)
+            table_name: The table containing the data to invalidate
+        """
+        if not invalidation_ranges:
+            return
+
+        try:
+            # First check if the table has the metadata column
+            self.cursor.execute(
+                """
+                SELECT COUNT(*) as count
+                FROM INFORMATION_SCHEMA.COLUMNS
+                WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND COLUMN_NAME = '_META_BLOCK_RANGES'
+                """,
+                (self.config.schema, table_name.upper()),
+            )
+
+            result = self.cursor.fetchone()
+            if not result or result['COUNT'] == 0:
+                self.logger.warning(
+                    f"Table '{table_name}' doesn't have '_meta_block_ranges' column, skipping reorg handling"
+                )
+                return
+
+            # Build DELETE statement with conditions for each invalidation range
+            # Snowflake's PARSE_JSON and ARRAY_SIZE functions help work with JSON data
+            delete_conditions = []
+
+            for range_obj in invalidation_ranges:
+                network = range_obj.network
+                reorg_start = range_obj.start
+
+                # Create condition for this network's reorg
+                # Delete rows where any range in the JSON array for this network has end >= reorg_start
+                condition = f"""
+                EXISTS (
+                    SELECT 1
+                    FROM TABLE(FLATTEN(input => PARSE_JSON("_META_BLOCK_RANGES"))) f
+                    WHERE f.value:network::STRING = '{network}'
+                    AND f.value:end::NUMBER >= {reorg_start}
+                )
+                """
+                delete_conditions.append(condition)
+
+            # Combine conditions with OR
+            if delete_conditions:
+                where_clause = ' OR '.join(f'({cond})' for cond in delete_conditions)
+
+                # Execute deletion
+                delete_sql = f'DELETE FROM {table_name} WHERE {where_clause}'
+
+                self.logger.info(
+                    f'Executing blockchain reorg deletion for {len(invalidation_ranges)} networks '
+                    f"in Snowflake table '{table_name}'"
+                )
+
+                # Execute the delete and get row count
+                self.cursor.execute(delete_sql)
+                deleted_rows = self.cursor.rowcount
+
+                # Commit the transaction
+                self.connection.commit()
+
+                self.logger.info(f"Blockchain reorg deleted {deleted_rows} rows from table '{table_name}'")
+
+        except Exception as e:
+            self.logger.error(f"Failed to handle blockchain reorg for table '{table_name}': {str(e)}")
+            # Rollback on error
+            if self.connection:
+                self.connection.rollback()
+            raise

tests/integration/test_snowflake_loader.py

@@ -387,3 +387,202 @@ def test_schema_with_special_characters(self, snowflake_config, test_table_name,
             assert row['first name'] == 'Alice'
             assert abs(row['total$amount'] - 100.0) < 0.001
             assert row['2024_data'] == 'a'
+
+    def test_handle_reorg_no_metadata_column(self, snowflake_config, test_table_name, cleanup_tables):
+        """Test reorg handling when table lacks metadata column"""
+        from src.amp.streaming.types import BlockRange
+
+        cleanup_tables.append(test_table_name)
+        loader = SnowflakeLoader(snowflake_config)
+
+        with loader:
+            # Create table without metadata column
+            data = pa.table({'id': [1, 2, 3], 'block_num': [100, 150, 200], 'value': [10.0, 20.0, 30.0]})
+            loader.load_table(data, test_table_name, create_table=True)
+
+            # Call handle reorg
+            invalidation_ranges = [BlockRange(network='ethereum', start=150, end=250)]
+
+            # Should log warning and not modify data
+            loader._handle_reorg(invalidation_ranges, test_table_name)
+
+            # Verify data unchanged
+            loader.cursor.execute(f'SELECT COUNT(*) FROM {test_table_name}')
+            count = loader.cursor.fetchone()['COUNT(*)']
+            assert count == 3
+
+    def test_handle_reorg_single_network(self, snowflake_config, test_table_name, cleanup_tables):
+        """Test reorg handling for single network data"""
+        import json
+
+        from src.amp.streaming.types import BlockRange
+
+        cleanup_tables.append(test_table_name)
+        loader = SnowflakeLoader(snowflake_config)
+
+        with loader:
+            # Create table with metadata
+            block_ranges = [
+                [{'network': 'ethereum', 'start': 100, 'end': 110}],
+                [{'network': 'ethereum', 'start': 150, 'end': 160}],
+                [{'network': 'ethereum', 'start': 200, 'end': 210}],
+            ]
+
+            data = pa.table(
+                {
+                    'id': [1, 2, 3],
+                    'block_num': [105, 155, 205],
+                    '_meta_block_ranges': [json.dumps(ranges) for ranges in block_ranges],
+                }
+            )
+
+            # Load initial data
+            result = loader.load_table(data, test_table_name, create_table=True)
+            assert result.success
+            assert result.rows_loaded == 3
+
+            # Verify all data exists
+            loader.cursor.execute(f'SELECT COUNT(*) FROM {test_table_name}')
+            count = loader.cursor.fetchone()['COUNT(*)']
+            assert count == 3
+
+            # Reorg from block 155 - should delete rows 2 and 3
+            invalidation_ranges = [BlockRange(network='ethereum', start=155, end=300)]
+            loader._handle_reorg(invalidation_ranges, test_table_name)
+
+            # Verify only first row remains
+            loader.cursor.execute(f'SELECT COUNT(*) FROM {test_table_name}')
+            count = loader.cursor.fetchone()['COUNT(*)']
+            assert count == 1
+
+            loader.cursor.execute(f'SELECT id FROM {test_table_name}')
+            remaining_id = loader.cursor.fetchone()['ID']
+            assert remaining_id == 1
+
+    def test_handle_reorg_multi_network(self, snowflake_config, test_table_name, cleanup_tables):
+        """Test reorg handling preserves data from unaffected networks"""
+        import json
+
+        from src.amp.streaming.types import BlockRange
+
+        cleanup_tables.append(test_table_name)
+        loader = SnowflakeLoader(snowflake_config)
+
+        with loader:
+            # Create data from multiple networks
+            block_ranges = [
+                [{'network': 'ethereum', 'start': 100, 'end': 110}],
+                [{'network': 'polygon', 'start': 100, 'end': 110}],
+                [{'network': 'ethereum', 'start': 150, 'end': 160}],
+                [{'network': 'polygon', 'start': 150, 'end': 160}],
+            ]
+
+            data = pa.table(
+                {
+                    'id': [1, 2, 3, 4],
+                    'network': ['ethereum', 'polygon', 'ethereum', 'polygon'],
+                    '_meta_block_ranges': [json.dumps([r]) for r in block_ranges],
+                }
+            )
+
+            # Load initial data
+            result = loader.load_table(data, test_table_name, create_table=True)
+            assert result.success
+            assert result.rows_loaded == 4
+
+            # Reorg only ethereum from block 150
+            invalidation_ranges = [BlockRange(network='ethereum', start=150, end=200)]
+            loader._handle_reorg(invalidation_ranges, test_table_name)
+
+            # Verify ethereum row 3 deleted, but polygon rows preserved
+            loader.cursor.execute(f'SELECT id FROM {test_table_name} ORDER BY id')
+            remaining_ids = [row['ID'] for row in loader.cursor.fetchall()]
+            assert remaining_ids == [1, 2, 4]  # Row 3 deleted
+
+    def test_handle_reorg_overlapping_ranges(self, snowflake_config, test_table_name, cleanup_tables):
+        """Test reorg with overlapping block ranges"""
+        import json
+
+        from src.amp.streaming.types import BlockRange
+
+        cleanup_tables.append(test_table_name)
+        loader = SnowflakeLoader(snowflake_config)
+
+        with loader:
+            # Create data with overlapping ranges
+            block_ranges = [
+                [{'network': 'ethereum', 'start': 90, 'end': 110}],  # Overlaps with reorg
+                [{'network': 'ethereum', 'start': 140, 'end': 160}],  # Overlaps with reorg
+                [{'network': 'ethereum', 'start': 170, 'end': 190}],  # After reorg
+            ]
+
+            data = pa.table({'id': [1, 2, 3], '_meta_block_ranges': [json.dumps(ranges) for ranges in block_ranges]})
+
+            # Load initial data
+            result = loader.load_table(data, test_table_name, create_table=True)
+            assert result.success
+            assert result.rows_loaded == 3
+
+            # Reorg from block 150 - should delete rows where end >= 150
+            invalidation_ranges = [BlockRange(network='ethereum', start=150, end=200)]
+            loader._handle_reorg(invalidation_ranges, test_table_name)
+
+            # Only first row should remain (ends at 110 < 150)
+            loader.cursor.execute(f'SELECT COUNT(*) FROM {test_table_name}')
+            count = loader.cursor.fetchone()['COUNT(*)']
+            assert count == 1
+
+            loader.cursor.execute(f'SELECT id FROM {test_table_name}')
+            remaining_id = loader.cursor.fetchone()['ID']
+            assert remaining_id == 1
+
+    def test_streaming_with_reorg(self, snowflake_config, test_table_name, cleanup_tables):
+        """Test streaming data with reorg support"""
+        from src.amp.streaming.types import BatchMetadata, BlockRange, ResponseBatch, ResponseBatchWithReorg
+
+        cleanup_tables.append(test_table_name)
+        loader = SnowflakeLoader(snowflake_config)
+
+        with loader:
+            # Create streaming data with metadata
+            data1 = pa.RecordBatch.from_pydict({'id': [1, 2], 'value': [100, 200]})
+
+            data2 = pa.RecordBatch.from_pydict({'id': [3, 4], 'value': [300, 400]})
+
+            # Create response batches
+            response1 = ResponseBatchWithReorg(
+                is_reorg=False,
+                data=ResponseBatch(
+                    data=data1, metadata=BatchMetadata(ranges=[BlockRange(network='ethereum', start=100, end=110)])
+                ),
+            )
+
+            response2 = ResponseBatchWithReorg(
+                is_reorg=False,
+                data=ResponseBatch(
+                    data=data2, metadata=BatchMetadata(ranges=[BlockRange(network='ethereum', start=150, end=160)])
+                ),
+            )
+
+            # Simulate reorg event
+            reorg_response = ResponseBatchWithReorg(
+                is_reorg=True, invalidation_ranges=[BlockRange(network='ethereum', start=150, end=200)]
+            )
+
+            # Process streaming data
+            stream = [response1, response2, reorg_response]
+            results = list(loader.load_stream_continuous(iter(stream), test_table_name))
+
+            # Verify results
+            assert len(results) == 3
+            assert results[0].success
+            assert results[0].rows_loaded == 2
+            assert results[1].success
+            assert results[1].rows_loaded == 2
+            assert results[2].success
+            assert results[2].is_reorg
+
+            # Verify reorg deleted the second batch
+            loader.cursor.execute(f'SELECT id FROM {test_table_name} ORDER BY id')
+            remaining_ids = [row['ID'] for row in loader.cursor.fetchall()]
+            assert remaining_ids == [1, 2]  # 3 and 4 deleted by reorg