import boto3 from typing import List, Dict, Any, Optional, Iterator from botocore.exceptions import ClientError, NoCredentialsError from botocore.config import Config import logging from concurrent.futures import ThreadPoolExecutor, as_completed from threading import Lock import time from tqdm import tqdm from collections import defaultdict # Configure logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) # DynamoDB Table Configuration TABLE_NAME = "clip-gen-config" REGION = "us-east-2" # Performance tuning parameters MAX_WORKERS = 50 # Number of parallel threads BATCH_SIZE = 1000 # Process UUIDs in batches CONNECTION_POOL_SIZE = 50 # Connection pool size MAX_RETRIES = 3 # Max retries for failed queries RETRY_DELAY = 0.1 # Initial retry delay in seconds class OptimizedDynamoDBQuerier: """Optimized DynamoDB querier with connection pooling and parallel processing.""" def __init__(self, profile_name: Optional[str] = "default"): """Initialize the querier with connection pool.""" # Configure boto3 client with connection pooling config = Config( region_name=REGION, max_pool_connections=CONNECTION_POOL_SIZE, retries={"max_attempts": MAX_RETRIES, "mode": "adaptive"}, ) # Create session and client if profile_name: session = boto3.Session(profile_name=profile_name) self.client = session.client("dynamodb", config=config) else: self.client = boto3.client("dynamodb", config=config) # Thread-safe results storage self.results_lock = Lock() self.failed_uuids_lock = Lock() def _query_single_uuid(self, uuid: str) -> List[Dict[str, Any]]: """Query a single UUID with retry logic.""" for attempt in range(MAX_RETRIES): try: response = self.client.query( TableName=TABLE_NAME, KeyConditionExpression="clipId = :clipId", ExpressionAttributeValues={":clipId": {"S": uuid}}, ConsistentRead=True, ) # Process and return the items items = [] for item in response.get("Items", []): converted_item = _convert_dynamodb_item(item) items.append(converted_item) return items except ClientError as e: if attempt < MAX_RETRIES - 1: time.sleep(RETRY_DELAY * (2**attempt)) # Exponential backoff continue logger.error( f"Error querying clipId {uuid} after {MAX_RETRIES} attempts: {e}" ) raise def query_dynamodb_parallel( self, uuids: List[str], show_progress: bool = True ) -> List[Dict[str, Any]]: """ Query DynamoDB in parallel for all UUIDs. Args: uuids: List of UUID strings show_progress: Show progress bar Returns: List of all records from DynamoDB """ if not uuids: raise ValueError("UUID list cannot be empty") logger.info( f"Starting parallel query for {len(uuids)} UUIDs with {MAX_WORKERS} workers" ) all_results = [] failed_uuids = [] # Create progress bar if requested pbar = ( tqdm(total=len(uuids), desc="Querying DynamoDB") if show_progress else None ) with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor: # Submit all tasks future_to_uuid = { executor.submit(self._query_single_uuid, uuid): uuid for uuid in uuids } # Process completed tasks for future in as_completed(future_to_uuid): uuid = future_to_uuid[future] try: items = future.result() with self.results_lock: all_results.extend(items) except Exception as e: logger.error(f"Failed to query UUID {uuid}: {e}") with self.failed_uuids_lock: failed_uuids.append(uuid) if pbar: pbar.update(1) if pbar: pbar.close() if failed_uuids: logger.warning(f"Failed to query {len(failed_uuids)} UUIDs") logger.info(f"Retrieved {len(all_results)} total records from DynamoDB") return all_results def query_dynamodb_in_batches( self, uuids: List[str], batch_callback: Optional[callable] = None, save_intermediate: bool = False, intermediate_file_prefix: str = "dynamo_results_batch", ) -> List[Dict[str, Any]]: """ Query DynamoDB in batches to handle very large UUID lists. Args: uuids: List of UUID strings batch_callback: Optional callback function called after each batch save_intermediate: Save intermediate results to files intermediate_file_prefix: Prefix for intermediate files Returns: List of all records from DynamoDB """ import json all_results = [] total_batches = (len(uuids) + BATCH_SIZE - 1) // BATCH_SIZE logger.info( f"Processing {len(uuids)} UUIDs in {total_batches} batches of {BATCH_SIZE}" ) for batch_idx in range(0, len(uuids), BATCH_SIZE): batch_uuids = uuids[batch_idx : batch_idx + BATCH_SIZE] batch_num = batch_idx // BATCH_SIZE + 1 logger.info(f"Processing batch {batch_num}/{total_batches}") # Query this batch batch_results = self.query_dynamodb_parallel( batch_uuids, show_progress=True ) all_results.extend(batch_results) # Save intermediate results if requested if save_intermediate: filename = f"{intermediate_file_prefix}_{batch_num}.json" with open(filename, "w") as f: json.dump(batch_results, f) logger.info(f"Saved batch {batch_num} to {filename}") # Call callback if provided if batch_callback: batch_callback(batch_num, batch_results) # Small delay between batches to avoid overwhelming DynamoDB if batch_num < total_batches: time.sleep(0.5) return all_results def _convert_dynamodb_item(item: Dict[str, Any]) -> Dict[str, Any]: """Convert DynamoDB item format to regular Python dict.""" converted = {} for key, value in item.items(): if "S" in value: # String converted[key] = value["S"] elif "N" in value: # Number try: converted[key] = int(value["N"]) except ValueError: converted[key] = float(value["N"]) elif "B" in value: # Binary converted[key] = value["B"] elif "BOOL" in value: # Boolean converted[key] = value["BOOL"] elif "NULL" in value: # Null converted[key] = None elif "L" in value: # List converted[key] = [_convert_dynamodb_value(v) for v in value["L"]] elif "M" in value: # Map converted[key] = _convert_dynamodb_item(value["M"]) elif "SS" in value: # String Set converted[key] = value["SS"] elif "NS" in value: # Number Set converted[key] = [float(n) for n in value["NS"]] elif "BS" in value: # Binary Set converted[key] = value["BS"] return converted def _convert_dynamodb_value(value: Dict[str, Any]) -> Any: """Helper function to convert individual DynamoDB values.""" if "S" in value: return value["S"] elif "N" in value: try: return int(value["N"]) except ValueError: return float(value["N"]) elif "B" in value: return value["B"] elif "BOOL" in value: return value["BOOL"] elif "NULL" in value: return None elif "L" in value: return [_convert_dynamodb_value(v) for v in value["L"]] elif "M" in value: return _convert_dynamodb_item(value["M"]) elif "SS" in value: return value["SS"] elif "NS" in value: return [float(n) for n in value["NS"]] elif "BS" in value: return value["BS"] return value def query_dynamodb_by_uuids_optimized( uuids: List[str], profile_name: Optional[str] = "default", use_batches: bool = True, save_intermediate: bool = False, ) -> List[Dict[str, Any]]: """ Optimized function to query DynamoDB for large numbers of UUIDs. Args: uuids: List of UUID strings profile_name: AWS profile name use_batches: Process in batches (recommended for > 10k UUIDs) save_intermediate: Save intermediate results to files Returns: List of all records from DynamoDB """ querier = OptimizedDynamoDBQuerier(profile_name=profile_name) if use_batches and len(uuids) > 10000: return querier.query_dynamodb_in_batches( uuids, save_intermediate=save_intermediate ) else: return querier.query_dynamodb_parallel(uuids) def main(): """Demo usage with performance comparison.""" # Test with a larger set of UUIDs test_uuids = [ "22acf020-35d4-4412-9d5d-907fea873bca", "8ded8a65-4a77-4b98-97d7-4e8f49118072", ] * 100 # 200 UUIDs for testing # Time the optimized version start_time = time.time() results = query_dynamodb_by_uuids_optimized(test_uuids) elapsed = time.time() - start_time print( f"Optimized version: Retrieved {len(results)} records in {elapsed:.2f} seconds" ) print(f"Average time per UUID: {elapsed/len(test_uuids)*1000:.2f} ms") # For 1M UUIDs, you would use: # results = query_dynamodb_by_uuids_optimized( # your_million_uuids, # use_batches=True, # save_intermediate=True # Save progress in case of interruption # ) if __name__ == "__main__": main()