import sys import time from datetime import datetime, timedelta from typing import Optional import boto3 from utils.postgres.postgres_client import get_postgres_connection from pyspark.context import SparkContext # type: ignore from awsglue.utils import getResolvedOptions # type: ignore from awsglue.context import GlueContext # type: ignore from awsglue.job import Job # type: ignore # Initialize Glue context args = getResolvedOptions(sys.argv, ["JOB_NAME"]) sc = SparkContext() glueContext = GlueContext(sc) spark = glueContext.spark_session job = Job(glueContext) job.init(args["JOB_NAME"], args) # Get configurable parameters with defaults def get_job_parameter(key: str, default: str) -> str: """Get job parameter with fallback to default value.""" try: job_args = getResolvedOptions(sys.argv, [key]) return job_args[key] except: return default # Configuration parameters START_DATE = get_job_parameter("START_DATE", "2025-07-21") END_DATE = get_job_parameter("END_DATE", "2025-08-14") BATCH_SIZE = int(get_job_parameter("BATCH_SIZE", "10000")) BATCH_DELAY_SECONDS = float(get_job_parameter("BATCH_DELAY_SECONDS", "0.2")) DRY_RUN = get_job_parameter("DRY_RUN", "false").lower() == "true" print(f"=== Billing ClipOwnership Cleanup Job Started ===") print(f"Start Date: {START_DATE}") print(f"End Date: {END_DATE}") print(f"Batch Size: {BATCH_SIZE}") print(f"Batch Delay: {BATCH_DELAY_SECONDS} seconds") print(f"Dry Run: {DRY_RUN}") def parse_date(date_str: str) -> datetime: """Parse date string to datetime object.""" return datetime.strptime(date_str, "%Y-%m-%d") def get_date_range(start_date: str, end_date: str) -> list: """Generate list of dates between start and end (inclusive).""" start_dt = parse_date(start_date) end_dt = parse_date(end_date) dates = [] current_dt = start_dt while current_dt <= end_dt: dates.append(current_dt.strftime("%Y-%m-%d")) current_dt += timedelta(days=1) return dates def count_records_for_date(cursor, date_str: str) -> int: """Count records for a specific date.""" count_query = """ SELECT COUNT(*) FROM billing_clipownership WHERE DATE(created_at) = %s """ cursor.execute(count_query, (date_str,)) result = cursor.fetchone() return result[0] if result else 0 def delete_batch_for_date(cursor, date_str: str, batch_size: int, dry_run: bool = False) -> int: """Delete a batch of records for a specific date.""" if dry_run: print(f"[DRY RUN] Would delete batch of {batch_size} records for {date_str}") return min(batch_size, count_records_for_date(cursor, date_str)) delete_query = """ DELETE FROM billing_clipownership WHERE ctid IN ( SELECT ctid FROM billing_clipownership WHERE DATE(created_at) = %s LIMIT %s ) """ cursor.execute(delete_query, (date_str, batch_size)) return cursor.rowcount def cleanup_data_for_date(connection, date_str: str, batch_size: int, batch_delay: float, dry_run: bool = False) -> dict: """Clean up all data for a specific date in batches.""" cursor = connection.cursor() # Get initial count initial_count = count_records_for_date(cursor, date_str) if initial_count == 0: print(f"No records found for {date_str}") return { "date": date_str, "initial_count": 0, "deleted_count": 0, "batches_processed": 0, "final_count": 0 } print(f"Processing {date_str}: {initial_count:,} records to delete") total_deleted = 0 batch_num = 0 while True: batch_num += 1 batch_start_time = time.time() # Delete batch deleted_in_batch = delete_batch_for_date(cursor, date_str, batch_size, dry_run) if deleted_in_batch == 0: break total_deleted += deleted_in_batch batch_time = time.time() - batch_start_time if not dry_run: connection.commit() # Log progress remaining = initial_count - total_deleted progress_pct = (total_deleted / initial_count) * 100 print(f" Batch {batch_num}: Deleted {deleted_in_batch:,} records " f"({batch_time:.2f}s) - Progress: {progress_pct:.1f}% " f"({remaining:,} remaining)") # Apply delay between batches if batch_delay > 0: time.sleep(batch_delay) # Safety check - avoid infinite loop if batch_num > 10000: # Max 100M records with 10k batches print(f"WARNING: Maximum batch count reached for {date_str}") break # Get final count final_count = count_records_for_date(cursor, date_str) cursor.close() return { "date": date_str, "initial_count": initial_count, "deleted_count": total_deleted, "batches_processed": batch_num, "final_count": final_count } def main(): """Main cleanup process.""" start_time = time.time() # Get date range dates_to_process = get_date_range(START_DATE, END_DATE) print(f"Processing {len(dates_to_process)} dates: {dates_to_process[0]} to {dates_to_process[-1]}") # Connect to database connection = get_postgres_connection("billing_clipownership_cleanup_manual") # Process each date total_stats = { "dates_processed": 0, "total_initial_count": 0, "total_deleted": 0, "total_batches": 0, "date_results": [] } try: for i, date_str in enumerate(dates_to_process, 1): print(f"\n=== Processing Date {i}/{len(dates_to_process)}: {date_str} ===") date_start_time = time.time() date_stats = cleanup_data_for_date( connection, date_str, BATCH_SIZE, BATCH_DELAY_SECONDS, DRY_RUN ) date_time = time.time() - date_start_time # Update totals total_stats["dates_processed"] += 1 total_stats["total_initial_count"] += date_stats["initial_count"] total_stats["total_deleted"] += date_stats["deleted_count"] total_stats["total_batches"] += date_stats["batches_processed"] total_stats["date_results"].append(date_stats) print(f"Completed {date_str}: {date_stats['deleted_count']:,} deleted " f"in {date_stats['batches_processed']} batches ({date_time:.1f}s)") # Progress summary overall_progress = (i / len(dates_to_process)) * 100 print(f"Overall Progress: {overall_progress:.1f}% ({i}/{len(dates_to_process)} dates)") finally: connection.close() # Final summary total_time = time.time() - start_time print(f"\n=== CLEANUP SUMMARY ===") print(f"Total Runtime: {total_time:.1f} seconds ({total_time/60:.1f} minutes)") print(f"Dates Processed: {total_stats['dates_processed']}") print(f"Total Records Found: {total_stats['total_initial_count']:,}") print(f"Total Records Deleted: {total_stats['total_deleted']:,}") print(f"Total Batches Processed: {total_stats['total_batches']}") if total_stats['total_initial_count'] > 0: success_rate = (total_stats['total_deleted'] / total_stats['total_initial_count']) * 100 print(f"Success Rate: {success_rate:.1f}%") # Per-date summary print(f"\n=== PER-DATE RESULTS ===") for result in total_stats['date_results']: if result['initial_count'] > 0: print(f"{result['date']}: {result['initial_count']:,} -> " f"{result['deleted_count']:,} deleted " f"({result['batches_processed']} batches)") print(f"\n=== Job Completed Successfully ===") if __name__ == "__main__": main() # Commit the Glue job job.commit()