giglez/scripts/cleanup_database.py

#!/usr/bin/env python3
"""
Cleanup Database - Remove test/mock data

Remove test and mock captures from the database
"""

import sys
import argparse
import requests
from pathlib import Path


def cleanup_by_source(api_url, data_source):
    """Delete captures by data source"""
    print("=" * 80)
    print(f"Cleaning up captures with data_source='{data_source}'")
    print("=" * 80)

    endpoint = f"{api_url}/api/v1/admin/cleanup?data_source={data_source}"

    try:
        response = requests.delete(endpoint)

        if response.status_code == 200:
            result = response.json()
            print(f"\n✅ {result['message']}")
            print(f"   Deleted: {result['deleted_count']}")
            print(f"   Remaining: {result['remaining_count']}")
            return 0
        else:
            print(f"\n❌ HTTP {response.status_code}: {response.text}")
            return 1

    except Exception as e:
        print(f"\n❌ Error: {e}")
        return 1


def cleanup_session(api_url, session_id):
    """Delete captures by session ID"""
    print("=" * 80)
    print(f"Cleaning up session: {session_id}")
    print("=" * 80)

    endpoint = f"{api_url}/api/v1/admin/cleanup/session/{session_id}"

    try:
        response = requests.delete(endpoint)

        if response.status_code == 200:
            result = response.json()
            print(f"\n✅ {result['message']}")
            print(f"   Deleted: {result['deleted_count']}")
            print(f"   Remaining: {result['remaining_count']}")
            return 0
        else:
            print(f"\n❌ HTTP {response.status_code}: {response.text}")
            return 1

    except Exception as e:
        print(f"\n❌ Error: {e}")
        return 1


def cleanup_all(api_url, confirm=False):
    """Delete ALL captures"""
    print("=" * 80)
    print("⚠️  WARNING: Deleting ALL captures!")
    print("=" * 80)

    if not confirm:
        print("\n❌ Confirmation required. Use --confirm to proceed.")
        return 1

    endpoint = f"{api_url}/api/v1/admin/cleanup/all"

    try:
        response = requests.delete(endpoint)

        if response.status_code == 200:
            result = response.json()
            print(f"\n✅ {result['message']}")
            print(f"   Deleted: {result['deleted_count']}")
            return 0
        else:
            print(f"\n❌ HTTP {response.status_code}: {response.text}")
            return 1

    except Exception as e:
        print(f"\n❌ Error: {e}")
        return 1


def show_stats(api_url):
    """Show current database stats"""
    try:
        response = requests.get(f"{api_url}/api/v1/stats/summary")

        if response.status_code == 200:
            stats = response.json()
            print("\n📊 Current Database Stats:")
            print(f"   Total captures: {stats.get('total_captures', 0)}")

            if 'data_sources' in stats:
                print("\n   By data source:")
                for source, count in stats['data_sources'].items():
                    print(f"     {source}: {count}")

            print()

    except Exception as e:
        print(f"⚠️  Could not fetch stats: {e}")


def main():
    parser = argparse.ArgumentParser(
        description="Cleanup test/mock data from GigLez database"
    )

    parser.add_argument(
        '--api-url',
        default='http://localhost:8000',
        help='GigLez API URL (default: http://localhost:8000)'
    )

    parser.add_argument(
        '--source',
        choices=['test', 'mock', 'production'],
        help='Delete captures by data source'
    )

    parser.add_argument(
        '--session',
        help='Delete captures by session ID'
    )

    parser.add_argument(
        '--all',
        action='store_true',
        help='Delete ALL captures (requires --confirm)'
    )

    parser.add_argument(
        '--confirm',
        action='store_true',
        help='Confirm deletion (required for --all)'
    )

    parser.add_argument(
        '--stats',
        action='store_true',
        help='Show database stats before and after cleanup'
    )

    args = parser.parse_args()

    # Show initial stats
    if args.stats:
        print("BEFORE cleanup:")
        show_stats(args.api_url)

    # Perform cleanup
    result = 0

    if args.all:
        result = cleanup_all(args.api_url, args.confirm)
    elif args.source:
        result = cleanup_by_source(args.api_url, args.source)
    elif args.session:
        result = cleanup_session(args.api_url, args.session)
    else:
        print("❌ Error: Must specify --source, --session, or --all")
        parser.print_help()
        return 1

    # Show final stats
    if args.stats and result == 0:
        print("\nAFTER cleanup:")
        show_stats(args.api_url)

    return result


if __name__ == '__main__':
    sys.exit(main())