Formatting

Author: fordN

apps/snowflake_parallel_loader.py

@@ -61,9 +61,7 @@ def configure_logging(verbose: bool = False):
     """
     # Configure root logger first
     logging.basicConfig(
-        level=logging.INFO,
-        format='%(asctime)s | %(levelname)s | %(name)s | %(message)s',
-        datefmt='%Y-%m-%d %H:%M:%S'
+        level=logging.INFO, format='%(asctime)s | %(levelname)s | %(name)s | %(message)s', datefmt='%Y-%m-%d %H:%M:%S'
     )
 
     if not verbose:
@@ -223,8 +221,16 @@ def print_configuration(args, min_block: int, max_block: int, has_labels: bool):
         print(f'🏷️  Label Joining: ENABLED ({args.label_name})')
 
 
-def print_results(results, table_name: str, min_block: int, max_block: int,
-                  duration: float, num_workers: int, has_labels: bool, label_columns: str = ''):
+def print_results(
+    results,
+    table_name: str,
+    min_block: int,
+    max_block: int,
+    duration: float,
+    num_workers: int,
+    has_labels: bool,
+    label_columns: str = '',
+):
     """Print execution results and sample queries."""
     # Calculate statistics
     total_rows = sum(r.rows_loaded for r in results if r.success)
@@ -268,131 +274,81 @@ def main():
     parser = argparse.ArgumentParser(
         description='Load data into Snowflake using parallel streaming with custom SQL queries',
         formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog=__doc__
+        epilog=__doc__,
     )
 
     # Required arguments
     required = parser.add_argument_group('required arguments')
-    required.add_argument(
-        '--query-file',
-        required=True,
-        help='Path to SQL query file to execute'
-    )
-    required.add_argument(
-        '--table-name',
-        required=True,
-        help='Destination Snowflake table name'
-    )
+    required.add_argument('--query-file', required=True, help='Path to SQL query file to execute')
+    required.add_argument('--table-name', required=True, help='Destination Snowflake table name')
 
     # Block range arguments (mutually exclusive groups)
     block_range = parser.add_argument_group('block range')
-    block_range.add_argument(
-        '--blocks',
-        type=int,
-        help='Number of recent blocks to load (auto-detect range)'
-    )
-    block_range.add_argument(
-        '--min-block',
-        type=int,
-        help='Explicit start block (requires --max-block)'
-    )
-    block_range.add_argument(
-        '--max-block',
-        type=int,
-        help='Explicit end block (requires --min-block)'
-    )
+    block_range.add_argument('--blocks', type=int, help='Number of recent blocks to load (auto-detect range)')
+    block_range.add_argument('--min-block', type=int, help='Explicit start block (requires --max-block)')
+    block_range.add_argument('--max-block', type=int, help='Explicit end block (requires --min-block)')
     block_range.add_argument(
         '--source-table',
         default='eth_firehose.logs',
-        help='Table for block range detection (default: eth_firehose.logs)'
+        help='Table for block range detection (default: eth_firehose.logs)',
     )
     block_range.add_argument(
-        '--block-column',
-        default='block_num',
-        help='Column name for block partitioning (default: block_num)'
+        '--block-column', default='block_num', help='Column name for block partitioning (default: block_num)'
     )
 
     # Label configuration (all optional)
     labels = parser.add_argument_group('label configuration (optional)')
-    labels.add_argument(
-        '--label-csv',
-        help='Path to CSV file with label data'
-    )
-    labels.add_argument(
-        '--label-name',
-        help='Label identifier (required if --label-csv provided)'
-    )
-    labels.add_argument(
-        '--label-key',
-        help='CSV column for joining (required if --label-csv provided)'
-    )
-    labels.add_argument(
-        '--stream-key',
-        help='Stream column for joining (required if --label-csv provided)'
-    )
+    labels.add_argument('--label-csv', help='Path to CSV file with label data')
+    labels.add_argument('--label-name', help='Label identifier (required if --label-csv provided)')
+    labels.add_argument('--label-key', help='CSV column for joining (required if --label-csv provided)')
+    labels.add_argument('--stream-key', help='Stream column for joining (required if --label-csv provided)')
 
     # Snowflake configuration
     snowflake = parser.add_argument_group('snowflake configuration')
     snowflake.add_argument(
-        '--connection-name',
-        help='Snowflake connection name (default: auto-generated from table name)'
+        '--connection-name', help='Snowflake connection name (default: auto-generated from table name)'
     )
     snowflake.add_argument(
         '--loading-method',
         choices=['snowpipe_streaming', 'stage', 'insert'],
         default='snowpipe_streaming',
-        help='Snowflake loading method (default: snowpipe_streaming)'
+        help='Snowflake loading method (default: snowpipe_streaming)',
     )
     snowflake.add_argument(
         '--preserve-reorg-history',
         action='store_true',
         default=True,
-        help='Enable reorg history preservation (default: enabled)'
+        help='Enable reorg history preservation (default: enabled)',
     )
     snowflake.add_argument(
         '--no-preserve-reorg-history',
         action='store_false',
         dest='preserve_reorg_history',
-        help='Disable reorg history preservation'
-    )
-    snowflake.add_argument(
-        '--disable-state',
-        action='store_true',
-        help='Disable state management (job resumption)'
-    )
-    snowflake.add_argument(
-        '--pool-size',
-        type=int,
-        help='Connection pool size (default: workers + 2)'
+        help='Disable reorg history preservation',
     )
+    snowflake.add_argument('--disable-state', action='store_true', help='Disable state management (job resumption)')
+    snowflake.add_argument('--pool-size', type=int, help='Connection pool size (default: workers + 2)')
 
     # Parallel execution configuration
     parallel = parser.add_argument_group('parallel execution')
-    parallel.add_argument(
-        '--workers',
-        type=int,
-        default=4,
-        help='Number of parallel workers (default: 4)'
-    )
+    parallel.add_argument('--workers', type=int, default=4, help='Number of parallel workers (default: 4)')
     parallel.add_argument(
         '--flush-interval',
         type=float,
         default=1.0,
-        help='Snowpipe Streaming buffer flush interval in seconds (default: 1.0)'
+        help='Snowpipe Streaming buffer flush interval in seconds (default: 1.0)',
     )
 
     # Server configuration
     parser.add_argument(
         '--server',
         default=os.getenv('AMP_SERVER_URL', 'grpc://34.27.238.174:80'),
-        help='AMP server URL (default: from AMP_SERVER_URL env or grpc://34.27.238.174:80)'
+        help='AMP server URL (default: from AMP_SERVER_URL env or grpc://34.27.238.174:80)',
     )
 
     # Logging configuration
     parser.add_argument(
-        '--verbose',
-        action='store_true',
-        help='Enable verbose logging from Snowflake libraries (default: suppressed)'
+        '--verbose', action='store_true', help='Enable verbose logging from Snowflake libraries (default: suppressed)'
     )
 
     args = parser.parse_args()
@@ -445,8 +401,7 @@ def main():
 
         # Print results
         label_columns = f'{args.label_key} joined columns' if has_labels else ''
-        print_results(results, args.table_name, min_block, max_block, duration,
-                     args.workers, has_labels, label_columns)
+        print_results(results, args.table_name, min_block, max_block, duration, args.workers, has_labels, label_columns)
 
         return args.table_name, sum(r.rows_loaded for r in results if r.success), duration
 
@@ -456,6 +411,7 @@ def main():
     except Exception as e:
         print(f'\n\n❌ Error: {e}')
         import traceback
+
         traceback.print_exc()
         sys.exit(1)
 

apps/test_erc20_labeled_parallel.py

@@ -55,9 +55,7 @@ def get_recent_block_range(client: Client, num_blocks: int = 100_000):
     return min_block, max_block
 
 
-def load_erc20_transfers_with_labels(
-    num_blocks: int = 100_000, num_workers: int = 4, flush_interval: float = 1.0
-):
+def load_erc20_transfers_with_labels(num_blocks: int = 100_000, num_workers: int = 4, flush_interval: float = 1.0):
     """Load ERC20 transfers with token labels using Snowpipe Streaming and parallel streaming."""
 
     # Initialize client

src/amp/client.py

@@ -34,7 +34,7 @@ def load(
         destination: str,
         config: Dict[str, Any] = None,
         label_config: Optional[LabelJoinConfig] = None,
-        **kwargs
+        **kwargs,
     ) -> Union[LoadResult, Iterator[LoadResult]]:
         """
         Load query results to specified destination

src/amp/loaders/base.py

@@ -6,7 +6,6 @@
 import time
 from abc import ABC, abstractmethod
 from dataclasses import fields, is_dataclass
-from datetime import UTC, datetime
 from logging import Logger
 from typing import Any, Dict, Generic, Iterator, List, Optional, Set, TypeVar
 
@@ -261,10 +260,7 @@ def _try_load_batch(self, batch: pa.RecordBatch, table_name: str, **kwargs) -> L
             if label_config:
                 # Perform the join
                 batch = self._join_with_labels(
-                    batch,
-                    label_config.label_name,
-                    label_config.label_key_column,
-                    label_config.stream_key_column
+                    batch, label_config.label_name, label_config.label_key_column, label_config.stream_key_column
                 )
                 self.logger.debug(
                     f'Joined batch with label {label_config.label_name}: {batch.num_rows} rows after join '
@@ -478,9 +474,7 @@ def load_stream_continuous(
 
                     # Choose processing strategy: transactional vs non-transactional
                     use_transactional = (
-                        hasattr(self, 'load_batch_transactional')
-                        and self.state_enabled
-                        and response.metadata.ranges
+                        hasattr(self, 'load_batch_transactional') and self.state_enabled and response.metadata.ranges
                     )
 
                     if use_transactional:
@@ -636,9 +630,7 @@ def _process_batch_transactional(
         try:
             # Delegate to loader-specific transactional implementation
             # Loaders that support transactions implement load_batch_transactional()
-            rows_loaded_batch = self.load_batch_transactional(
-                batch_data, table_name, connection_name, ranges
-            )
+            rows_loaded_batch = self.load_batch_transactional(batch_data, table_name, connection_name, ranges)
             duration = time.time() - start_time
 
             # Mark batches as processed in state store after successful transaction
@@ -703,7 +695,9 @@ def _process_batch_non_transactional(
 
                 if is_duplicate:
                     # Skip this batch - already processed
-                    self.logger.info(f'Skipping duplicate batch: {len(ranges)} ranges already processed for {table_name}')
+                    self.logger.info(
+                        f'Skipping duplicate batch: {len(ranges)} ranges already processed for {table_name}'
+                    )
                     return LoadResult(
                         rows_loaded=0,
                         duration=0.0,
@@ -731,7 +725,6 @@ def _process_batch_non_transactional(
 
         return result
 
-
     def _augment_streaming_result(
         self, result: LoadResult, batch_count: int, ranges: Optional[List[BlockRange]], ranges_complete: bool
     ) -> LoadResult:
@@ -808,23 +801,26 @@ def _add_metadata_columns(self, data: pa.RecordBatch, block_ranges: List[BlockRa
             # Convert BlockRanges to BatchIdentifiers and get compact unique IDs
             batch_ids = [BatchIdentifier.from_block_range(br) for br in block_ranges]
             # Combine multiple batch IDs with "|" separator for multi-network batches
-            batch_id_str = "|".join(bid.unique_id for bid in batch_ids)
+            batch_id_str = '|'.join(bid.unique_id for bid in batch_ids)
             batch_id_array = pa.array([batch_id_str] * num_rows, type=pa.string())
             result = result.append_column('_amp_batch_id', batch_id_array)
 
         # Optionally add full JSON for debugging/auditing
         if self.store_full_metadata:
             import json
-            ranges_json = json.dumps([
-                {
-                    'network': br.network,
-                    'start': br.start,
-                    'end': br.end,
-                    'hash': br.hash,
-                    'prev_hash': br.prev_hash
-                }
-                for br in block_ranges
-            ])
+
+            ranges_json = json.dumps(
+                [
+                    {
+                        'network': br.network,
+                        'start': br.start,
+                        'end': br.end,
+                        'hash': br.hash,
+                        'prev_hash': br.prev_hash,
+                    }
+                    for br in block_ranges
+                ]
+            )
             ranges_array = pa.array([ranges_json] * num_rows, type=pa.string())
             result = result.append_column('_amp_block_ranges', ranges_array)
 
@@ -966,7 +962,6 @@ def _join_with_labels(
 
         # 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
 
         type_conversion_time_ms = 0.0
         if stream_key_type != label_key_type:
@@ -1032,14 +1027,14 @@ def hex_to_binary(value):
             timing_msg = (
                 f'⏱️  Label join: {input_rows} → {output_rows} rows in {total_time_ms:.2f}ms '
                 f'(type_conv={type_conversion_time_ms:.2f}ms, join={join_time_ms:.2f}ms, '
-                f'{output_rows/total_time_ms*1000:.0f} rows/sec) '
-                f'[label={label_name}, retained={output_rows/input_rows*100:.1f}%]\n'
+                f'{output_rows / total_time_ms * 1000:.0f} rows/sec) '
+                f'[label={label_name}, retained={output_rows / input_rows * 100:.1f}%]\n'
             )
         else:
             timing_msg = (
                 f'⏱️  Label join: {input_rows} → {output_rows} rows in {total_time_ms:.2f}ms '
-                f'(join={join_time_ms:.2f}ms, {output_rows/total_time_ms*1000:.0f} rows/sec) '
-                f'[label={label_name}, retained={output_rows/input_rows*100:.1f}%]\n'
+                f'(join={join_time_ms:.2f}ms, {output_rows / total_time_ms * 1000:.0f} rows/sec) '
+                f'[label={label_name}, retained={output_rows / input_rows * 100:.1f}%]\n'
             )
 
         sys.stderr.write(timing_msg)

src/amp/loaders/implementations/deltalake_loader.py

@@ -1,6 +1,5 @@
 # src/amp/loaders/implementations/deltalake_loader.py
 
-import json
 import os
 import time
 from dataclasses import dataclass, field

src/amp/loaders/implementations/lmdb_loader.py

@@ -1,7 +1,6 @@
 # amp/loaders/implementations/lmdb_loader.py
 
 import hashlib
-import json
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Dict, List, Optional