# Airflow: Orchestrates WHERE and WHEN tasks run, not HOW they process data

from airflow import DAG
from airflow.operators.bash import BashOperator
from airflow.providers.apache.spark.operators.spark_submit import SparkSubmitOperator
from airflow.providers.google.cloud.operators.bigquery import BigQueryInsertJobOperator
from datetime import datetime

with DAG(
    'etl_orchestration',
    start_date=datetime(2024, 1, 1),
    schedule='@daily',
    catchup=False
) as dag:

    # Airflow orchestrates external tools
    extract_postgres = BashOperator(
        task_id='extract_from_postgres',
        bash_command='python /scripts/extract_postgres.py'
    )

    # Spark does the heavy transformation
    transform_spark = SparkSubmitOperator(
        task_id='transform_data',
        application='/spark_jobs/transform.py',
        conn_id='spark_cluster'
    )

    # BigQuery loads the data
    load_bigquery = BigQueryInsertJobOperator(
        task_id='load_to_bigquery',
        configuration={
            'load': {
                'sourceUris': ['gs://bucket/transformed/*'],
                'destinationTable': {'projectId': 'my-project', 'datasetId': 'analytics', 'tableId': 'sales'}
            }
        }
    )

    extract_postgres >> transform_spark >> load_bigquery

# ETL Tools: Define HOW data is transformed (visual or code)
# They are the WORKERS, not the ORCHESTRATOR

# Example: Talend Job (conceptual)
# 1. GUI-based job design
# 2. Drag-and-drop components
# 3. Built-in transformations
# 4. Usually handles: Extract → Transform → Load in one tool

# Limitations:
# - Often proprietary and expensive
# - Limited scheduling capabilities
# - Harder to version control (XML-based)
# - Less flexible for complex workflows

# Detailed execution flow

"""
1. DAG FILE PARSING (Scheduler)
   - Scheduler scans dag_folder every min_file_process_interval
   - Parses Python files to discover DAG objects
   - Updates metadata DB with DAG structure
"""

# Example DAG file
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

def extract():
    print("Extracting data...")

def transform():
    print("Transforming data...")

with DAG('etl_pipeline', start_date=datetime(2024, 1, 1), schedule='@daily') as dag:
    extract_task = PythonOperator(task_id='extract', python_callable=extract)
    transform_task = PythonOperator(task_id='transform', python_callable=transform)
    extract_task >> transform_task


"""
2. DAG RUN CREATION (Scheduler)
   - Scheduler checks if schedule interval has passed
   - Creates DagRun in metadata DB
   - DagRun represents one execution of the entire DAG
"""

# Scheduler logic (simplified)
if current_time >= (last_run_time + schedule_interval):
    dag_run = DagRun(
        dag_id='etl_pipeline',
        execution_date=logical_date,
        state=State.RUNNING
    )
    save_to_db(dag_run)


"""
3. TASK SCHEDULING (Scheduler)
   - For each DagRun, scheduler creates TaskInstances
   - Checks dependencies (upstream tasks completed?)
   - Queues tasks that are ready
"""

for task in dag.tasks:
    if all(upstream_task.state == State.SUCCESS for upstream_task in task.upstream_tasks):
        task_instance = TaskInstance(
            task=task,
            execution_date=dag_run.execution_date,
            state=State.QUEUED
        )
        executor.queue_task(task_instance)


"""
4. TASK EXECUTION (Executor → Worker)
   - Executor sends task to worker
   - Worker runs task logic
   - Updates task state in DB
"""

# Worker execution
task_instance.state = State.RUNNING
try:
    task_instance.task.execute(context={'execution_date': dag_run.execution_date})
    task_instance.state = State.SUCCESS
except Exception:
    task_instance.state = State.FAILED

"""
5. DAG RUN COMPLETION (Scheduler)
   - When all tasks complete, DagRun marked as SUCCESS or FAILED
   - Triggers cleanup, callbacks, and next schedule
"""

"""
DAG: A collection of tasks with defined dependencies

DIRECTED: Tasks have order (A → B → C)
ACYCLIC: No loops (can't go back to previous task)
GRAPH: Visual representation of workflow
"""

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

# The DAG is the container
with DAG(
    dag_id='my_pipeline',           # Unique identifier
    start_date=datetime(2024, 1, 1), # When DAG becomes active
    schedule='@daily',               # How often to run
    catchup=False                    # Don't backfill historical runs
) as dag:

    # Tasks are the nodes
    task_a = PythonOperator(task_id='task_a', python_callable=lambda: print('A'))
    task_b = PythonOperator(task_id='task_b', python_callable=lambda: print('B'))
    task_c = PythonOperator(task_id='task_c', python_callable=lambda: print('C'))

    # Dependencies define the edges
    task_a >> task_b >> task_c  # A then B then C

    # Equivalent to:
    # task_a.set_downstream(task_b)
    # task_b.set_downstream(task_c)

    ┌─────────┐
    │ Task A  │ (Extract)
    └────┬────┘
         │
         ▼
    ┌─────────┐
    │ Task B  │ (Transform)
    └────┬────┘
         │
         ▼
    ┌─────────┐
    │ Task C  │ (Load)
    └─────────┘

"""
Operators define WHAT each task does
Airflow has 100+ built-in operators for different systems
"""

from airflow.operators.bash import BashOperator
from airflow.operators.python import PythonOperator
from airflow.providers.postgres.operators.postgres import PostgresOperator
from airflow.providers.http.operators.http import SimpleHttpOperator

with DAG('operator_examples', start_date=datetime(2024, 1, 1), schedule=None) as dag:

    # Execute shell commands
    bash_task = BashOperator(
        task_id='run_script',
        bash_command='python /scripts/process.py'
    )

    # Execute Python functions
    def process_data(**context):
        print(f"Processing for {context['ds']}")  # ds = execution date

    python_task = PythonOperator(
        task_id='process',
        python_callable=process_data
    )

    # Execute SQL queries
    sql_task = PostgresOperator(
        task_id='insert_data',
        postgres_conn_id='my_postgres',
        sql="INSERT INTO table VALUES (...)"
    )

    # Call HTTP APIs
    api_task = SimpleHttpOperator(
        task_id='call_api',
        http_conn_id='my_api',
        endpoint='/trigger',
        method='POST'
    )

"""
TaskInstance = Task + Execution Date
Records one execution of a task for a specific DAG run
"""

# Conceptual example
task_instance = TaskInstance(
    task=extract_task,
    execution_date=datetime(2024, 1, 15),
    state='running',  # queued, running, success, failed, up_for_retry
    start_date=datetime(2024, 1, 15, 2, 0, 0),
    end_date=None,
    duration=None,
    try_number=1,
    max_tries=3
)

# Query task instances
from airflow.models import TaskInstance as TI

# Get all failed task instances
failed_tasks = TI.query.filter(
    TI.state == 'failed',
    TI.execution_date >= datetime(2024, 1, 1)
).all()

"""
Complete E-commerce Data Pipeline
Runs daily at 2 AM to process previous day's data
"""

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.postgres.operators.postgres import PostgresOperator
from airflow.providers.amazon.aws.transfers.s3_to_redshift import S3ToRedshiftOperator
from airflow.providers.slack.operators.slack_webhook import SlackWebhookOperator
from datetime import datetime, timedelta

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

with DAG(
    'ecommerce_daily_pipeline',
    default_args=default_args,
    description='Daily ETL for e-commerce analytics',
    schedule='0 2 * * *',  # 2 AM daily
    start_date=datetime(2024, 1, 1),
    catchup=False,
    tags=['production', 'analytics', 'daily']
) as dag:

    # Extract from operational databases
    extract_orders = PostgresOperator(
        task_id='extract_orders',
        postgres_conn_id='prod_db',
        sql="""
            COPY (
                SELECT * FROM orders
                WHERE order_date = '{{ ds }}'  -- Templated execution date
            ) TO '/tmp/orders_{{ ds }}.csv' WITH CSV HEADER;
        """
    )

    extract_customers = PostgresOperator(
        task_id='extract_customers',
        postgres_conn_id='prod_db',
        sql="COPY (SELECT * FROM customers WHERE updated_at >= '{{ ds }}') TO '/tmp/customers_{{ ds }}.csv';"
    )

    # Upload to S3 (data lake)
    def upload_to_s3(**context):
        import boto3
        s3 = boto3.client('s3')
        ds = context['ds']
        s3.upload_file(f'/tmp/orders_{ds}.csv', 'data-lake', f'raw/orders/{ds}/orders.csv')
        s3.upload_file(f'/tmp/customers_{ds}.csv', 'data-lake', f'raw/customers/{ds}/customers.csv')

    upload = PythonOperator(
        task_id='upload_to_s3',
        python_callable=upload_to_s3
    )

    # Transform with dbt (running in separate container)
    transform = BashOperator(
        task_id='transform_dbt',
        bash_command='cd /dbt && dbt run --models tag:daily'
    )

    # Load to Redshift data warehouse
    load_orders = S3ToRedshiftOperator(
        task_id='load_orders_to_redshift',
        s3_bucket='data-lake',
        s3_key='transformed/orders/{{ ds }}/orders.parquet',
        redshift_conn_id='redshift',
        schema='analytics',
        table='fact_orders',
        copy_options=['PARQUET']
    )

    # Data quality checks
    def check_data_quality(**context):
        from airflow.providers.postgres.hooks.postgres import PostgresHook
        hook = PostgresHook(postgres_conn_id='redshift')

        # Check record count
        result = hook.get_first("SELECT COUNT(*) FROM analytics.fact_orders WHERE date = '{{ ds }}'")
        if result[0] < 100:
            raise ValueError(f"Only {result[0]} orders found - expected at least 100")

        # Check for nulls
        null_check = hook.get_first("SELECT COUNT(*) FROM analytics.fact_orders WHERE date = '{{ ds }}' AND customer_id IS NULL")
        if null_check[0] > 0:
            raise ValueError(f"Found {null_check[0]} orders with null customer_id")

    quality_check = PythonOperator(
        task_id='quality_check',
        python_callable=check_data_quality
    )

    # Notify on success
    notify_success = SlackWebhookOperator(
        task_id='notify_success',
        http_conn_id='slack_webhook',
        message=':white_check_mark: Daily pipeline completed for {{ ds }}',
        channel='#data-alerts'
    )

    # Define dependencies
    [extract_orders, extract_customers] >> upload >> transform >> load_orders >> quality_check >> notify_success

┌──────────────────┐  ┌──────────────────┐
│ Extract Orders   │  │ Extract Customers│
└────────┬─────────┘  └────────┬─────────┘
         │                     │
         └──────────┬──────────┘
                    ▼
            ┌──────────────┐
            │ Upload to S3 │
            └──────┬───────┘
                   ▼
            ┌──────────────┐
            │ Transform dbt│
            └──────┬───────┘
                   ▼
         ┌─────────────────┐
         │ Load to Redshift│
         └────────┬────────┘
                  ▼
         ┌─────────────────┐
         │ Quality Check   │
         └────────┬────────┘
                  ▼
         ┌─────────────────┐
         │ Notify Success  │
         └─────────────────┘

# Install Airflow (Python 3.8+)
pip install "apache-airflow==2.8.1" \
  --constraint "https://raw.githubusercontent.com/apache/airflow/constraints-2.8.1/constraints-3.11.txt"

# Initialize database
airflow db init

# Create admin user
airflow users create \
  --username admin \
  --firstname Admin \
  --lastname User \
  --role Admin \
  --email admin@example.com

# Start web server (port 8080)
airflow webserver

# Start scheduler (in separate terminal)
airflow scheduler