data_quality.py

#!/usr/bin/env python3
"""
Data Quality Checks for Mini Data Warehouse
Validates data integrity and provides quality metrics
"""

import pandas as pd
import psycopg2
import logging
from datetime import datetime
import json
import os

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class DataQualityChecker:
    def __init__(self, connection_params=None):
        if connection_params is None:
            self.connection_params = {
                'host': 'localhost',
                'port': '5432',
                'database': 'warehouse',
                'user': 'admin',
                'password': 'secret'
            }
        else:
            self.connection_params = connection_params
        
        self.quality_results = {}
        
    def connect_db(self):
        """Connect to PostgreSQL database"""
        try:
            conn = psycopg2.connect(**self.connection_params)
            return conn
        except psycopg2.Error as e:
            logger.error(f"Database connection failed: {e}")
            return None
    
    def check_table_completeness(self):
        """Check for null values and completeness"""
        logger.info("Checking table completeness...")
        
        conn = self.connect_db()
        if not conn:
            return
            
        try:
            with conn.cursor() as cur:
                tables = ['customers', 'products', 'orders', 'order_items']
                completeness = {}
                
                for table in tables:
                    # Get total record count
                    cur.execute(f"SELECT COUNT(*) FROM {table}")
                    total_records = cur.fetchone()[0]
                    
                    # Get column info
                    cur.execute(f"""
                        SELECT column_name, data_type 
                        FROM information_schema.columns 
                        WHERE table_name = '{table}'
                        ORDER BY ordinal_position
                    """)
                    columns = cur.fetchall()
                    
                    table_stats = {
                        'total_records': total_records,
                        'columns': {}
                    }
                    
                    for col_name, data_type in columns:
                        # Count null values
                        cur.execute(f"SELECT COUNT(*) FROM {table} WHERE {col_name} IS NULL")
                        null_count = cur.fetchone()[0]
                        
                        completeness_rate = ((total_records - null_count) / total_records * 100) if total_records > 0 else 0
                        
                        table_stats['columns'][col_name] = {
                            'null_count': null_count,
                            'completeness_rate': round(completeness_rate, 2)
                        }
                    
                    completeness[table] = table_stats
                
                self.quality_results['completeness'] = completeness
                
        except psycopg2.Error as e:
            logger.error(f"Completeness check failed: {e}")
        finally:
            conn.close()
    
    def check_referential_integrity(self):
        """Check foreign key constraints"""
        logger.info("Checking referential integrity...")
        
        conn = self.connect_db()
        if not conn:
            return
            
        try:
            with conn.cursor() as cur:
                integrity_issues = {}
                
                # Check orders -> customers
                cur.execute("""
                    SELECT COUNT(*) FROM orders o 
                    LEFT JOIN customers c ON o.customer_id = c.customer_id 
                    WHERE c.customer_id IS NULL
                """)
                orphaned_orders = cur.fetchone()[0]
                
                # Check order_items -> orders
                cur.execute("""
                    SELECT COUNT(*) FROM order_items oi 
                    LEFT JOIN orders o ON oi.order_id = o.order_id 
                    WHERE o.order_id IS NULL
                """)
                orphaned_order_items_orders = cur.fetchone()[0]
                
                # Check order_items -> products
                cur.execute("""
                    SELECT COUNT(*) FROM order_items oi 
                    LEFT JOIN products p ON oi.product_id = p.product_id 
                    WHERE p.product_id IS NULL
                """)
                orphaned_order_items_products = cur.fetchone()[0]
                
                integrity_issues = {
                    'orphaned_orders': orphaned_orders,
                    'orphaned_order_items_orders': orphaned_order_items_orders,
                    'orphaned_order_items_products': orphaned_order_items_products
                }
                
                self.quality_results['referential_integrity'] = integrity_issues
                
        except psycopg2.Error as e:
            logger.error(f"Referential integrity check failed: {e}")
        finally:
            conn.close()
    
    def check_data_consistency(self):
        """Check for data consistency issues"""
        logger.info("Checking data consistency...")
        
        conn = self.connect_db()
        if not conn:
            return
            
        try:
            with conn.cursor() as cur:
                consistency_issues = {}
                
                # Check if order total matches sum of order items
                cur.execute("""
                    SELECT o.order_id, o.total_amount, 
                           COALESCE(SUM(oi.quantity * oi.unit_price), 0) as calculated_total
                    FROM orders o
                    LEFT JOIN order_items oi ON o.order_id = oi.order_id
                    GROUP BY o.order_id, o.total_amount
                    HAVING ABS(o.total_amount - COALESCE(SUM(oi.quantity * oi.unit_price), 0)) > 0.01
                """)
                
                inconsistent_totals = cur.fetchall()
                
                # Check for negative prices
                cur.execute("SELECT COUNT(*) FROM products WHERE price < 0")
                negative_prices = cur.fetchone()[0]
                
                # Check for negative quantities
                cur.execute("SELECT COUNT(*) FROM order_items WHERE quantity <= 0")
                negative_quantities = cur.fetchone()[0]
                
                # Check for duplicate emails
                cur.execute("""
                    SELECT email, COUNT(*) 
                    FROM customers 
                    GROUP BY email 
                    HAVING COUNT(*) > 1
                """)
                duplicate_emails = cur.fetchall()
                
                consistency_issues = {
                    'inconsistent_order_totals': len(inconsistent_totals),
                    'negative_prices': negative_prices,
                    'negative_quantities': negative_quantities,
                    'duplicate_emails': len(duplicate_emails)
                }
                
                self.quality_results['consistency'] = consistency_issues
                
        except psycopg2.Error as e:
            logger.error(f"Data consistency check failed: {e}")
        finally:
            conn.close()
    
    def generate_quality_report(self):
        """Generate comprehensive data quality report"""
        logger.info("Generating data quality report...")
        
        self.check_table_completeness()
        self.check_referential_integrity()
        self.check_data_consistency()
        
        # Calculate overall quality score
        total_score = 100
        
        # Deduct points for quality issues
        if 'referential_integrity' in self.quality_results:
            ri = self.quality_results['referential_integrity']
            if ri['orphaned_orders'] > 0:
                total_score -= 20
            if ri['orphaned_order_items_orders'] > 0:
                total_score -= 15
            if ri['orphaned_order_items_products'] > 0:
                total_score -= 15
        
        if 'consistency' in self.quality_results:
            c = self.quality_results['consistency']
            if c['inconsistent_order_totals'] > 0:
                total_score -= 20
            if c['negative_prices'] > 0:
                total_score -= 10
            if c['negative_quantities'] > 0:
                total_score -= 10
            if c['duplicate_emails'] > 0:
                total_score -= 10
        
        self.quality_results['overall_score'] = max(0, total_score)
        self.quality_results['timestamp'] = datetime.now().isoformat()
        
        return self.quality_results
    
    def save_report(self, filename='data_quality_report.json'):
        """Save quality report to file"""
        with open(filename, 'w') as f:
            json.dump(self.quality_results, f, indent=2)
        logger.info(f"Quality report saved to {filename}")
    
    def print_report(self):
        """Print formatted quality report"""
        report = self.quality_results
        
        print("\n" + "="*50)
        print("DATA QUALITY REPORT")
        print("="*50)
        print(f"Generated: {report.get('timestamp', 'Unknown')}")
        print(f"Overall Quality Score: {report.get('overall_score', 'N/A')}/100")
        print()
        
        # Completeness Report
        if 'completeness' in report:
            print("COMPLETENESS ANALYSIS")
            print("-" * 30)
            for table, stats in report['completeness'].items():
                print(f"{table.upper()}: {stats['total_records']} records")
                for col, col_stats in stats['columns'].items():
                    status = "✓" if col_stats['completeness_rate'] == 100 else "⚠"
                    print(f"  {status} {col}: {col_stats['completeness_rate']}% complete")
                print()
        
        # Referential Integrity Report
        if 'referential_integrity' in report:
            print("REFERENTIAL INTEGRITY")
            print("-" * 30)
            ri = report['referential_integrity']
            print(f"Orphaned orders: {ri['orphaned_orders']}")
            print(f"Orphaned order items (orders): {ri['orphaned_order_items_orders']}")
            print(f"Orphaned order items (products): {ri['orphaned_order_items_products']}")
            print()
        
        # Consistency Report
        if 'consistency' in report:
            print("DATA CONSISTENCY")
            print("-" * 30)
            c = report['consistency']
            print(f"Inconsistent order totals: {c['inconsistent_order_totals']}")
            print(f"Negative prices: {c['negative_prices']}")
            print(f"Negative quantities: {c['negative_quantities']}")
            print(f"Duplicate emails: {c['duplicate_emails']}")

def main():
    """Main function to run data quality checks"""
    checker = DataQualityChecker()
    report = checker.generate_quality_report()
    
    checker.print_report()
    checker.save_report()
    
    # Return exit code based on quality score
    score = report.get('overall_score', 0)
    if score >= 95:
        exit_code = 0
    elif score >= 80:
        exit_code = 1
    else:
        exit_code = 2
    
    logger.info(f"Data quality check completed with score: {score}/100")
    return exit_code

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