Backfill Factory Sensor History into Delta Lake with Apache Spark and delta-rs

"""
Production implementation for backfilling factory sensor history into Delta Lake.
Provides secure, scalable operations using Apache Spark and delta-rs.
"""
from typing import Dict, Any, List, Tuple
import os
import logging
import time
import delta
from pyspark.sql import SparkSession, DataFrame

# Setting up logging for the application
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class Config:
    # Configuration class for environment variables
    delta_table_path: str = os.getenv('DELTA_TABLE_PATH')
    spark_master: str = os.getenv('SPARK_MASTER', 'local[*]')

def create_spark_session() -> SparkSession:
    """
    Create a Spark session configured for Delta Lake.
    
    Returns:
        SparkSession: Configured Spark session
    """
    logger.info('Creating Spark session...')
    spark = SparkSession.builder \
        .appName('Backfill Sensor History') \
        .master(Config.spark_master) \
        .config('spark.sql.extensions', 'delta.sql.DeltaSparkSessionExtensions') \
        .config('spark.sql.catalog.spark_catalog', 'org.apache.spark.sql.delta.catalog.DeltaCatalog') \
        .getOrCreate()
    logger.info('Spark session created successfully.')
    return spark

async def validate_input(data: Dict[str, Any]) -> bool:
    """Validate the input data for processing.
    
    Args:
        data: Input data dictionary to validate
    Returns:
        bool: True if valid
    Raises:
        ValueError: If validation fails
    """
    if 'timestamp' not in data or 'sensor_id' not in data:
        raise ValueError('Missing required fields: timestamp and sensor_id')
    logger.debug('Input data validated successfully.')
    return True

async def sanitize_fields(data: Dict[str, Any]) -> Dict[str, Any]:
    """Sanitize input fields to prevent SQL injection and other issues.
    
    Args:
        data: Raw input data
    Returns:
        Dict[str, Any]: Sanitized data
    """
    logger.info('Sanitizing input fields...')
    sanitized_data = {key: str(value).strip() for key, value in data.items()}
    logger.debug('Input fields sanitized.')
    return sanitized_data

async def transform_records(data: List[Dict[str, Any]]) -> DataFrame:
    """Transform raw sensor data into a Spark DataFrame.
    
    Args:
        data: List of dictionaries containing sensor records
    Returns:
        DataFrame: Spark DataFrame ready for processing
    """
    logger.info('Transforming records into DataFrame...')
    spark = create_spark_session()
    df = spark.createDataFrame(data)
    logger.debug('Records transformed into DataFrame.')
    return df

async def save_to_delta(df: DataFrame) -> None:
    """Save the processed DataFrame to Delta Lake.
    
    Args:
        df: Spark DataFrame to save
    """
    logger.info('Saving DataFrame to Delta Lake...')
    df.write.format('delta') \
        .mode('append') \
        .save(Config.delta_table_path)
    logger.info('DataFrame saved successfully to Delta Lake.')

async def process_batch(batch: List[Dict[str, Any]]) -> None:
    """Process a batch of sensor data.
    
    Args:
        batch: List of sensor data records
    """
    try:
        logger.info(f'Processing batch of {len(batch)} records...')
        for record in batch:
            await validate_input(record)  # Validate each record
            sanitized_record = await sanitize_fields(record)  # Sanitize fields
            logger.debug(f'Sanitized record: {sanitized_record}')
        df = await transform_records(batch)  # Transform into DataFrame
        await save_to_delta(df)  # Save to Delta Lake
    except ValueError as e:
        logger.error(f'Validation error: {e}')
    except Exception as e:
        logger.error(f'Error processing batch: {e}')

async def fetch_data() -> List[Dict[str, Any]]:
    """Fetch data from an external source (placeholder).
    
    Returns:
        List[Dict[str, Any]]: Fetched sensor data
    """
    logger.info('Fetching data...')
    # Simulated fetching data
    data = [
        {'timestamp': '2023-10-01T00:00:00Z', 'sensor_id': 'sensor_1', 'value': 10},
        {'timestamp': '2023-10-01T00:01:00Z', 'sensor_id': 'sensor_2', 'value': 15}
    ]
    logger.debug('Data fetched successfully.')
    return data

async def main() -> None:
    """Main function to orchestrate the backfill process.
    """
    try:
        logger.info('Starting the backfill process...')
        data = await fetch_data()  # Fetch data
        batch_size = 100  # Define the batch size
        for i in range(0, len(data), batch_size):
            batch = data[i:i + batch_size]  # Create batch
            await process_batch(batch)  # Process each batch
        logger.info('Backfill process completed successfully.')
    except Exception as e:
        logger.error(f'Error in main process: {e}')

if __name__ == '__main__':
    # Entry point for the script
    import asyncio
    asyncio.run(main())