Process IIoT Sensor Streams at the Edge with Bytewax and Polars

                      
                     
"""
Production implementation for processing IIoT sensor streams at the edge.
Provides secure, scalable operations using Bytewax and Polars.
"""
from typing import Dict, Any, List
import os
import logging
import polars as pl
from bytewax.dataflow import DataFlow
from bytewax.dataflow import run
import time

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Configuration class to handle environment variables
class Config:
    database_url: str = os.getenv('DATABASE_URL', 'sqlite:///:memory:')  # Default to in-memory for testing

# Validate input data
async def validate_input(data: Dict[str, Any]) -> bool:
    """Validate incoming sensor data.
    
    Args:
        data: Sensor data to validate
    Returns:
        True if valid
    Raises:
        ValueError: If validation fails
    """
    if 'sensor_id' not in data:
        raise ValueError('Missing sensor_id')  # Ensure sensor_id is present
    if 'value' not in data:
        raise ValueError('Missing value')  # Ensure value is present
    return True

# Sanitize fields in the input data
def sanitize_fields(data: Dict[str, Any]) -> Dict[str, Any]:
    """Sanitize input data fields.
    
    Args:
        data: Input data
    Returns:
        Sanitized data
    """
    return {key: str(value).strip() for key, value in data.items()}  # Strip whitespace

# Normalize data before processing
def normalize_data(data: Dict[str, Any]) -> Dict[str, Any]:
    """Normalize the input data fields.
    
    Args:
        data: Input data
    Returns:
        Normalized data
    """
    # Example normalization logic
    data['value'] = float(data['value'])  # Convert value to float
    return data

# Function to process a batch of records
async def process_batch(records: List[Dict[str, Any]]) -> None:
    """Process a batch of sensor records.
    
    Args:
        records: List of records to process
    """
    for record in records:
        logger.info(f'Processing record: {record}')  # Log processing
        # Simulate processing time
        time.sleep(0.1)

# Aggregate metrics from processed data
def aggregate_metrics(data: pl.DataFrame) -> pl.DataFrame:
    """Aggregate metrics from the processed data.
    
    Args:
        data: DataFrame containing processed data
    Returns:
        Aggregated metrics DataFrame
    """
    return data.groupby('sensor_id').agg(pl.col('value').mean()).collect()  # Mean value per sensor

# Fetch data from the source (e.g., API, database)
async def fetch_data() -> List[Dict[str, Any]]:
    """Fetch data from a source.
    
    Returns:
        List of fetched data dictionaries
    """
    # Simulated fetching
    return [{'sensor_id': 'sensor_1', 'value': '23.5'}, {'sensor_id': 'sensor_2', 'value': '20.8'}]

# Save aggregated metrics to the database
async def save_to_db(data: pl.DataFrame) -> None:
    """Save aggregated metrics to the database.
    
    Args:
        data: Aggregated metrics DataFrame
    """
    logger.info('Saving data to database...')  # Log saving
    # Simulated saving logic
    time.sleep(0.1)  # Simulate write time

# Call an external API (e.g., to notify or trigger an action)
async def call_api(data: Dict[str, Any]) -> None:
    """Call an external API with the processed data.
    
    Args:
        data: Data to send to the API
    """
    logger.info(f'Calling external API with data: {data}')  # Log API call
    # Simulated API call
    time.sleep(0.1)  # Simulate API call time

# Main orchestrator class
class SensorStreamProcessor:
    def __init__(self):
        self.config = Config()  # Load configuration

    async def process_stream(self) -> None:
        """Main method to process the sensor stream.
        
        Returns:
            None
        """
        raw_data = await fetch_data()  # Fetch raw data
        valid_data = []  # Store valid data
        for data in raw_data:
            try:
                await validate_input(data)  # Validate data
                sanitized_data = sanitize_fields(data)  # Sanitize data
                normalized_data = normalize_data(sanitized_data)  # Normalize data
                valid_data.append(normalized_data)  # Add to valid data
            except ValueError as e:
                logger.error(f'Validation error: {e}')  # Log validation errors

        await process_batch(valid_data)  # Process valid data
        aggregated = aggregate_metrics(pl.DataFrame(valid_data))  # Aggregate metrics
        await save_to_db(aggregated)  # Save aggregated data
        await call_api({'status': 'completed'})  # Notify completion

# Entry point
if __name__ == '__main__':
    processor = SensorStreamProcessor()  # Instantiate processor
    import asyncio
    asyncio.run(processor.process_stream())  # Run processing loop