dataset_processor.py

#!/usr/bin/env python3
"""
Centralized dataset processing module for both local and server use.
"""

import os
import json
import pandas as pd
from io import BytesIO
from typing import Dict, List, Any, Tuple
from utils import (
    get_drive_service,
    download_file_from_drive,
    extract_metadata_from_dataframe,
    suggest_dq_rules,
    create_contract_excel,
    generate_dq_report
)

def create_dataset_readme(metadata: Dict[str, Any], dq_rules: List[Dict[str, Any]], 
                         dq_report: Dict[str, Any], readme_path: str):
    """Create a comprehensive README file for the processed dataset."""
    content = f"""# Dataset Processing Report

## Dataset Information
- **Filename**: {metadata['filename']}
- **Rows**: {metadata['row_count']:,}
- **Columns**: {metadata['column_count']}
- **Processing Date**: {pd.Timestamp.now().strftime('%Y-%m-%d %H:%M:%S')}

## Column Details
| Column Name | Data Type | Null Count | Null % | Unique Count | Min Value | Max Value | Mean |
|-------------|-----------|------------|--------|--------------|-----------|-----------|------|
"""
    
    for col in metadata['columns']:
        min_val = f"{col.get('min_value', 'N/A'):.4f}" if col.get('min_value') is not None else "N/A"
        max_val = f"{col.get('max_value', 'N/A'):.4f}" if col.get('max_value') is not None else "N/A"
        mean_val = f"{col.get('mean_value', 'N/A'):.4f}" if col.get('mean_value') is not None else "N/A"
        
        content += f"| {col['name']} | {col['data_type']} | {col['null_count']} | {col['null_percentage']:.1f}% | {col['unique_count']} | {min_val} | {max_val} | {mean_val} |\n"
    
    content += f"""
## Data Quality Summary
- **Total Rules**: {dq_report['data_quality_summary']['total_rules']}
- **Error Rules**: {dq_report['data_quality_summary']['error_rules']}
- **Warning Rules**: {dq_report['data_quality_summary']['warning_rules']}

## Data Quality Rules
"""
    
    for rule in dq_rules:
        severity_icon = "🔴" if rule['severity'] == 'error' else "🟡"
        content += f"- {severity_icon} **{rule['rule_type'].upper()}**: {rule['description']}\n"
    
    content += f"""
## Column Quality Metrics
"""
    
    for col_quality in dq_report['column_quality']:
        content += f"- **{col_quality['column_name']}**: {col_quality['completeness']:.1f}% complete, {col_quality['uniqueness']:.1f}% unique\n"
    
    content += f"""
## Files Generated
- `{metadata['filename']}` - Original dataset
- `{metadata['filename'].split('.')[0]}_metadata.json` - Detailed metadata
- `{metadata['filename'].split('.')[0]}_contract.xlsx` - Data contract with schema and DQ rules
- `{metadata['filename'].split('.')[0]}_dq_report.json` - Comprehensive data quality report
- `README.md` - This summary document

## Usage
This dataset has been processed through the MCP Dataset Onboarding pipeline. All artifacts are ready for catalog publication or further analysis.
"""
    
    with open(readme_path, 'w', encoding='utf-8') as f:
        f.write(content)

def process_dataset_with_organization(file_id: str, output_folder: str = "processed_datasets") -> Dict[str, Any]:
    """
    Process dataset and organize all artifacts in a dedicated folder structure.
    
    Args:
        file_id: Google Drive file ID
        output_folder: Base folder for organizing processed datasets
        
    Returns:
        Dictionary with processing results and file paths
    """
    try:
        print(f"Processing file ID: {file_id}")
        
        # Step 1: Download and extract metadata
        print("Step 1: Downloading file and extracting metadata...")
        drive_service = get_drive_service()
        
        # Download file
        file_content = download_file_from_drive(file_id, drive_service)
        
        # Get file info
        file_info = drive_service.files().get(fileId=file_id).execute()
        filename = file_info['name']
        print(f"Downloaded: {filename}")
        
        # Create output directory structure
        base_name = filename.split('.')[0]
        dataset_folder = os.path.join(output_folder, base_name)
        os.makedirs(dataset_folder, exist_ok=True)
        print(f"Created output folder: {dataset_folder}")
        
        # Read file based on extension
        if filename.lower().endswith('.csv'):
            df = pd.read_csv(BytesIO(file_content))
        elif filename.lower().endswith(('.xlsx', '.xls')):
            df = pd.read_excel(BytesIO(file_content))
        else:
            raise Exception("Unsupported file format")
        
        # Extract metadata
        metadata = extract_metadata_from_dataframe(df, filename)
        print(f"[OK] Extracted metadata: {metadata['row_count']} rows, {metadata['column_count']} columns")
        
        # Step 2: Generate DQ rules
        print("Step 2: Generating data quality rules...")
        dq_rules = suggest_dq_rules(metadata)
        print(f"[OK] Generated {len(dq_rules)} DQ rules")
        
        # Step 3: Create contract Excel
        print("Step 3: Creating contract Excel file...")
        contract_filename = os.path.join(dataset_folder, f"{base_name}_contract.xlsx")
        create_contract_excel(metadata, dq_rules, contract_filename)
        print(f"[OK] Created contract: {contract_filename}")
        
        # Step 4: Generate DQ report
        print("Step 4: Generating DQ report...")
        dq_report = generate_dq_report(metadata, dq_rules)
        
        # Save artifacts in organized folder
        metadata_filename = os.path.join(dataset_folder, f"{base_name}_metadata.json")
        dq_report_filename = os.path.join(dataset_folder, f"{base_name}_dq_report.json")
        
        # Save original dataset for reference
        original_filename = os.path.join(dataset_folder, filename)
        with open(original_filename, 'wb') as f:
            f.write(file_content)
        
        with open(metadata_filename, 'w') as f:
            json.dump(metadata, f, indent=2)
        
        with open(dq_report_filename, 'w') as f:
            json.dump(dq_report, f, indent=2)
        
        # Create a summary README for the dataset
        readme_filename = os.path.join(dataset_folder, "README.md")
        create_dataset_readme(metadata, dq_rules, dq_report, readme_filename)
        
        print(f"[OK] Saved artifacts in folder: {dataset_folder}")
        print(f"  - {os.path.basename(original_filename)} (original dataset)")
        print(f"  - {os.path.basename(metadata_filename)} (metadata)")
        print(f"  - {os.path.basename(contract_filename)} (contract)")
        print(f"  - {os.path.basename(dq_report_filename)} (DQ report)")
        print(f"  - {os.path.basename(readme_filename)} (summary)")
        
        return {
            "status": "success",
            "output_folder": dataset_folder,
            "files_created": [
                original_filename,
                metadata_filename, 
                contract_filename, 
                dq_report_filename,
                readme_filename
            ],
            "metadata": metadata,
            "dq_rules": dq_rules,
            "dq_report": dq_report
        }
        
    except Exception as e:
        print(f"Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return {"status": "error", "message": str(e)}

def list_processed_datasets(output_folder: str = "processed_datasets") -> List[Dict[str, Any]]:
    """List all processed datasets in the output folder."""
    datasets = []
    
    if not os.path.exists(output_folder):
        return datasets
    
    for item in os.listdir(output_folder):
        item_path = os.path.join(output_folder, item)
        if os.path.isdir(item_path):
            # Check if it's a valid dataset folder
            readme_path = os.path.join(item_path, "README.md")
            metadata_path = os.path.join(item_path, f"{item}_metadata.json")
            
            if os.path.exists(readme_path) and os.path.exists(metadata_path):
                try:
                    with open(metadata_path, 'r') as f:
                        metadata = json.load(f)
                    
                    datasets.append({
                        "dataset_name": item,
                        "folder_path": item_path,
                        "filename": metadata.get('filename', 'Unknown'),
                        "row_count": metadata.get('row_count', 0),
                        "column_count": metadata.get('column_count', 0),
                        "processed_date": os.path.getctime(item_path)
                    })
                except Exception as e:
                    print(f"Error reading metadata for {item}: {e}")
    
    return sorted(datasets, key=lambda x: x['processed_date'], reverse=True)