Event-Driven Database Ingestion from S3 Data with Amazon MWAA and AWS Glue

def lambda_handler(event, context):
    """Lambda handler triggered by S3 events"""
    mwaa_env_name = os.environ['MWAA_ENVIRONMENT_NAME']
    dag_id = os.environ['AIRFLOW_DAG_ID']
    
    for record in event.get('Records', []):
        s3_info = record.get('s3', {})
        s3_bucket = s3_info.get('bucket', {}).get('name')
        s3_key = urllib.parse.unquote_plus(s3_info.get('object', {}).get('key'))
        
        # Skip non-CSV files and archived files
        if not s3_key.endswith('.csv') or '/archive/' in s3_key:
            continue
        
        # Extract directory from the uploaded file
        s3_directory = '/'.join(s3_key.split('/')[:-1]) + '/' if '/' in s3_key else ''
        
        # Build DAG trigger command with directory configuration
        dag_conf = json.dumps({
            's3_bucket': s3_bucket,
            's3_directory': s3_directory,
            'triggered_by_file': s3_key
        })
        
        # Use MWAA CLI token method for triggering
        cli_token_response = mwaa_client.create_cli_token(Name=mwaa_env_name)
        cli_command = f'dags trigger {dag_id} --conf "{dag_conf_escaped}"'
        
        # Execute via MWAA CLI endpoint
        trigger_dag_via_cli(cli_token_response, cli_command)

dag = DAG(
    's3_healthcare_to_postgres_elt',
    default_args=default_args,
    schedule_interval=None,  # Triggered by events, not scheduled
    max_active_runs=1,
    catchup=False,
    tags=['healthcare', 'elt', 's3', 'postgresql', 'glue'],
)

def list_and_prepare_files(**context):
    """List CSV files in S3 directory and prepare for batch processing"""
    s3_bucket = context['ti'].xcom_pull(key='s3_bucket', task_ids='validate_config')
    s3_directory = context['ti'].xcom_pull(key='s3_directory', task_ids='validate_config')
    
    # List all CSV files in directory, excluding archived files
    s3_hook = S3Hook(aws_conn_id='aws_default')
    objects = s3_hook.list_keys(bucket_name=s3_bucket, prefix=s3_directory)
    
    csv_files = [
        obj for obj in objects 
        if obj.endswith('.csv') and '/archive/' not in obj
    ]
    
    # Limit to 50 files per batch
    if len(csv_files) > 50:
        csv_files = csv_files[:50]
    
    # Handle empty directory
    if len(csv_files) == 0:
        context['ti'].xcom_push(key='skip_processing', value=True)
        return {'skip_processing': True}
    
    return {'csv_files': csv_files, 'file_count': len(csv_files)}

# Glue job processes multiple files in a single run
load_data = GlueJobOperator(
    task_id='glue_s3_to_postgres',
    job_name='healthcare-s3-to-postgresql',
    script_args={
        '--S3_BUCKET': '{{ ti.xcom_pull(task_ids="list_files", key="glue_s3_bucket") }}',
        '--S3_DIRECTORY': '{{ ti.xcom_pull(task_ids="list_files", key="glue_s3_directory") }}',
        '--CSV_FILES': '{{ ti.xcom_pull(task_ids="list_files", key="glue_csv_files") | tojson }}',
    },
    aws_conn_id='aws_default',
    dag=dag,
)

-- ============================================
-- 1. Admissions Table (Raw Data)
-- ============================================
CREATE TABLE admissions (
    admission_id VARCHAR(50) PRIMARY KEY,
    patient_id VARCHAR(50) NOT NULL,
    patient_name VARCHAR(200),
    age INTEGER,
    gender VARCHAR(20),
    admission_date TIMESTAMP NOT NULL,
    discharge_date TIMESTAMP,
    diagnosis_code VARCHAR(20),
    diagnosis_description TEXT,
    department VARCHAR(100),
    attending_physician VARCHAR(200),
    admission_type VARCHAR(50),
    insurance_provider VARCHAR(50),
    total_charges DECIMAL(10, 2),
    length_of_stay_days INTEGER,
    readmission_30_days BOOLEAN DEFAULT FALSE,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);

# Parse CSV files list and read multiple files at once
csv_files = json.loads(args['CSV_FILES'])
file_paths = [f"s3://{args['S3_BUCKET']}/{file_key}" for file_key in csv_files]

# Spark can read multiple CSV files efficiently
df = spark.read \
    .option("header", "true") \
    .option("inferSchema", "true") \
    .csv(file_paths)

# Data transformations with type casting
df_transformed = df \
    .withColumn("admission_date", to_timestamp(col("admission_date"), "yyyy-MM-dd HH:mm:ss")) \
    .withColumn("total_charges", col("total_charges").cast("decimal(10,2)")) \
    .withColumn("readmission_30_days", 
               when(lower(col("readmission_30_days")).isin("true", "1", "yes"), True)
               .otherwise(False)) \
    .withColumn("created_at", current_timestamp())

# Create staging table and perform upsert
staging_table = f"admissions_staging_{int(time.time())}"

# Load data to staging table
df_transformed.write \
    .mode("overwrite") \
    .format("jdbc") \
    .option("url", jdbc_url) \
    .option("dbtable", staging_table) \
    .save()

# Use Spark's JVM for database operations
jvm = spark.sparkContext._jvm
connection = jvm.java.sql.DriverManager.getConnection(jdbc_url, username, password)
statement = connection.createStatement()

# Perform upsert: delete existing, insert all, cleanup
statement.executeUpdate(f"DELETE FROM admissions WHERE admission_id IN (SELECT admission_id FROM {staging_table})")
statement.executeUpdate(f"INSERT INTO admissions SELECT * FROM {staging_table}")
statement.executeUpdate(f"DROP TABLE {staging_table}")

transform_1_cleanse = PostgresOperator(
    task_id='transform_1_data_cleansing',
    postgres_conn_id='postgres_rds',
    sql='CALL staging.sp_cleanse_patient_admissions();',
    dag=dag,
)

-- Transform 1: Data Cleansing and Validation
CREATE OR REPLACE PROCEDURE staging.sp_cleanse_patient_admissions()
LANGUAGE plpgsql
AS $BODY$
BEGIN
    -- Update gender values to standardized format
    UPDATE admissions 
    SET gender = CASE 
        WHEN UPPER(gender) IN ('M', 'MALE') THEN 'Male'
        WHEN UPPER(gender) IN ('F', 'FEMALE') THEN 'Female'
        ELSE gender
    END
    WHERE gender IS NOT NULL;
    
    -- Validate age range (0-120)
    UPDATE admissions 
    SET age = NULL 
    WHERE age < 0 OR age > 120;
    
    -- Validate charges (>= 0)
    UPDATE admissions 
    SET total_charges = 0 
    WHERE total_charges < 0;
    
    -- Remove duplicates based on admission_id (keep latest)
    DELETE FROM admissions a1 
    USING admissions a2 
    WHERE a1.admission_id = a2.admission_id 
    AND a1.created_at < a2.created_at;
    
    RAISE NOTICE 'Data cleansing completed';
END;
$BODY$;

transform_2_aggregate = PostgresOperator(
    task_id='transform_2_business_aggregations',
    postgres_conn_id='postgres_rds',
    sql='CALL analytics.sp_aggregate_admission_metrics();',
    dag=dag,
)

    -- Transform 2: Business Aggregations
CREATE OR REPLACE PROCEDURE analytics.sp_aggregate_admission_metrics()
LANGUAGE plpgsql
AS $BODY$
BEGIN
    -- Create or replace aggregated metrics table
    DROP TABLE IF EXISTS analytics.daily_admission_metrics;
    
    CREATE TABLE analytics.daily_admission_metrics AS
    SELECT 
        DATE(admission_date) as admission_date,
        admission_type,
        department,
        COUNT(*) as total_admissions,
        ROUND(AVG(age)::NUMERIC, 1) as avg_age,
        ROUND(AVG(length_of_stay_days)::NUMERIC, 1) as avg_length_of_stay,
        ROUND(AVG(total_charges)::NUMERIC, 2) as avg_charges,
        COUNT(DISTINCT patient_id) as unique_patients,
        COUNT(DISTINCT attending_physician) as unique_physicians
    FROM admissions
    WHERE admission_date IS NOT NULL
    GROUP BY DATE(admission_date), admission_type, department
    ORDER BY admission_date DESC, total_admissions DESC;
    
    RAISE NOTICE 'Business aggregations completed';
END;
$BODY$;

dataBucket.addEventNotification(
  s3.EventType.OBJECT_CREATED,
  new s3n.LambdaDestination(lambdaFunction),
  {
    prefix: 'admissions/',
    suffix: '.csv',
  }
);