ucs_data_pipeline.py

#!/usr/bin/env python3
"""
UCS Data Pipeline - Enhanced version

A modular data pipeline for extracting data from Superset and loading it into PostgreSQL.
"""

import argparse
import sys
import os
from typing import Dict, Any

# Load environment variables from .env file
try:
    from dotenv import load_dotenv
    load_dotenv()
    print(f"Loaded .env file. SUPERSET_API_KEY present: {'SUPERSET_API_KEY' in os.environ}")
except ImportError:
    print("Warning: python-dotenv not installed. Environment variables may not be loaded properly.")

from ucs_pipeline.pipeline import Pipeline
from ucs_pipeline.utils.error_handling.exceptions import PipelineError

def parse_args() -> argparse.Namespace:
    """Parse command line arguments.
    
    Returns:
        Parsed arguments
    """
    parser = argparse.ArgumentParser(description='UCS Data Pipeline')
    parser.add_argument(
        '-c', '--config',
        help='Path to configuration file',
        type=str,
        default=None
    )
    parser.add_argument(
        '-v', '--verbose',
        help='Enable verbose logging',
        action='store_true'
    )
    
    return parser.parse_args()

def main() -> int:
    """Main entry point for the UCS Data Pipeline.
    
    Returns:
        Exit code (0 for success, non-zero for failure)
    """
    args = parse_args()
    
    try:
        # Initialize and run the pipeline
        pipeline = Pipeline(config_file=args.config)
        
        # Set log level to DEBUG if verbose flag is set
        if args.verbose:
            pipeline.logger.logger.setLevel('DEBUG')
        
        # Run the pipeline
        metrics = pipeline.run()
        
        # Clean up resources
        pipeline.cleanup()
        
        # Print execution summary
        print("\nExecution Summary:")
        print(f"- Execution time: {metrics['execution_time']:.2f} seconds")
        print(f"- Records extracted: {metrics['records_extracted']}")
        print(f"- Records loaded: {metrics.get('records_loaded', 0)}")
        
        if metrics.get('errors', 0) > 0:
            print(f"- Errors: {metrics['errors']}")
            print(f"- Error message: {metrics.get('error_message', 'Unknown error')}")
            return 1
        
        return 0
        
    except PipelineError as e:
        print(f"Pipeline error: {str(e)}", file=sys.stderr)
        return 1
    except Exception as e:
        print(f"Unexpected error: {str(e)}", file=sys.stderr)
        return 1

if __name__ == "__main__":
    sys.exit(main())