fix_postgres_sequences.py

#!/usr/bin/env python3
"""
Fix PostgreSQL sequences
This script updates the sequences to start from the highest existing ID + 1
"""

import psycopg2
import logging

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

# Database configuration
POSTGRES_CONFIG = {
    'host': 'localhost',
    'port': 5432,
    'database': 'arrest_data',
    'user': 'arrest_user',
    'password': 'arrest_password'
}

def fix_sequences():
    """Fix all sequences in the arrest_data schema"""
    try:
        conn = psycopg2.connect(**POSTGRES_CONFIG)
        cursor = conn.cursor()
        
        # Set search path
        cursor.execute("SET search_path TO arrest_data, public;")
        
        # Fix clients sequence
        cursor.execute("SELECT MAX(id) FROM clients;")
        max_client_id = cursor.fetchone()[0]
        if max_client_id:
            cursor.execute(f"SELECT setval('clients_id_seq', {max_client_id});")
            logger.info(f"✅ Fixed clients sequence to start from {max_client_id + 1}")
        
        # Fix sessions sequence
        cursor.execute("SELECT MAX(id) FROM sessions;")
        max_session_id = cursor.fetchone()[0]
        if max_session_id:
            cursor.execute(f"SELECT setval('sessions_id_seq', {max_session_id});")
            logger.info(f"✅ Fixed sessions sequence to start from {max_session_id + 1}")
        
        # Fix queries sequence
        cursor.execute("SELECT MAX(id) FROM queries;")
        max_query_id = cursor.fetchone()[0]
        if max_query_id:
            cursor.execute(f"SELECT setval('queries_id_seq', {max_query_id});")
            logger.info(f"✅ Fixed queries sequence to start from {max_query_id + 1}")
        
        # Fix messages sequence
        cursor.execute("SELECT MAX(id) FROM messages;")
        max_message_id = cursor.fetchone()[0]
        if max_message_id:
            cursor.execute(f"SELECT setval('messages_id_seq', {max_message_id});")
            logger.info(f"✅ Fixed messages sequence to start from {max_message_id + 1}")
        
        conn.commit()
        cursor.close()
        conn.close()
        
        logger.info("🎉 All sequences have been fixed!")
        return True
        
    except Exception as e:
        logger.error(f"❌ Failed to fix sequences: {e}")
        return False

def show_current_data():
    """Show current data counts"""
    try:
        conn = psycopg2.connect(**POSTGRES_CONFIG)
        cursor = conn.cursor()
        
        # Set search path
        cursor.execute("SET search_path TO arrest_data, public;")
        
        tables = ['clients', 'sessions', 'queries', 'messages']
        for table in tables:
            cursor.execute(f"SELECT COUNT(*) FROM {table};")
            count = cursor.fetchone()[0]
            logger.info(f"📊 {table}: {count} records")
        
        cursor.close()
        conn.close()
        
    except Exception as e:
        logger.error(f"❌ Failed to show data: {e}")

if __name__ == "__main__":
    logger.info("🔧 Fixing PostgreSQL sequences after migration...")
    
    show_current_data()
    print()
    
    if fix_sequences():
        logger.info("✅ Sequences fixed successfully!")
        logger.info("You can now run the test script: python test_postgres.py")
    else:
        logger.error("❌ Failed to fix sequences")
        exit(1)