Airflow Tools

airflow_tools

# 1. List existing DAGs to verify connectivity
result = airflow_action(action="list_dags")
print(f"Connected to Airflow. Found {len(result['dags'])} DAGs")

# 2. Set configuration variables
airflow_action(
    action="set_variable",
    variable_key="source_database",
    variable_value="CUSTOMER_DB"
)

airflow_action(
    action="set_variable",
    variable_key="target_schema",
    variable_value="ANALYTICS"
)

# 3. Create database connection
airflow_action(
    action="create_connection",
    conn_id="postgres_source",
    conn_type="postgres",
    host="source-db.company.com",
    login="etl_user",
    password="secure_password",
    schema="customers",
    port=5432
)

# 4. Test the connection
test_result = airflow_action(
    action="test_connection",
    conn_id="postgres_source"
)
print(f"Connection test: {test_result['status']}")

# 5. Create comprehensive ETL DAG
dag_code = """
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.postgres.operators.postgres import PostgresOperator
from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator
from airflow.models import Variable
from datetime import datetime, timedelta
import logging

default_args = {
    'owner': 'data_engineering',
    'depends_on_past': False,
    'start_date': datetime(2024, 1, 1),
    'email': ['data-team@company.com'],
    'email_on_failure': True,
    'email_on_retry': False,
    'retries': 2,
    'retry_delay': timedelta(minutes=5),
}

def validate_data(**context):
    source_db = Variable.get("source_database")
    logging.info(f"Validating data from {source_db}")
    return {"status": "validated", "row_count": 15000}

def send_notification(**context):
    ti = context['task_instance']
    validation = ti.xcom_pull(task_ids='validate_data')
    logging.info(f"Pipeline complete. Processed {validation['row_count']} rows")

with DAG(
    'customer_data_pipeline',
    default_args=default_args,
    description='Extract customer data from Postgres, transform, load to Snowflake',
    schedule_interval='0 2 * * *',
    catchup=False,
    max_active_runs=1,
    tags=['etl', 'customers', 'daily']
) as dag:
    
    extract = PostgresOperator(
        task_id='extract_customer_data',
        postgres_conn_id='postgres_source',
        sql='''
            SELECT customer_id, name, email, created_date, status
            FROM customers
            WHERE updated_at >= '{{ ds }}'
        ''',
        dag=dag
    )
    
    validate = PythonOperator(
        task_id='validate_data',
        python_callable=validate_data,
        provide_context=True,
        dag=dag
    )
    
    transform = PostgresOperator(
        task_id='transform_data',
        postgres_conn_id='postgres_source',
        sql='''
            CREATE TEMP TABLE transformed_customers AS
            SELECT 
                customer_id,
                UPPER(name) as name,
                LOWER(email) as email,
                created_date,
                CASE 
                    WHEN status = 'A' THEN 'ACTIVE'
                    WHEN status = 'I' THEN 'INACTIVE'
                    ELSE 'UNKNOWN'
                END as status_desc
            FROM customers_staging;
        ''',
        dag=dag
    )
    
    load = SnowflakeOperator(
        task_id='load_to_snowflake',
        snowflake_conn_id='snowflake_prod',
        sql='''
            MERGE INTO {{ var.value.target_schema }}.CUSTOMERS dst
            USING staging.customers_temp src
            ON dst.customer_id = src.customer_id
            WHEN MATCHED THEN UPDATE SET
                dst.name = src.name,
                dst.email = src.email,
                dst.status = src.status_desc,
                dst.updated_at = CURRENT_TIMESTAMP()
            WHEN NOT MATCHED THEN INSERT
                (customer_id, name, email, status, created_at)
            VALUES
                (src.customer_id, src.name, src.email, src.status_desc, src.created_date);
        ''',
        dag=dag
    )
    
    notify = PythonOperator(
        task_id='send_notification',
        python_callable=send_notification,
        provide_context=True,
        dag=dag
    )
    
    extract >> validate >> transform >> load >> notify
"""

# 6. Deploy the DAG
airflow_action(
    action="create_dag",
    dag_id="customer_data_pipeline",
    dag_code=dag_code
)
print("DAG created successfully")

# 7. Wait for DAG to be parsed
import time
time.sleep(30)

# 8. Unpause the DAG to enable scheduling
airflow_action(
    action="unpause_dag",
    dag_id="customer_data_pipeline"
)

# 9. Trigger initial run
run_result = airflow_action(
    action="trigger_dag",
    dag_id="customer_data_pipeline",
    conf={"mode": "initial_load", "date_range": "2024-01-01"}
)
print(f"DAG triggered. Run ID: {run_result['dag_run_id']}")

# 10. Monitor the run
time.sleep(10)
status = airflow_action(
    action="get_dag_status",
    dag_id="customer_data_pipeline"
)
print(f"Current status: {status['state']}")

# 11. Get run history
runs = airflow_action(
    action="get_dag_runs",
    dag_id="customer_data_pipeline",
    limit=5
)
for run in runs['dag_runs']:
    print(f"Run {run['execution_date']}: {run['state']} ({run['duration']}s)")

# 12. Get logs for debugging
logs = airflow_action(
    action="get_task_logs",
    dag_id="customer_data_pipeline",
    task_id="validate_data",
    execution_date=run_result['execution_date']
)
print("Validation task logs:")
print(logs['logs'])

from datetime import datetime, timedelta

# Trigger backfill for last 7 days
start_date = datetime.now() - timedelta(days=7)
for i in range(7):
    exec_date = (start_date + timedelta(days=i)).strftime("%Y-%m-%d")
    airflow_action(
        action="trigger_dag",
        dag_id="customer_data_pipeline",
        execution_date=exec_date,
        conf={"backfill": True}
    )
    print(f"Triggered backfill for {exec_date}")

def generate_dag_code(source_table, target_table, schedule):
    return f"""
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

with DAG(
    '{source_table}_to_{target_table}',
    start_date=datetime(2024, 1, 1),
    schedule_interval='{schedule}',
    catchup=False
) as dag:
    
    def process_data():
        print("Processing {source_table} to {target_table}")
    
    task = PythonOperator(
        task_id='process',
        python_callable=process_data
    )
"""

# Generate multiple DAGs
tables = [
    ("customers", "dim_customers", "0 3 * * *"),
    ("orders", "fact_orders", "0 4 * * *"),
    ("products", "dim_products", "0 2 * * *")
]

for source, target, schedule in tables:
    dag_code = generate_dag_code(source, target, schedule)
    airflow_action(
        action="create_dag",
        dag_id=f"{source}_pipeline",
        dag_code=dag_code
    )

from datetime import datetime, timedelta

def should_trigger_dag():
    last_update = get_last_data_update()
    if datetime.now() - last_update < timedelta(hours=1):
        return True
    return False

if should_trigger_dag():
    airflow_action(
        action="trigger_dag",
        dag_id="realtime_pipeline"
    )
else:
    print("Source data not fresh enough, skipping trigger")

def check_dag_health(dag_id, max_failures=3):
    """Check recent DAG run health"""
    runs = airflow_action(
        action="get_dag_runs",
        dag_id=dag_id,
        limit=10
    )
    
    failed_runs = [r for r in runs['dag_runs'] if r['state'] == 'failed']
    
    if len(failed_runs) >= max_failures:
        send_alert(f"DAG {dag_id} has {len(failed_runs)} failures in last 10 runs")
        return False
    
    return True

# Monitor critical DAGs
critical_dags = ["customer_data_pipeline", "financial_reporting", "ml_training"]
for dag_id in critical_dags:
    health = check_dag_health(dag_id)
    print(f"{dag_id}: {'Healthy' if health else 'Unhealthy'}")