retention_policy.py

#!/usr/bin/env python3
"""
Data Retention Policy Implementation for Mini Data Warehouse
Manages data lifecycle and automated cleanup
"""

import psycopg2
from datetime import datetime, timedelta
import logging
import json
import os

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

class RetentionPolicyManager:
    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
        
        # Default retention policies (in days)
        self.retention_policies = {
            'raw_data': {
                'orders': 1095,  # 3 years
                'order_items': 1095,  # 3 years
                'customers': -1,  # Never delete (keep forever)
                'products': -1,  # Never delete (keep forever)
            },
            'aggregated_data': {
                'fact_sales': 2555,  # 7 years
                'dim_customer': -1,  # Never delete
                'dim_product': -1,  # Never delete
                'dim_date': -1,  # Never delete
            },
            'logs_and_metadata': {
                'data_quality_logs': 90,  # 90 days
                'load_metadata': 365,  # 1 year
            }
        }
        
    def connect_db(self):
        """Connect to PostgreSQL database"""
        try:
            conn = psycopg2.connect(**self.connection_params)
            conn.autocommit = True
            return conn
        except psycopg2.Error as e:
            logger.error(f"Database connection failed: {e}")
            return None
    
    def create_audit_table(self):
        """Create audit table to track retention policy executions"""
        conn = self.connect_db()
        if not conn:
            return False
        
        try:
            with conn.cursor() as cur:
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS retention_audit (
                        audit_id SERIAL PRIMARY KEY,
                        table_name VARCHAR(100) NOT NULL,
                        policy_type VARCHAR(50) NOT NULL,
                        retention_days INTEGER,
                        records_before INTEGER,
                        records_deleted INTEGER,
                        execution_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        status VARCHAR(20) DEFAULT 'SUCCESS'
                    )
                """)
                
                # Create index for better query performance
                cur.execute("""
                    CREATE INDEX IF NOT EXISTS idx_retention_audit_table_date 
                    ON retention_audit(table_name, execution_date)
                """)
                
                logger.info("Retention audit table created/verified")
                return True
                
        except psycopg2.Error as e:
            logger.error(f"Error creating audit table: {e}")
            return False
        finally:
            conn.close()
    
    def create_archive_tables(self):
        """Create archive tables for storing deleted data"""
        conn = self.connect_db()
        if not conn:
            return False
        
        try:
            with conn.cursor() as cur:
                # Archive table for orders
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS archived_orders (
                        archived_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        order_id INTEGER,
                        customer_id INTEGER,
                        order_date DATE,
                        total_amount NUMERIC(12,2)
                    )
                """)
                
                # Archive table for order_items
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS archived_order_items (
                        archived_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        order_item_id INTEGER,
                        order_id INTEGER,
                        product_id INTEGER,
                        quantity INTEGER,
                        unit_price NUMERIC(10,2)
                    )
                """)
                
                # Archive table for fact_sales
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS archived_fact_sales (
                        archived_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        sales_key INTEGER,
                        date_key INTEGER,
                        customer_key INTEGER,
                        product_key INTEGER,
                        order_id INTEGER,
                        order_item_id INTEGER,
                        quantity INTEGER,
                        unit_price NUMERIC(10,2),
                        total_amount NUMERIC(12,2),
                        profit_margin NUMERIC(5,2),
                        created_date TIMESTAMP
                    )
                """)
                
                logger.info("Archive tables created/verified")
                return True
                
        except psycopg2.Error as e:
            logger.error(f"Error creating archive tables: {e}")
            return False
        finally:
            conn.close()
    
    def get_table_record_count(self, table_name):
        """Get current record count for a table"""
        conn = self.connect_db()
        if not conn:
            return 0
        
        try:
            with conn.cursor() as cur:
                cur.execute(f"SELECT COUNT(*) FROM {table_name}")
                return cur.fetchone()[0]
        except psycopg2.Error as e:
            logger.error(f"Error counting records in {table_name}: {e}")
            return 0
        finally:
            conn.close()
    
    def archive_and_delete_orders(self, retention_days):
        """Archive and delete old orders and related order items"""
        if retention_days <= 0:
            logger.info("Orders retention policy: Keep forever")
            return True
        
        cutoff_date = datetime.now() - timedelta(days=retention_days)
        logger.info(f"Applying retention policy to orders: deleting records older than {cutoff_date.date()}")
        
        conn = self.connect_db()
        if not conn:
            return False
        
        try:
            with conn.cursor() as cur:
                # Get count before deletion
                records_before = self.get_table_record_count('orders')
                
                # Archive orders to be deleted
                cur.execute("""
                    INSERT INTO archived_orders (order_id, customer_id, order_date, total_amount)
                    SELECT order_id, customer_id, order_date, total_amount
                    FROM orders
                    WHERE order_date < %s
                """, (cutoff_date.date(),))
                
                archived_orders = cur.rowcount
                
                # Archive related order items
                cur.execute("""
                    INSERT INTO archived_order_items (order_item_id, order_id, product_id, quantity, unit_price)
                    SELECT oi.order_item_id, oi.order_id, oi.product_id, oi.quantity, oi.unit_price
                    FROM order_items oi
                    JOIN orders o ON oi.order_id = o.order_id
                    WHERE o.order_date < %s
                """, (cutoff_date.date(),))
                
                archived_order_items = cur.rowcount
                
                # Delete order items first (foreign key constraint)
                cur.execute("""
                    DELETE FROM order_items
                    WHERE order_id IN (
                        SELECT order_id FROM orders WHERE order_date < %s
                    )
                """, (cutoff_date.date(),))
                
                deleted_order_items = cur.rowcount
                
                # Delete orders
                cur.execute("DELETE FROM orders WHERE order_date < %s", (cutoff_date.date(),))
                deleted_orders = cur.rowcount
                
                # Record audit information
                cur.execute("""
                    INSERT INTO retention_audit (table_name, policy_type, retention_days, records_before, records_deleted)
                    VALUES (%s, %s, %s, %s, %s)
                """, ('orders', 'time_based', retention_days, records_before, deleted_orders))
                
                cur.execute("""
                    INSERT INTO retention_audit (table_name, policy_type, retention_days, records_before, records_deleted)
                    VALUES (%s, %s, %s, %s, %s)
                """, ('order_items', 'cascading', retention_days, 0, deleted_order_items))
                
                logger.info(f"Archived and deleted {deleted_orders} orders and {deleted_order_items} order items")
                return True
                
        except psycopg2.Error as e:
            logger.error(f"Error in orders retention policy: {e}")
            return False
        finally:
            conn.close()
    
    def archive_and_delete_fact_sales(self, retention_days):
        """Archive and delete old sales facts"""
        if retention_days <= 0:
            logger.info("Fact sales retention policy: Keep forever")
            return True
        
        cutoff_date = datetime.now() - timedelta(days=retention_days)
        cutoff_date_key = int(cutoff_date.strftime('%Y%m%d'))
        
        logger.info(f"Applying retention policy to fact_sales: deleting records older than {cutoff_date.date()}")
        
        conn = self.connect_db()
        if not conn:
            return False
        
        try:
            with conn.cursor() as cur:
                # Get count before deletion
                records_before = self.get_table_record_count('fact_sales')
                
                # Archive sales facts to be deleted
                cur.execute("""
                    INSERT INTO archived_fact_sales (
                        sales_key, date_key, customer_key, product_key, order_id, 
                        order_item_id, quantity, unit_price, total_amount, profit_margin, created_date
                    )
                    SELECT 
                        sales_key, date_key, customer_key, product_key, order_id,
                        order_item_id, quantity, unit_price, total_amount, profit_margin, created_date
                    FROM fact_sales
                    WHERE date_key < %s
                """, (cutoff_date_key,))
                
                archived_records = cur.rowcount
                
                # Delete old sales facts
                cur.execute("DELETE FROM fact_sales WHERE date_key < %s", (cutoff_date_key,))
                deleted_records = cur.rowcount
                
                # Record audit information
                cur.execute("""
                    INSERT INTO retention_audit (table_name, policy_type, retention_days, records_before, records_deleted)
                    VALUES (%s, %s, %s, %s, %s)
                """, ('fact_sales', 'time_based', retention_days, records_before, deleted_records))
                
                logger.info(f"Archived and deleted {deleted_records} sales fact records")
                return True
                
        except psycopg2.Error as e:
            logger.error(f"Error in fact_sales retention policy: {e}")
            return False
        finally:
            conn.close()
    
    def cleanup_audit_logs(self, retention_days=90):
        """Clean up old audit log entries"""
        cutoff_date = datetime.now() - timedelta(days=retention_days)
        logger.info(f"Cleaning up audit logs older than {cutoff_date.date()}")
        
        conn = self.connect_db()
        if not conn:
            return False
        
        try:
            with conn.cursor() as cur:
                cur.execute(
                    "DELETE FROM retention_audit WHERE execution_date < %s", 
                    (cutoff_date,)
                )
                
                deleted_records = cur.rowcount
                logger.info(f"Deleted {deleted_records} old audit log entries")
                return True
                
        except psycopg2.Error as e:
            logger.error(f"Error cleaning up audit logs: {e}")
            return False
        finally:
            conn.close()
    
    def get_retention_statistics(self):
        """Get statistics about retention policy executions"""
        conn = self.connect_db()
        if not conn:
            return {}
        
        try:
            with conn.cursor() as cur:
                # Get recent retention policy executions
                cur.execute("""
                    SELECT 
                        table_name,
                        MAX(execution_date) as last_execution,
                        SUM(records_deleted) as total_deleted,
                        COUNT(*) as execution_count
                    FROM retention_audit
                    WHERE execution_date >= CURRENT_DATE - INTERVAL '30 days'
                    GROUP BY table_name
                    ORDER BY last_execution DESC
                """)
                
                recent_executions = cur.fetchall()
                
                # Get archive table sizes
                cur.execute("SELECT COUNT(*) FROM archived_orders")
                archived_orders_count = cur.fetchone()[0]
                
                cur.execute("SELECT COUNT(*) FROM archived_order_items")
                archived_order_items_count = cur.fetchone()[0]
                
                cur.execute("SELECT COUNT(*) FROM archived_fact_sales")
                archived_fact_sales_count = cur.fetchone()[0]
                
                statistics = {
                    'recent_executions': [
                        {
                            'table_name': row[0],
                            'last_execution': row[1].isoformat() if row[1] else None,
                            'total_deleted': row[2],
                            'execution_count': row[3]
                        } for row in recent_executions
                    ],
                    'archive_counts': {
                        'archived_orders': archived_orders_count,
                        'archived_order_items': archived_order_items_count,
                        'archived_fact_sales': archived_fact_sales_count
                    }
                }
                
                return statistics
                
        except psycopg2.Error as e:
            logger.error(f"Error getting retention statistics: {e}")
            return {}
        finally:
            conn.close()
    
    def execute_retention_policies(self):
        """Execute all configured retention policies"""
        logger.info("Starting retention policy execution...")
        
        # Ensure infrastructure is in place
        if not self.create_audit_table():
            return False
        
        if not self.create_archive_tables():
            return False
        
        success = True
        
        # Execute raw data retention policies
        raw_policies = self.retention_policies['raw_data']
        
        if not self.archive_and_delete_orders(raw_policies['orders']):
            success = False
        
        # Execute aggregated data retention policies
        agg_policies = self.retention_policies['aggregated_data']
        
        if not self.archive_and_delete_fact_sales(agg_policies['fact_sales']):
            success = False
        
        # Cleanup audit logs
        if not self.cleanup_audit_logs():
            success = False
        
        if success:
            logger.info("All retention policies executed successfully")
        else:
            logger.error("Some retention policies failed")
        
        return success
    
    def load_retention_config(self, config_file='retention_config.json'):
        """Load retention policies from configuration file"""
        if os.path.exists(config_file):
            with open(config_file, 'r') as f:
                config = json.load(f)
                self.retention_policies.update(config)
                logger.info(f"Loaded retention configuration from {config_file}")
        else:
            # Create default config file
            with open(config_file, 'w') as f:
                json.dump(self.retention_policies, f, indent=2)
                logger.info(f"Created default retention configuration: {config_file}")
    
    def print_retention_report(self):
        """Print a formatted retention policy report"""
        stats = self.get_retention_statistics()
        
        print("\n" + "="*50)
        print("DATA RETENTION POLICY REPORT")
        print("="*50)
        print(f"Generated: {datetime.now().isoformat()}")
        print()
        
        print("RECENT POLICY EXECUTIONS (Last 30 days)")
        print("-" * 40)
        for execution in stats.get('recent_executions', []):
            print(f"Table: {execution['table_name']}")
            print(f"  Last execution: {execution['last_execution']}")
            print(f"  Records deleted: {execution['total_deleted']}")
            print(f"  Execution count: {execution['execution_count']}")
            print()
        
        print("ARCHIVE TABLE SIZES")
        print("-" * 20)
        archive_counts = stats.get('archive_counts', {})
        for table, count in archive_counts.items():
            print(f"{table}: {count:,} records")
        
        print("\nCURRENT RETENTION POLICIES")
        print("-" * 30)
        for category, policies in self.retention_policies.items():
            print(f"{category.upper().replace('_', ' ')}:")
            for table, days in policies.items():
                if days == -1:
                    policy = "Keep forever"
                else:
                    policy = f"{days} days"
                print(f"  {table}: {policy}")
            print()

def main():
    """Main function to run retention policy management"""
    import argparse
    
    parser = argparse.ArgumentParser(description='Data retention policy manager for Mini Data Warehouse')
    parser.add_argument('--config', help='Path to retention configuration file')
    parser.add_argument('--report', action='store_true', help='Generate retention policy report')
    parser.add_argument('--execute', action='store_true', help='Execute retention policies')
    
    args = parser.parse_args()
    
    manager = RetentionPolicyManager()
    
    # Load configuration if provided
    if args.config:
        manager.load_retention_config(args.config)
    else:
        manager.load_retention_config()  # Load default config
    
    success = True
    
    # Execute retention policies if requested
    if args.execute:
        if not manager.execute_retention_policies():
            success = False
    
    # Generate report if requested or as default action
    if args.report or not args.execute:
        manager.print_retention_report()
    
    return 0 if success else 1

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