import os import json import boto3 # type: ignore import subprocess import zipfile import glob from pathlib import Path env = os.getenv("ENVIRONMENT", "default") if env != 'staging': s3_bucket = "aws-glue-assets-734185074900-us-east-2" glue_job_role = "arn:aws:iam::734185074900:role/suno-glue-job-execution" trigger_start_on_creation = "true" else: s3_bucket = "aws-glue-assets-590183763515-us-east-2" glue_job_role = "arn:aws:iam::590183763515:role/testing-rds-glue-role" trigger_start_on_creation = "false" # Initialize AWS clients glue_client = boto3.client('glue') s3_client = boto3.client('s3') cloudwatch = boto3.client('cloudwatch') # Get changed files from environment variable changed_files = os.environ.get('CHANGED_FILES', '').strip().split('\n') # Process each changed directory that has a config.json processed_dirs = set() events_client = boto3.client('events') if env != 'staging': sns_topic_arn = 'arn:aws:sns:us-east-2:734185074900:GlueJobStack-GlueJobNotificationTopic5F523CE6-o6b9kc6v0B5J' rule_name = 'GlueJobStack-GlueJobStateChangeRuleBEBC3FD6-GSdv7dY0sjCx' def update_glue_failure_rule(): # existing_rule = events_client.describe_rule(Name=rule_name) # existing_pattern = json.loads(existing_rule['EventPattern']) # Extract existing job names from the pattern # existing_job_names = set() # Use a set to enforce uniqueness # job_names = existing_pattern['detail']['jobName'] # for name in job_names: # existing_job_names.add(name) # print(f"Found {len(existing_job_names)} existing job names in rule") # Merge existing job names with new ones (avoiding duplicates) # existing_job_names.add(job_name) # print(f"Combined job list contains {len(existing_job_names)} jobs") # Create the updated event pattern event_pattern = { "source": ["aws.glue"], "detail-type": ["Glue Job State Change"], "detail": { "jobName": [{ "prefix": "rds_to_s3" }], "state": ["FAILED", "TIMEOUT"] } } # Update the rule events_client.put_rule( Name=rule_name, EventPattern=json.dumps(event_pattern), State='ENABLED', Description='Rule for Glue job failures' ) existing_targets = events_client.list_targets_by_rule(Rule=rule_name)['Targets'] # Update targets events_client.put_targets( Rule=rule_name, Targets=existing_targets ) print(f"Updated EventBridge rule '{rule_name}'") # return existing_job_names def adding_new_job_to_failure_monitoring(job_name: str): alarm_name = f'GlueJobFailureAlarm-{job_name}' response = cloudwatch.put_metric_alarm( AlarmName=alarm_name, AlarmDescription=f'Alarm for Glue job failures: {job_name}', ActionsEnabled=True, AlarmActions=[sns_topic_arn], MetricName='GlueJobRunErrors', Namespace='AWS/Glue', Statistic='Sum', Dimensions=[ { 'Name': 'JobName', 'Value': job_name } ], Period=60, EvaluationPeriods=1, Threshold=1.0, ComparisonOperator='GreaterThanOrEqualToThreshold', TreatMissingData='notBreaching' ) print(f"Created alarm: {alarm_name}") # update_glue_failure_rule() return response def process_triggers(glue_client, job_name, triggers_config): existing_triggers = [] existing_trigger_details = {} # Add a dictionary to store full trigger details try: # List all triggers paginator = glue_client.get_paginator('get_triggers') for page in paginator.paginate(): for trigger in page['Triggers']: # Check if this trigger targets our job for action in trigger.get('Actions', []): if action.get('JobName') == job_name: existing_triggers.append(trigger['Name']) existing_trigger_details[trigger['Name']] = trigger # Store full trigger details break except Exception as e: print(f"Error listing existing triggers: {e}") # Process each trigger configuration for i, trigger_config in enumerate(triggers_config): # Generate trigger name if not provided if 'Name' not in trigger_config or not trigger_config['Name']: trigger_name = f"{job_name}-trigger-{i+1}" else: trigger_name = trigger_config['Name'] # Check if this trigger already exists trigger_exists = trigger_name in existing_triggers # Set up base trigger parameters trigger_params = { 'Name': trigger_name, 'Type': trigger_config.get('Type', 'SCHEDULED'), 'Actions': [ { 'JobName': job_name } ] } # Add description if provided if 'Description' in trigger_config: trigger_params['Description'] = trigger_config['Description'] # Add start on create flag if provided if 'StartOnCreation' in trigger_config: trigger_params['StartOnCreation'] = trigger_config['StartOnCreation'] else: trigger_params['StartOnCreation'] = True if trigger_params['StartOnCreation'] == True: if env != 'staging': adding_new_job_to_failure_monitoring(job_name) # Add trigger-specific parameters based on type if trigger_params['Type'] == 'SCHEDULED': # Add schedule expression for scheduled triggers if 'ScheduleExpression' in trigger_config: trigger_params['Schedule'] = trigger_config['ScheduleExpression'] else: print(f"Warning: No ScheduleExpression for scheduled trigger {trigger_name}, skipping") continue elif trigger_params['Type'] == 'CONDITIONAL': # Add predicate for conditional triggers predicate = { 'Logical': trigger_config.get('Logical', 'AND'), 'Conditions': [] } # Process conditions (job completion, etc.) if 'Conditions' in trigger_config: for condition in trigger_config['Conditions']: if 'JobName' in condition: # Job completion condition job_condition = { 'LogicalOperator': condition.get('LogicalOperator', 'EQUALS'), 'JobName': condition['JobName'], 'State': condition.get('State', 'SUCCEEDED') } predicate['Conditions'].append(job_condition) if not predicate['Conditions']: print(f"Warning: No conditions for conditional trigger {trigger_name}, skipping") continue trigger_params['Predicate'] = predicate # Add arguments if provided if 'Arguments' in trigger_config: trigger_params['Actions'][0]['Arguments'] = trigger_config['Arguments'] # Create or update the trigger if trigger_exists: update_params = { 'Name': trigger_name, 'TriggerUpdate': {} } # Populate the TriggerUpdate dictionary with allowed parameters if 'Description' in trigger_params: update_params['TriggerUpdate']['Description'] = trigger_params['Description'] if 'Schedule' in trigger_params: update_params['TriggerUpdate']['Schedule'] = trigger_params['Schedule'] if 'Actions' in trigger_params: update_params['TriggerUpdate']['Actions'] = trigger_params['Actions'] if 'Predicate' in trigger_params: update_params['TriggerUpdate']['Predicate'] = trigger_params['Predicate'] glue_client.update_trigger(**update_params) print(f"Successfully updated trigger {trigger_name} for job {job_name}") # Remove from existing_triggers list to track which ones we've processed if trigger_name in existing_triggers: existing_triggers.remove(trigger_name) else: # Create new trigger glue_client.create_trigger(**trigger_params) print(f"Successfully created trigger {trigger_name} for job {job_name}") for trigger_name in existing_triggers: glue_client.delete_trigger(Name=trigger_name) print(f"Deleted trigger {trigger_name}") def should_package_modules(config): """ Check if a job needs custom module packaging based on config.json. Returns True if the job has --extra-py-files in DefaultArguments. """ default_args = config.get('Job', {}).get('DefaultArguments', {}) return '--extra-py-files' in default_args def package_job_modules(job_dir, job_name, main_script_name): """ Package custom modules for a Glue job into a zip file. Args: job_dir: Directory containing the job files job_name: Name of the Glue job main_script_name: Name of the main script file to exclude Returns: Path to the created zip file, or None if no modules to package """ job_path = Path(job_dir) zip_filename = f"{job_name}_modules.zip" zip_path = job_path / zip_filename # Remove existing zip if it exists if zip_path.exists(): zip_path.unlink() print(f"Removed existing module zip: {zip_path}") # Find all Python files except the main script python_files = [] # Get all .py files recursively for py_file in job_path.rglob("*.py"): # Skip the main script file if py_file.name == main_script_name: continue # Skip common build/test files if any(pattern in py_file.name for pattern in ['test_', '_test.py', 'package_']): continue # Calculate relative path from job directory rel_path = py_file.relative_to(job_path) python_files.append((py_file, rel_path)) if not python_files: print(f"No custom modules found for job {job_name}") return None # Create zip file with modules with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf: for file_path, archive_path in python_files: zipf.write(file_path, archive_path) print(f"Added to module zip: {archive_path}") print(f"Created module package: {zip_path} ({zip_path.stat().st_size / 1024:.1f} KB)") return str(zip_path) def upload_module_zip_to_s3(zip_path, job_name, s3_bucket): """ Upload the module zip file to S3. Args: zip_path: Local path to the zip file job_name: Name of the Glue job s3_bucket: S3 bucket name Returns: S3 path to the uploaded zip file """ s3_key = f"libs/{job_name}_modules.zip" try: subprocess.run([ "aws", "s3", "cp", zip_path, f"s3://{s3_bucket}/{s3_key}" ], check=True) s3_path = f"s3://{s3_bucket}/{s3_key}" print(f"Uploaded module zip to: {s3_path}") return s3_path except subprocess.CalledProcessError as e: print(f"Error uploading module zip to S3: {e}") raise def find_job_directory(file_path): """ Find the job directory containing config.json for a changed file. This function walks up the directory tree from the changed file until it finds a directory containing config.json, which indicates a Glue job directory. Args: file_path: Path to the changed file Returns: str: Path to the job directory containing config.json, or None if not found """ current_path = Path(file_path).parent # Walk up the directory tree while current_path != current_path.parent: # Stop at filesystem root config_path = current_path / "config.json" if config_path.exists(): print(f"Found job directory for {file_path}: {current_path}") return str(current_path) current_path = current_path.parent # No job directory found print(f"No job directory found for {file_path}") return None for file_path in changed_files: print(f"Processing file: {file_path}") # Find the job directory containing config.json for this file job_dir = find_job_directory(file_path) if not job_dir: print(f"Skipping {file_path} - no job directory found") continue # Skip if we've already processed this job directory if job_dir in processed_dirs: print(f"Already processed job directory {job_dir}, skipping") continue # Use the job directory as our working directory dir_path = job_dir config_path = os.path.join(job_dir, 'config.json') # Read the configuration with open(config_path, 'r') as f: content = f.read() content = content.replace('{s3bucket}', s3_bucket) content = content.replace('{gluejobrole}', glue_job_role) content = content.replace('"{triggerstartoncreat}"', trigger_start_on_creation) config = json.loads(content) print(f"Found config.json in {os.path.dirname(config_path)}") job_name = config['Job']['Name'] # Find Python script(s) in the directory py_file = f"{dir_path}/{job_name}.py" if not os.path.exists(py_file): print(f"No Python scripts found in {dir_path}, skipping...") continue # Use the first Python file as the main script script_path = py_file script_name = os.path.basename(py_file) s3_prefix = "scripts" s3_key = f"{s3_prefix}/{script_name}" print(f"Processing script: {script_path} for job: {job_name}") # Upload the script to S3 try: subprocess.run(["aws", "s3", "cp", script_path, f"s3://{s3_bucket}/{s3_key}"], check=True) print(f"Uploaded {script_path} to s3://{s3_bucket}/{s3_key}") except subprocess.CalledProcessError as e: print(f"Error uploading script to S3: {e}") continue # Check if this job needs custom module packaging if should_package_modules(config): print(f"Job {job_name} needs custom module packaging") # Package custom modules zip_path = package_job_modules(dir_path, job_name, script_name) if zip_path: # Upload module zip to S3 s3_module_path = upload_module_zip_to_s3(zip_path, job_name, s3_bucket) # Append the module package to existing --extra-py-files job_config = config['Job'] if 'DefaultArguments' in job_config and '--extra-py-files' in job_config['DefaultArguments']: existing_files = job_config['DefaultArguments']['--extra-py-files'] # Append the new module package to existing files (comma-separated) job_config['DefaultArguments']['--extra-py-files'] = f"{existing_files},{s3_module_path}" print(f"Updated --extra-py-files from: {existing_files}") print(f"Updated --extra-py-files to: {existing_files},{s3_module_path}") # Clean up local zip file try: os.remove(zip_path) print(f"Cleaned up local zip file: {zip_path}") except OSError as e: print(f"Warning: Could not remove local zip file {zip_path}: {e}") else: print(f"No modules found to package for job {job_name}") # Set job update parameters job_update = config['Job'] # Check if job exists and update or create accordingly try: glue_client.get_job(JobName=job_name) # Job exists, update it print(f"Updating existing job: {job_name}") job_update.pop('Name') response = glue_client.update_job( JobName=job_name, JobUpdate=job_update ) print(f"Successfully updated job: {job_name}") except glue_client.exceptions.EntityNotFoundException: # Job doesn't exist, create it print(f"Creating new job: {job_name}") response = glue_client.create_job(**job_update) print(f"Successfully created job: {job_name}") if 'Triggers' in config: process_triggers(glue_client, job_name, config['Triggers']) # Add job directory to processed list to avoid reprocessing processed_dirs.add(job_dir)