whos_your_daddy_common 0.2.0

import argparse
import common
import psycopg
from typing import Optional
import yaml

def build_clap() -> Optional[argparse.Namespace]:
    parser = argparse.ArgumentParser(prog="delete_db")

    parser.add_argument(
        "-f",
        help="YAML file containing database config data."
    )

    args = parser.parse_args()

    return args

def delete_database(
        db: common.Database
):
    """Delete the database"""
    conn = None
    try:
        conn = psycopg.connect(
            host=db.host,
            port=db.port,
            user=db.admin_user,
            password=db.admin_password,
            autocommit = True
        )

        cur = conn.cursor()
        cur.execute(f"DROP DATABASE {db.name};")
        
    #     # Check if database exists
    #     cur.execute(
    #         "SELECT 1 FROM pg_database WHERE datname = %s",
    #         (DB_CONFIG['new_database'],)
    #     )
    #     exists = cur.fetchone()
        
    #     if exists:
    #         # Terminate all connections to the database
    #         cur.execute(sql.SQL("""
    #             SELECT pg_terminate_backend(pid)
    #             FROM pg_stat_activity
    #             WHERE datname = %s AND pid <> pg_backend_pid()
    #         """), (DB_CONFIG['new_database'],))
            
    #         # Drop database
    #         cur.execute(sql.SQL(
    #             "DROP DATABASE {}"
    #         ).format(sql.Identifier(DB_CONFIG['new_database'])))
    #         print(f"✓ Database '{DB_CONFIG['new_database']}' deleted successfully")
    #     else:
    #         print(f"✓ Database '{DB_CONFIG['new_database']}' does not exist (already deleted)")
        
    #     cur.close()
    except Exception as e:
        print(f"✗ Error deleting database: {e}")
        if conn:
            conn.close()
        return
    finally:
        if conn:
            conn.close()

def confirm_deletion(
        db: common.Database
) -> bool:
    """Ask user to confirm deletion"""
    print("=" * 50)
    print("WARNING: This will delete:")
    print(f"  - Database: {db.name}")
    print(f"  - Tables: {', '.join(str(obj.name) for obj in db.tables)}")
    print(f"  - Users: {', '.join(str(obj.username) for obj in db.users)}")
    print("=" * 50)
    
    # Prompt the user for confirmation about deleting data, and strip
    # whitespace and make the response only lowercase.
    response = input("\nAre you sure you want to proceed? (yes/no): ").strip().lower()
    return response in ['yes', 'y']

def main():
    """Main teardown function"""

    args = build_clap()

    if args is None:
        print("Error - Failed to parse arguments")
        return
    
    #===========================================================================
    # Get database configuration.
    databases: list[common.Database] = []

    with open(args.f, 'r') as f:
        config_dict = yaml.safe_load(f)

        for db in config_dict['db']:
            databases.append(common.Database.from_dict(db))

    print(databases)

    #===========================================================================
    # Begin the deletion process.
    try:
        for db in databases:
            # Prompt the user to confirm deletion.
            if not confirm_deletion(db):
                print("\n✗ Teardown cancelled by user\n")
                return
            
            print(f"Starting PostgreSQL database teardown for {db.name}\n")

            delete_database(db)

    except Exception as e:
        print("\n" + "=" * 50)
        print("✗ Teardown failed!")
        print("=" * 50)

if __name__ == "__main__":
    main()