WIP: fixes I don't think we should have needed 🤷

Author: fordN

src/amp/loaders/implementations/snowflake_loader.py

@@ -566,10 +566,12 @@ def _get_or_create_channel(self, table_name: str, channel_suffix: str = 'default
             client = self.streaming_clients[table_name]
 
             # Open channel - returns (channel, status) tuple
-            channel, status = client.open_channel(channel_name=channel_name)
-
-            if status != 'OPEN':
-                raise RuntimeError(f'Failed to open streaming channel {channel_name}: status={status}')
+            # Status indicates channel state but appears to be a channel object representation
+            result = client.open_channel(channel_name=channel_name)
+            if isinstance(result, tuple):
+                channel, status = result
+            else:
+                channel = result
 
             self.streaming_channels[channel_key] = channel
             self.logger.info(f'Opened Snowpipe Streaming channel: {channel_name}')
@@ -584,9 +586,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()
 
@@ -1132,8 +1134,11 @@ def _create_table_from_schema(self, schema: pa.Schema, table_name: str) -> None:
         # Build CREATE TABLE statement
         columns = []
         for field in schema:
+            # Special case: _meta_block_ranges should be VARIANT for optimal JSON querying
+            if field.name == '_meta_block_ranges':
+                snowflake_type = 'VARIANT'
             # Handle complex types
-            if pa.types.is_timestamp(field.type):
+            elif pa.types.is_timestamp(field.type):
                 if field.type.tz is not None:
                     snowflake_type = 'TIMESTAMP_TZ'
                 else:
@@ -1184,7 +1189,7 @@ def _create_table_from_schema(self, schema: pa.Schema, table_name: str) -> None:
     def _get_loader_batch_metadata(self, batch: pa.RecordBatch, duration: float, **kwargs) -> Dict[str, Any]:
         """Get Snowflake-specific metadata for batch operation"""
         return {
-            'loading_method': 'stage' if self.use_stage else 'insert',
+            'loading_method': self.loading_method,
             'warehouse': self.config.warehouse,
             'database': self.config.database,
             'schema': self.config.schema,
@@ -1195,7 +1200,7 @@ def _get_loader_table_metadata(
     ) -> Dict[str, Any]:
         """Get Snowflake-specific metadata for table operation"""
         return {
-            'loading_method': 'stage' if self.use_stage else 'insert',
+            'loading_method': self.loading_method,
             'warehouse': self.config.warehouse,
             'database': self.config.database,
             'schema': self.config.schema,
@@ -1314,9 +1319,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(

src/amp/streaming/reorg.py

@@ -74,6 +74,7 @@ def __next__(self) -> ResponseBatchWithReorg:
                 return ResponseBatchWithReorg.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')

tests/integration/test_snowflake_loader.py

@@ -25,6 +25,7 @@
 try:
     from src.amp.loaders.base import LoadMode
     from src.amp.loaders.implementations.snowflake_loader import SnowflakeLoader
+    from src.amp.streaming.types import BatchMetadata, BlockRange, ResponseBatch, ResponseBatchWithReorg
 except ImportError:
     pytest.skip('amp modules not available', allow_module_level=True)
 
@@ -127,7 +128,8 @@ def test_batch_loading(self, snowflake_config, medium_test_table, test_table_nam
         loader = SnowflakeLoader(snowflake_config)
 
         with loader:
-            result = loader.load_table(medium_test_table, test_table_name, create_table=True)
+            # Use smaller batch size to force multiple batches (medium_test_table has 10000 rows)
+            result = loader.load_table(medium_test_table, test_table_name, create_table=True, batch_size=5000)
 
             assert result.success is True
             assert result.rows_loaded == medium_test_table.num_rows
@@ -333,11 +335,10 @@ def test_stage_and_compression_options(self, snowflake_config, medium_test_table
         """Test different stage and compression options"""
         cleanup_tables.append(test_table_name)
 
-        # Test with different compression
+        # Test with stage loading method
         config = {
             **snowflake_config,
             'loading_method': 'stage',
-            'compression': 'zstd',
         }
         loader = SnowflakeLoader(config)
 
@@ -614,23 +615,19 @@ def test_streaming_with_reorg(self, snowflake_config, test_table_name, cleanup_t
             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)])
-                ),
+            batch1 = ResponseBatch(
+                data=data1, metadata=BatchMetadata(ranges=[BlockRange(network='ethereum', start=100, end=110)])
             )
+            response1 = ResponseBatchWithReorg.data_batch(batch1)
 
-            response2 = ResponseBatchWithReorg(
-                is_reorg=False,
-                data=ResponseBatch(
-                    data=data2, metadata=BatchMetadata(ranges=[BlockRange(network='ethereum', start=150, end=160)])
-                ),
+            batch2 = ResponseBatch(
+                data=data2, metadata=BatchMetadata(ranges=[BlockRange(network='ethereum', start=150, end=160)])
             )
+            response2 = ResponseBatchWithReorg.data_batch(batch2)
 
             # Simulate reorg event
-            reorg_response = ResponseBatchWithReorg(
-                is_reorg=True, invalidation_ranges=[BlockRange(network='ethereum', start=150, end=200)]
+            reorg_response = ResponseBatchWithReorg.reorg_batch(
+                [BlockRange(network='ethereum', start=150, end=200)]
             )
 
             # Process streaming data
@@ -647,8 +644,8 @@ def test_streaming_with_reorg(self, snowflake_config, test_table_name, cleanup_t
             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()]
+            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
 
 
@@ -704,8 +701,8 @@ def test_streaming_connection(self, snowflake_streaming_config):
         loader.connect()
         assert loader._is_connected is True
         assert loader.connection is not None
-        # Streaming client is lazily initialized (created on first load, not at connection time)
-        assert loader.streaming_client is None
+        # Streaming clients dict is initialized empty (clients created on first load per table)
+        assert loader.streaming_clients == {}
 
         loader.disconnect()
         assert loader._is_connected is False