Snowflake Snowpipe Streaming High-performance Architecture

				
					#!/usr/bin/env python3
"""
Simple Kafka to Snowpipe Streaming - Single File Version

Streamlined version combining Kafka consumer and Snowflake ingestion.
Update configuration and run!

Usage:
    python3 kafka_snowpipe_simple.py
"""

from confluent_kafka import Consumer, KafkaError
import json
from datetime import datetime
import uuid
import os
import time

os.environ["SS_LOG_LEVEL"] = "warn"
from snowflake.ingest.streaming import StreamingIngestClient

# ============================================================================
# CONFIGURATION - UPDATE THESE
# ============================================================================

# Kafka Configuration
KAFKA_BROKER = 'localhost:9092'
KAFKA_TOPIC = 'user'
KAFKA_GROUP = 'snowpipe-consumer'

# Snowflake Configuration
SNOWFLAKE_DB = 'SANDBOX'
SNOWFLAKE_SCHEMA = 'DGANIGER'
SNOWFLAKE_TABLE = 'user_data'
SNOWFLAKE_PIPE = 'user_data_pipe'
PROFILE_JSON = 'profile.json'

# Processing Configuration
BATCH_SIZE = 50  # Number of messages per batch
IDLE_TIMEOUT_SECONDS = 5  # Exit if no messages for this many seconds

# ============================================================================
# GLOBAL STATE
# ============================================================================

class State:
    """Holds global state for processing."""
    def __init__(self):
        self.snowflake_client = None
        self.snowflake_channel = None
        self.kafka_consumer = None
        self.batch = []
        self.total_processed = 0
        self.last_message_time = None

state = State()

# ============================================================================
# SNOWFLAKE SETUP
# ============================================================================

def setup_snowflake():
    """Initialize Snowflake Streaming client and channel."""
    print("Connecting to Snowflake...")
    
    # Create client
    state.snowflake_client = StreamingIngestClient(
        client_name=f"KAFKA_CLIENT_{uuid.uuid4()}",
        db_name=SNOWFLAKE_DB,
        schema_name=SNOWFLAKE_SCHEMA,
        pipe_name=SNOWFLAKE_PIPE,
        profile_json=PROFILE_JSON
    )
    
    # Open channel
    state.snowflake_channel, status = state.snowflake_client.open_channel(
        f"KAFKA_CHANNEL_{uuid.uuid4()}"
    )
    
    print(f"  Snowflake connected: {state.snowflake_channel.channel_name}")
    print(f"   Database: {SNOWFLAKE_DB}")
    print(f"   Schema: {SNOWFLAKE_SCHEMA}")
    print(f"   Table: {SNOWFLAKE_TABLE}\n")

# ============================================================================
# KAFKA SETUP
# ============================================================================

def setup_kafka():
    """Initialize Kafka consumer."""
    print(" Connecting to Kafka...")
    
    config = {
        'bootstrap.servers': KAFKA_BROKER,
        'group.id': KAFKA_GROUP,
        'auto.offset.reset': 'earliest',
        'enable.auto.commit': True,
        'session.timeout.ms': 6000,
    }
    
    state.kafka_consumer = Consumer(config)
    state.kafka_consumer.subscribe([KAFKA_TOPIC])
    
    print(f"  Kafka connected: {KAFKA_BROKER}")
    print(f"   Topic: {KAFKA_TOPIC}\n")

# ============================================================================
# MESSAGE PROCESSING
# ============================================================================

def transform_message(msg):
    """
    Transform Kafka message to Snowflake row.
    
    Creates a JSON with structure:
    {
        "topic": "topic_name",
        "offset": 12345,
        "ingestion_timestamp": "2025-10-21T14:30:22.123456",
        "data": { ... original kafka message ... }
    }
    """
    try:
        # Parse JSON from Kafka
        kafka_data = json.loads(msg.value().decode('utf-8'))
    except json.JSONDecodeError:
        # Handle non-JSON messages - wrap as string
        kafka_data = msg.value().decode('utf-8', errors='ignore')
    
    # Create row with the requested structure
    # Best Practice - Include metadata fields like topic, offset and ingestion timestamp to track
    row = {
        'topic': msg.topic(),
        'offset': msg.offset(),
        'ingestion_timestamp': datetime.now().isoformat(),
        'data': kafka_data
    }
    
    return row

def process_message(msg):
    """Process each Kafka message."""
    try:
        # Update last message time
        state.last_message_time = time.time()
        
        # Transform message
        row = transform_message(msg)
        
        # Add to batch
        state.batch.append(row)
        state.total_processed += 1
        
        print(f" Message {state.total_processed} buffered (batch: {len(state.batch)}/{BATCH_SIZE})")
        
        # Flush when batch is full
        if len(state.batch) >= BATCH_SIZE:
            flush_to_snowflake()
            
    except Exception as e:
        print(f" Error processing message: {e}")

def check_idle_timeout():
    """Check if idle timeout has been reached."""
    if state.last_message_time is None:
        return False
    
    idle_time = time.time() - state.last_message_time
    if idle_time >= IDLE_TIMEOUT_SECONDS:
        print(f"\n  No messages received for {idle_time:.1f} seconds")
        print(f"   Idle timeout ({IDLE_TIMEOUT_SECONDS}s) reached - initiating graceful shutdown...")
        return True
    return False

def flush_to_snowflake():
    """Flush batch to Snowflake."""
    if not state.batch:
        return
        
    print(f"\n Flushing {len(state.batch)} messages to Snowflake...")
    
    try:
        # Append rows to Snowflake channel
        for i, row in enumerate(state.batch):
            # Use topic and offset to create unique row ID
            row_id = f"{row.get('topic', 'unknown')}_{row.get('offset', i)}"
            state.snowflake_channel.append_row(row, row_id)
        
        print(f" Batch ingested! Total processed: {state.total_processed}\n")
        
        # Clear batch
        state.batch = []
        
    except Exception as e:
        print(f" Error flushing to Snowflake: {e}")
        raise

# ============================================================================
# CLEANUP
# Flush any remaining messages and close kafka connection 
# ============================================================================

def cleanup():
    """Cleanup resources."""
    print("\n Cleaning up...")
    
    # Flush remaining messages
    if state.batch:
        print(f"   Flushing remaining {len(state.batch)} messages...")
        flush_to_snowflake()
    
    # Close Kafka consumer
    if state.kafka_consumer:
        state.kafka_consumer.close()
        print("   Kafka consumer closed")
    
    # Close Snowflake resources
    if state.snowflake_channel:
        state.snowflake_channel.close()
        print("Snowflake channel closed")
    
    if state.snowflake_client:
        state.snowflake_client.close()
        print("Snowflake client closed")
    
    print(f" Final Stats:")
    print(f"Total messages processed: {state.total_processed}")
    print(f"Messages successfully ingested: {state.total_processed - len(state.batch)}")

# ============================================================================
# MAIN
# ============================================================================

def main():
    """Main function."""
    print("\n" + "="*60)
    print("Kafka → Snowpipe Streaming - Simple Version")
    print("="*60 + "\n")
    
    try:
        # Setup Snowflake and Kafka
        setup_snowflake()
        setup_kafka()
        
        print("="*60)
        print(f" Batch Size: {BATCH_SIZE} messages")
        print(f"  Idle Timeout: {IDLE_TIMEOUT_SECONDS} seconds")
        print(" Press Ctrl+C to stop")
        print("="*60 + "\n")
        
        # Initialize last message time
        state.last_message_time = time.time()
        
        # Start consuming
        print(" Listening for messages...\n")
        
        while True:
            # Poll for a single message with 1 second timeout
            msg = state.kafka_consumer.poll(timeout=1.0)
            
            if msg is None:
                # No message received, check for idle timeout
                if check_idle_timeout():
                    break
                continue
            
            if msg.error():
                if msg.error().code() == KafkaError._PARTITION_EOF:
                    continue
                else:
                    print(f" Kafka error: {msg.error()}")
                    continue
            
            # Process the message
            process_message(msg)
        
        print(" Graceful shutdown initiated")
        
    except KeyboardInterrupt:
        print("  Stopped by user")
        
    except Exception as e:
        print(f"\n Error: {e}")
        import traceback
        traceback.print_exc()
        
    finally:
        cleanup()

if __name__ == '__main__':
    main()