WIP: snowpipe fixes/tests and label manager

Author: fordN

src/amp/client.py

@@ -7,6 +7,7 @@
 
 from . import FlightSql_pb2
 from .config.connection_manager import ConnectionManager
+from .config.label_manager import LabelManager
 from .loaders.registry import create_loader, get_available_loaders
 from .loaders.types import LoadConfig, LoadMode, LoadResult
 from .streaming import (
@@ -105,6 +106,7 @@ class Client:
     def __init__(self, url):
         self.conn = flight.connect(url)
         self.connection_manager = ConnectionManager()
+        self.label_manager = LabelManager()
         self.logger = logging.getLogger(__name__)
 
     def sql(self, query: str) -> QueryBuilder:
@@ -123,6 +125,18 @@ def configure_connection(self, name: str, loader: str, config: Dict[str, Any]) -
         """Configure a named connection for reuse"""
         self.connection_manager.add_connection(name, loader, config)
 
+    def configure_label(self, name: str, csv_path: str, binary_columns: Optional[List[str]] = None) -> None:
+        """
+        Configure a label dataset from a CSV file for joining with streaming data.
+
+        Args:
+            name: Unique name for this label dataset
+            csv_path: Path to the CSV file
+            binary_columns: List of column names containing hex addresses to convert to binary.
+                          If None, auto-detects columns with 'address' in the name.
+        """
+        self.label_manager.add_label(name, csv_path, binary_columns)
+
     def list_connections(self) -> Dict[str, str]:
         """List all configured connections"""
         return self.connection_manager.list_connections()
@@ -238,7 +252,7 @@ def _load_table(
     ) -> LoadResult:
         """Load a complete Arrow Table"""
         try:
-            loader_instance = create_loader(loader, config)
+            loader_instance = create_loader(loader, config, label_manager=self.label_manager)
 
             with loader_instance:
                 return loader_instance.load_table(table, table_name, **load_config.__dict__, **kwargs)
@@ -265,7 +279,7 @@ def _load_stream(
     ) -> Iterator[LoadResult]:
         """Load from a stream of batches"""
         try:
-            loader_instance = create_loader(loader, config)
+            loader_instance = create_loader(loader, config, label_manager=self.label_manager)
 
             with loader_instance:
                 yield from loader_instance.load_stream(batch_stream, table_name, **load_config.__dict__, **kwargs)
@@ -355,7 +369,7 @@ 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)
+        loader_instance = create_loader(loader_type, loader_config, label_manager=self.label_manager)
 
         # Load checkpoint and create resume watermark if enabled (default: enabled)
         if resume_watermark is None and kwargs.get('resume', True):

src/amp/loaders/base.py

@@ -50,11 +50,12 @@ class DataLoader(ABC, Generic[TConfig]):
     REQUIRES_SCHEMA_MATCH: bool = True
     SUPPORTS_TRANSACTIONS: bool = False
 
-    def __init__(self, config: Dict[str, Any]) -> None:
+    def __init__(self, config: Dict[str, Any], label_manager=None) -> None:
         self.logger: Logger = logging.getLogger(f'{self.__class__.__name__}')
         self._connection: Optional[Any] = None
         self._is_connected: bool = False
         self._created_tables: Set[str] = set()  # Track created tables
+        self.label_manager = label_manager  # For CSV label joining
 
         # Parse configuration into typed format
         self.config: TConfig = self._parse_config(config)
@@ -240,6 +241,7 @@ def _try_load_batch(self, batch: pa.RecordBatch, table_name: str, **kwargs) -> L
         This is called by load_batch() within the retry loop. It handles:
         - Connection management
         - Mode validation
+        - Label joining (if configured)
         - Table creation
         - Error handling and timing
         - Metadata generation
@@ -258,7 +260,26 @@ def _try_load_batch(self, batch: pa.RecordBatch, table_name: str, **kwargs) -> L
             if mode not in self.SUPPORTED_MODES:
                 raise ValueError(f'Unsupported mode {mode}. Supported modes: {self.SUPPORTED_MODES}')
 
-            # Handle table creation
+            # Apply label joining if requested
+            label_name = kwargs.get('label')
+            label_key_column = kwargs.get('label_key_column')
+            stream_key_column = kwargs.get('stream_key_column')
+
+            if label_name or label_key_column or stream_key_column:
+                # If any label param is provided, all must be provided
+                if not (label_name and label_key_column and stream_key_column):
+                    raise ValueError(
+                        'Label joining requires all three parameters: label, label_key_column, stream_key_column'
+                    )
+
+                # Perform the join
+                batch = self._join_with_labels(batch, label_name, label_key_column, stream_key_column)
+                self.logger.debug(
+                    f'Joined batch with label {label_name}: {batch.num_rows} rows after join '
+                    f'(columns: {", ".join(batch.schema.names)})'
+                )
+
+            # Handle table creation (use joined schema if applicable)
             if kwargs.get('create_table', True) and table_name not in self._created_tables:
                 if hasattr(self, '_create_table_from_schema'):
                     self._create_table_from_schema(batch.schema, table_name)
@@ -891,6 +912,133 @@ def _get_loader_table_metadata(
         """Override in subclasses to add loader-specific table metadata"""
         return {}
 
+    def _get_effective_schema(
+        self, original_schema: pa.Schema, label_name: Optional[str], label_key_column: Optional[str]
+    ) -> pa.Schema:
+        """
+        Get effective schema by merging label columns into original schema.
+
+        If label_name is None, returns original schema unchanged.
+        Otherwise, merges label columns (excluding the join key which is already in original).
+
+        Args:
+            original_schema: Original data schema
+            label_name: Name of the label dataset (None if no labels)
+            label_key_column: Column name in the label table to join on
+
+        Returns:
+            Schema with label columns merged in
+        """
+        if label_name is None or label_key_column is None:
+            return original_schema
+
+        if self.label_manager is None:
+            raise ValueError('Label manager not configured')
+
+        label_table = self.label_manager.get_label(label_name)
+        if label_table is None:
+            raise ValueError(f"Label '{label_name}' not found")
+
+        # Start with original schema fields
+        merged_fields = list(original_schema)
+
+        # Add label columns (excluding the join key which is already in original)
+        for field in label_table.schema:
+            if field.name != label_key_column and field.name not in original_schema.names:
+                merged_fields.append(field)
+
+        return pa.schema(merged_fields)
+
+    def _join_with_labels(
+        self, batch: pa.RecordBatch, label_name: str, label_key_column: str, stream_key_column: str
+    ) -> pa.RecordBatch:
+        """
+        Join batch data with labels using inner join.
+
+        Handles automatic type conversion between stream and label key columns
+        (e.g., string ↔ binary for Ethereum addresses).
+
+        Args:
+            batch: Original data batch
+            label_name: Name of the label dataset
+            label_key_column: Column name in the label table to join on
+            stream_key_column: Column name in the batch data to join on
+
+        Returns:
+            Joined RecordBatch with label columns added
+
+        Raises:
+            ValueError: If label_manager not configured, label not found, or invalid columns
+        """
+        if self.label_manager is None:
+            raise ValueError('Label manager not configured')
+
+        label_table = self.label_manager.get_label(label_name)
+        if label_table is None:
+            raise ValueError(f"Label '{label_name}' not found")
+
+        # Validate columns exist
+        if stream_key_column not in batch.schema.names:
+            raise ValueError(f"Stream key column '{stream_key_column}' not found in batch schema")
+
+        if label_key_column not in label_table.schema.names:
+            raise ValueError(f"Label key column '{label_key_column}' not found in label table")
+
+        # Convert batch to table for join operation
+        batch_table = pa.Table.from_batches([batch])
+
+        # Get column types for join keys
+        stream_key_type = batch_table.schema.field(stream_key_column).type
+        label_key_type = label_table.schema.field(label_key_column).type
+
+        # If types don't match, cast one to match the other
+        # Prefer casting to binary since that's more efficient
+        import pyarrow.compute as pc
+
+        if stream_key_type != label_key_type:
+            # Try to cast stream key to label key type
+            if pa.types.is_fixed_size_binary(label_key_type) and pa.types.is_string(stream_key_type):
+                # Cast string to binary (hex strings like "0xABCD...")
+                def hex_to_binary(value):
+                    if value is None:
+                        return None
+                    # Remove 0x prefix if present
+                    hex_str = value[2:] if value.startswith('0x') else value
+                    return bytes.fromhex(hex_str)
+
+                # Cast the stream column to binary
+                stream_column = batch_table.column(stream_key_column)
+                binary_length = label_key_type.byte_width
+                binary_values = pa.array(
+                    [hex_to_binary(v.as_py()) for v in stream_column], type=pa.binary(binary_length)
+                )
+                batch_table = batch_table.set_column(
+                    batch_table.schema.get_field_index(stream_key_column), stream_key_column, binary_values
+                )
+            elif pa.types.is_binary(stream_key_type) and pa.types.is_string(label_key_type):
+                # Cast binary to string (for test compatibility)
+                stream_column = batch_table.column(stream_key_column)
+                string_values = pa.array([v.as_py().hex() if v.as_py() else None for v in stream_column])
+                batch_table = batch_table.set_column(
+                    batch_table.schema.get_field_index(stream_key_column), stream_key_column, string_values
+                )
+
+        # Perform inner join using PyArrow compute
+        # Inner join will filter out rows where stream key doesn't match any label key
+        joined_table = batch_table.join(
+            label_table, keys=stream_key_column, right_keys=label_key_column, join_type='inner'
+        )
+
+        # Convert back to RecordBatch
+        if joined_table.num_rows == 0:
+            # Empty result - return empty batch with joined schema
+            # Need to create empty arrays for each column
+            empty_data = {field.name: pa.array([], type=field.type) for field in joined_table.schema}
+            return pa.RecordBatch.from_pydict(empty_data, schema=joined_table.schema)
+
+        # Return as a single batch (assuming batch sizes are manageable)
+        return joined_table.to_batches()[0]
+
     def __enter__(self) -> 'DataLoader':
         self.connect()
         return self

src/amp/loaders/implementations/snowflake_loader.py

@@ -351,8 +351,8 @@ class SnowflakeLoader(DataLoader[SnowflakeConnectionConfig]):
     REQUIRES_SCHEMA_MATCH = False
     SUPPORTS_TRANSACTIONS = True
 
-    def __init__(self, config: Dict[str, Any]) -> None:
-        super().__init__(config)
+    def __init__(self, config: Dict[str, Any], label_manager=None) -> None:
+        super().__init__(config, label_manager=label_manager)
         self.connection: Optional[SnowflakeConnection] = None
         self.cursor = None
         self._created_tables = set()  # Track created tables
@@ -625,9 +625,9 @@ def disconnect(self) -> None:
             for channel_key, channel in self.streaming_channels.items():
                 try:
                     channel.close()
-                    self.logger.debug(f'Closed channel: {channel.name}')
+                    self.logger.debug(f'Closed channel: {channel_key}')
                 except Exception as e:
-                    self.logger.warning(f'Error closing channel: {e}')
+                    self.logger.warning(f'Error closing channel {channel_key}: {e}')
 
             self.streaming_channels.clear()
 
@@ -736,13 +736,19 @@ def _load_via_stage(self, batch: pa.RecordBatch, table_name: str) -> int:
 
         # Identify binary columns and convert to hex for CSV compatibility
         binary_columns = {}
+        # Track VARIANT columns so we can use PARSE_JSON in COPY INTO
+        variant_columns = set()
         modified_arrays = []
         modified_fields = []
 
         t_conversion_start = time.time()
         for i, field in enumerate(batch.schema):
             col_array = batch.column(i)
 
+            # Track _meta_block_ranges as VARIANT column for JSON parsing
+            if field.name == '_meta_block_ranges':
+                variant_columns.add(field.name)
+
             # Check if this is a binary type that needs hex encoding
             if pa.types.is_binary(field.type) or pa.types.is_large_binary(field.type) or pa.types.is_fixed_size_binary(field.type):
                 binary_columns[field.name] = field.type
@@ -801,12 +807,15 @@ def _load_via_stage(self, batch: pa.RecordBatch, table_name: str) -> int:
         t_put_end = time.time()
         self.logger.debug(f'PUT command took {t_put_end - t_put_start:.2f}s')
 
-        # Build column list with transformations - convert hex strings back to binary
+        # Build column list with transformations - convert hex strings back to binary, parse JSON for VARIANT
         final_column_specs = []
         for i, field in enumerate(batch.schema, start=1):
             if field.name in binary_columns:
                 # Use TO_BINARY to convert hex string back to binary
                 final_column_specs.append(f'TO_BINARY(${i}, \'HEX\')')
+            elif field.name in variant_columns:
+                # Use PARSE_JSON to convert JSON string to VARIANT
+                final_column_specs.append(f'PARSE_JSON(${i})')
             else:
                 final_column_specs.append(f'${i}')
 
@@ -1468,9 +1477,9 @@ def _handle_reorg(self, invalidation_ranges: List[BlockRange], table_name: str)
                         try:
                             channel.close()
                             del self.streaming_channels[channel_key]
-                            self.logger.debug(f'Closed streaming channel: {channel.name}')
+                            self.logger.debug(f'Closed streaming channel: {channel_key}')
                         except Exception as e:
-                            self.logger.warning(f'Error closing channel {channel.name}: {e}')
+                            self.logger.warning(f'Error closing channel {channel_key}: {e}')
                             # Continue closing other channels even if one fails
 
                     self.logger.info(
@@ -1482,7 +1491,7 @@ def _handle_reorg(self, invalidation_ranges: List[BlockRange], table_name: str)
                 """
                 SELECT COUNT(*) as count
                 FROM INFORMATION_SCHEMA.COLUMNS
-                WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND COLUMN_NAME = '_META_BLOCK_RANGES'
+                WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND COLUMN_NAME = '_meta_block_ranges'
                 """,
                 (self.config.schema, table_name.upper()),
             )
@@ -1494,32 +1503,33 @@ def _handle_reorg(self, invalidation_ranges: List[BlockRange], table_name: str)
                 )
                 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 = []
+            # Build WHERE conditions for FLATTEN-based deletion
+            # Since Snowflake doesn't support complex subqueries in DELETE WHERE,
+            # we use a CTE-based approach with row identification
+            where_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}
+                where_conditions.append(f"""
+                    (f.value:network::STRING = '{network}' AND f.value:end::NUMBER >= {reorg_start})
+                """)
+
+            if where_conditions:
+                # Use a CTE to identify rows to delete, then delete using METADATA$ROW_ID
+                where_clause = ' OR '.join(where_conditions)
+
+                # Create DELETE SQL using CTE for row identification
+                delete_sql = f"""
+                DELETE FROM {table_name}
+                WHERE "_meta_block_ranges" IN (
+                    SELECT DISTINCT "_meta_block_ranges"
+                    FROM {table_name}, LATERAL FLATTEN(input => "_meta_block_ranges") f
+                    WHERE {where_clause}
                 )
                 """
-                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 '