data_processing_engine.py

#!/usr/bin/env python3
"""
Data Processing Engine Module

A comprehensive data preprocessing and postprocessing framework that provides
automated discovery and application of data transformations.

Author: Data Processing Team
Version: 1.0.0
"""

from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Any, Union
from dataclasses import dataclass
from enum import Enum
import pandas as pd
from mostlyai.sdk import MostlyAI
from mostlyai.sdk.domain import Generator


@dataclass
class DataProcessorConfig:
    """Configuration object for data processors"""
    processor_type: str
    table_name: str
    column_name: str
    parameters: Dict[str, Any]
    priority: int = 0
    enabled: bool = True


class BaseDataProcessor(ABC):
    """
    Abstract base class for all data processors.
    
    All data processors must inherit from this class and implement
    the required abstract methods.
    """
    
    def __init__(self, name: str):
        """
        Initialize the data processor.
        
        Args:
            name (str): Name of the processor
        """
        self.name = name
    
    @abstractmethod
    def discover(self, table_name, column_name) -> Optional[DataProcessorConfig]:
        """
        Analyze the data's specified column to determine if preprocessing/postprocessing is needed.
        
        Args:
            table_name (str): Name of the table containing the column
            column_name (str): Name of the column to analyze. If left blank, then the processor will analyze all columns of the table.
            
        Returns:
            Optional[DataProcessorConfig]: List of Configuration objects if processing is needed,
                                         None if no processing is required
        """
        pass
    
    @abstractmethod
    def pre_process(self, config: DataProcessorConfig) -> None:
        """
        Apply preprocessing transformations to the data.
        After the preprocessing, it will update the target metadata with the changes done.
        It can create temporary data objects like lookup tables, etc. 
        
        Args:
            config (DataProcessorConfig): Configuration for the processor
        """
        pass
    
    @abstractmethod
    def post_process(self, config: DataProcessorConfig) -> None:
        """
        Apply postprocessing transformations to the data.
        It can create a copy of the initially created synthetic data as backup, then modify it. 
        
        Args:
            config (DataProcessorConfig): Configuration for the processor
        """
        pass


class RelativeDataProcessor(BaseDataProcessor):
    """
    Data processor for handling relative data transformations.
    
    This processor handles cases where data values are relative to other
    values in the dataset, such as percentage changes, ratios, or deltas.
    """
    
    def __init__(self):
        super().__init__("RelativeDataProcessor")
    
    def discover(self, table_name, column_name) -> Optional[DataProcessorConfig]:
        pass
    
    def pre_process(self, config: DataProcessorConfig) -> None:
        pass
    
    def post_process(self, config: DataProcessorConfig) -> None:
        pass


class CompositeKeyDataProcessor(BaseDataProcessor):
    """
    Data processor for handling composite key transformations.
    
    This processor manages complex key relationships and composite identifiers
    that may need special handling during data processing.
    """
    
    def __init__(self, pk_table_name, pk_column_names, fk_list):
        super().__init__("CompositeKeyDataProcessor")
    
    def discover(self, table_name, column_name) -> Optional[DataProcessorConfig]:
        """
        This will not have a discovery logic as it is a composite key processor.
        """
        # TODO: Implement discovery logic
        # - Check if column is part of primary key
        # - Analyze for composite key patterns
        # - Detect foreign key relationships
        # - Identify uniqueness constraints
        return None
    
    def pre_process(self, config: DataProcessorConfig) -> None:
        pass
    
    def post_process(self, config: DataProcessorConfig) -> None:
        pass


class FormulaDataProcessor(BaseDataProcessor):
    """
    Data processor for handling formula-based transformations.
    
    This processor manages columns that are derived from mathematical formulas
    or expressions involving other columns.
    """
    
    def __init__(self):
        super().__init__("FormulaDataProcessor")
    
    def discover(self, table_name, column_name) -> Optional[DataProcessorConfig]:
        pass
    
    def pre_process(self, config: DataProcessorConfig) -> None:
        pass
    
    def post_process(self, config: DataProcessorConfig) -> None:
        pass


class ProcessingEngine:
    """
    Main processing engine that orchestrates data preprocessing and postprocessing.
    
    This class manages the entire data processing pipeline, including discovery,
    configuration, and execution of data transformations.
    """
    
    def __init__(self, mostly: MostlyAI, generator_config: GeneratorConfig ):
        """
        Initialize the processing engine.
        Alternatively, we could also pass a generator config object to create the generator on the fly. 

        """
        self.source_config =  generator_config
        self.target_config = generator_config
        self.mostly = mostly
        self.data_processor_configs: List[DataProcessorConfig] = []
        self.data: Dict[str, pd.DataFrame] = {}  # Central data object
        self.source_generator = None
        self.target_generator = None
        # Static collection of available data processors
        self.available_processors: List[BaseDataProcessor] = [
            RelativeDataProcessor(),
            CompositeKeyDataProcessor(),
            FormulaDataProcessor()
        ]
    
    def add_data_processor_step(self, config: DataProcessorConfig) -> None:
        """
        Add a data processor configuration to the processing pipeline.
        
        Args:
            config (DataProcessorConfig): Configuration object to add
        """
        # TODO: Implement step addition logic
        # - Validate configuration object
        # - Check for conflicts with existing steps
        # - Add to the processing pipeline
        # - Sort by priority if needed
        pass
    
    def auto_discover(self) -> None:
        """
        Automatically discover data processing needs across all tables and columns.
        
        This method traverses all tables and columns, calls the discover function
        of all registered DataProcessor classes, and adds configurations for
        any processing that is needed.
        """
        # TODO: Implement auto-discovery logic
        # - Iterate through all tables in source metadata
        # - For each table, iterate through all columns
        # - Call discover() method on each available processor
        # - Add non-null configurations to the processing pipeline
        # - Handle discovery conflicts and priorities
        pass
    
    def load_data(self) -> None:
        """
        Load data from the source generator configuration.
        Actually, it would be good, if the generator object had a method to load the data to a staging area. 
        It should work like the pull_data method, but without the joins. 
        If not possible, then we can implement one using the connector objects etc. 

        """
        # TODO: Implement data loading logic
        self.source_generator = self.mostly.generators.create(self.source_config)
        self.source_generator.pull_data(self.source_config)  # method not available in the SDK
        pass
    
    def pre_process(self) -> None:
        """
        Execute all preprocessing steps in the configured order.
        
        This method applies all preprocessing transformations to the central data object
        and updates the target metadata accordingly.
        """
        # TODO: Implement preprocessing logic
        # - Sort configurations by priority and stage
        # - Iterate through all PRE_PROCESS configurations
        # - Call the appropriate processor's pre_process method
        # - Update target metadata as transformations are applied
        # - Handle processing errors and rollback if needed
        pass
    
    def post_process(self) -> None:
        """
        Execute all postprocessing steps in the configured order.
        
        This method applies all postprocessing transformations to the synthetic data.
        """
        # TODO: Implement postprocessing logic
        # - Unzip the synthetic data to the staging area
        # - Sort configurations by priority and stage
        # - Iterate through all POST_PROCESS configurations
        # - Call the appropriate processor's post_process method
        pass
    
    def train(self) -> None:
        """
        Train the generator.
        """
        # change the config to the new config
        self.target_generator = self.mostly.generators.create(self.target_config)
        self.target_generator.training.start()
        self.target_generator.training.wait()

    def generate(self) -> None:
        """
        Generate data.
        """
        # compose the synthetic dataset config
        synthetic_dataset_config = "some config"
        mostly.synthetic_datasets.create(synthetic_dataset_config)
        mostly.synthetic_datasets.generation.start()
        mostly.synthetic_datasets.generation.wait()

    def export_data(self) -> None:
        """
        Upload data to the destination defined in the generator configuration.
        """
        pass
    
# Example usage and testing functions
def example_usage():
    """
    Example demonstrating how to use the ProcessingEngine.
    """

    mostly = MostlyAI()

    someGeneratorConfig = {
        'name': 'Processing Engine Example Generator',
        'tables': [
            {
                'name': 'customers',
                'primary_key': 'CustomerId',
                'columns': [
                    {
                        'name': 'CustomerId',
                        'model_encoding_type': 'ModelEncodingType.tabular_categorical'
                    },
                    {
                        'name': 'CustomerType',
                        'model_encoding_type': 'ModelEncodingType.tabular_categorical'
                    },
                    {
                        'name': 'CustomerName',
                        'model_encoding_type': 'ModelEncodingType.language_text'
                    },
                    {
                        'name': 'JoinedAt',
                        'model_encoding_type': 'ModelEncodingType.tabular_datetime'
                    }
                ],
                'foreign_keys': None
            },
            {
                'name': 'orders',
                'primary_key': None,
                'columns': [
                    {
                        'name': 'CustomerId',
                        'model_encoding_type': 'ModelEncodingType.tabular_categorical'
                    },
                    {
                        'name': 'OrderDate',
                        'model_encoding_type': 'ModelEncodingType.tabular_datetime'
                    },
                    {
                        'name': 'OrderStatus',
                        'model_encoding_type': 'ModelEncodingType.tabular_categorical'
                    },
                    {
                        'name': 'Amount',
                        'model_encoding_type': 'ModelEncodingType.tabular_numeric_auto'
                    }
                ],
                'foreign_keys': [
                    {
                        'column': 'CustomerId',
                        'referenced_table': 'customers',
                        'is_context': True
                    }
                ]
            },
            {
                'name': 'warehouses',
                'primary_key': 'WarehouseId',
                'columns': [
                    {
                        'name': 'WarehouseId',
                        'model_encoding_type': 'ModelEncodingType.tabular_categorical'
                    },
                    {
                        'name': 'City',
                        'model_encoding_type': 'ModelEncodingType.tabular_categorical'
                    },
                    {
                        'name': 'Capacity',
                        'model_encoding_type': 'ModelEncodingType.tabular_categorical'
                    }
                ],
                'foreign_keys': None
            },
            {
                'name': 'wh_transactions',
                'primary_key': 'WarehouseTxId',
                'columns': [
                    {
                        'name': 'WarehouseTxId',
                        'model_encoding_type': 'ModelEncodingType.tabular_numeric_auto'
                    },
                    {
                        'name': 'WarehouseId',
                        'model_encoding_type': 'ModelEncodingType.tabular_categorical'
                    },
                    {
                        'name': 'CustomerId',
                        'model_encoding_type': 'ModelEncodingType.tabular_categorical'
                    },
                    {
                        'name': 'OrderDate',
                        'model_encoding_type': 'ModelEncodingType.tabular_datetime'
                    },
                    {
                        'name': 'Amount',
                        'model_encoding_type': 'ModelEncodingType.tabular_numeric_auto'
                    }
                ],
                'foreign_keys': [
                    {
                        'column': 'WarehouseId',
                        'referenced_table': 'warehouses',
                        'is_context': True
                    }
                ]
            }
        ]
    }

    g = mostly.train(config=someGeneratorConfig)
    s = mostly.generate(g, size=1000)
    s.data()


    # Initialize the processing engine
    processor = ProcessingEngine(mostly, "some generator config")
    
    # Load data
    processor.load_data()

    # Auto-discover processing steps for Relative, Formula. Composite key will be added manually.
    processor.auto_discover()

    # Add a data processor step for composite key processing
    processor.add_data_processor_step( 
        CompositeKeyDataProcessor( 
            pk_table_name = "orders", 
            pk_column_names = ["CustomerId", "OrderDate"], 
            pk_fallback_method = "REJECT",
            fk_list = [ {
                "fk_table_name": "wh_transactions",
                "fk_column_names": ["CustomerId", "OrderDate"]
            }
            # {
            #     "fk_table_name": "payment_tx",
            #     "fk_column_names": ["PAYMENT_CUSTOMER_ID", "ORDER_DATE"]
            # }
            ]
        ) 
        )

    # Execute preprocessing
    processor.pre_process()

    # train the generator
    processor.() 

    # generate data
    processor.generate()
    
    # Execute postprocessing
    processor.post_process()
    
    # finally export the data to the destination defined in the original generator configuration
    processor.export_data()
    
    

if __name__ == "__main__":
    example_usage()