# mwaa_environment.tf (Terraform)
resource "aws_mwaa_environment" "airflow" {
  name               = "production-airflow"
  airflow_version    = "2.8.1"
  environment_class  = "mw1.medium"

  source_bucket_arn  = aws_s3_bucket.dags.arn
  dag_s3_path        = "dags/"

  execution_role_arn = aws_iam_role.mwaa_role.arn

  network_configuration {
    security_group_ids = [aws_security_group.mwaa.id]
    subnet_ids         = var.private_subnet_ids
  }

  logging_configuration {
    dag_processing_logs {
      enabled   = true
      log_level = "INFO"
    }
    scheduler_logs {
      enabled   = true
      log_level = "INFO"
    }
    task_logs {
      enabled   = true
      log_level = "INFO"
    }
    webserver_logs {
      enabled   = true
      log_level = "WARNING"
    }
    worker_logs {
      enabled   = true
      log_level = "INFO"
    }
  }

  webserver_access_mode = "PUBLIC_ONLY"

  max_workers = 10
  min_workers = 2

  environment_variables = {
    AIRFLOW__CORE__LOAD_EXAMPLES = "False"
    AIRFLOW__WEBSERVER__EXPOSE_CONFIG = "True"
    AIRFLOW__DATABASE__SQL_ALCHEMY_CONN = "postgresql+psycopg2://..."
  }

  webserver_url = aws_mwaa_environment.airflow.webserver_url
  arn           = aws_mwaa_environment.airflow.arn
}

# composer_environment.tf (Terraform)
resource "google_composer_environment" "airflow" {
  name   = "production-airflow"
  region = "us-central1"

  config {
    node_count = 4

    software_config {
      image_version = "composer-2.6.1-airflow-2.8.1"
      pypi_packages = {
        "apache-airflow-providers-google" = ">=10.0.0"
        "apache-airflow-providers-amazon" = ">=8.0.0"
        "pandas" = ">=2.0.0"
      }
      env_variables = {
        AIRFLOW__CORE__LOAD_EXAMPLES = "False"
      }
    }

    workloads_config {
      scheduler {
        cpu        = 2
        memory_gb  = 4
        storage_gb = 10
        count      = 2
      }
      worker {
        cpu        = 2
        memory_gb  = 8
        storage_gb = 20
        min_count  = 2
        max_count  = 10
      }
      triggerer {
        cpu        = 0.5
        memory_gb  = 1
        count      = 2
      }
    }

    database_config {
      machine_type = "db-n1-standard-2"
    }

    web_server_config {
      machine_type = "composer-web-server-medium"
    }

    environment_size = "ENVIRONMENT_SIZE_MEDIUM"

    private_environment_config {
      enable_private_endpoint = true
      master_ipv4_cidr_block  = "172.16.0.0/28"
    }
  }
}

# dags/mwaa_example_dag.py
from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.providers.amazon.aws.transfers.s3_to_redshift import S3ToRedshiftOperator


def extract_to_s3(**context):
    """Extract data and upload to S3 DAG bucket."""
    import pandas as pd
    import io

    data = pd.DataFrame({
        'date': [context['ds']],
        'record_count': [1000],
        'status': ['completed'],
    })

    csv_buffer = io.StringIO()
    data.to_csv(csv_buffer, index=False)

    hook = S3Hook(aws_conn_id='aws_default')
    hook.load_string(
        string_data=csv_buffer.getvalue(),
        key=f'extracted/{context["ds"]}/data.csv',
        bucket_name='data-lake',
        replace=True,
    )


with DAG(
    dag_id='mwaa_managed_pipeline',
    start_date=datetime(2024, 1, 1),
    schedule_interval='@daily',
    catchup=False,
    default_args={
        'retries': 2,
        'retry_delay': timedelta(minutes=5),
    },
    tags=['mwaa', 'managed'],
) as dag:

    extract = PythonOperator(
        task_id='extract',
        python_callable=extract_to_s3,
    )

    load = S3ToRedshiftOperator(
        task_id='load_to_redshift',
        schema='public',
        table='daily_metrics',
        s3_bucket='data-lake',
        s3_key='extracted/{{ ds }}/data.csv',
        copy_options=['FORMAT AS CSV', 'IGNOREHEADER 1'],
        aws_conn_id='aws_default',
        redshift_conn_id='redshift_default',
    )

    extract >> load