Airflow

# dags/etl_pipeline.py
# 일일 ETL 파이프라인 DAG 예제

from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.operators.bash import BashOperator
from airflow.providers.google.cloud.operators.bigquery import (
    BigQueryCreateEmptyTableOperator,
    BigQueryInsertJobOperator
)
from airflow.providers.google.cloud.transfers.gcs_to_bigquery import (
    GCSToBigQueryOperator
)
from airflow.utils.task_group import TaskGroup
from airflow.models import Variable

# DAG 기본 설정
default_args = {
    'owner': 'data-team',
    'depends_on_past': False,
    'email': ['data-alerts@company.com'],
    'email_on_failure': True,
    'email_on_retry': False,
    'retries': 3,
    'retry_delay': timedelta(minutes=5),
    'execution_timeout': timedelta(hours=2),
}

# DAG 정의
with DAG(
    dag_id='daily_sales_etl',
    default_args=default_args,
    description='일일 매출 데이터 ETL 파이프라인',
    schedule_interval='0 2 * * *',  # 매일 오전 2시
    start_date=datetime(2024, 1, 1),
    catchup=False,
    tags=['etl', 'sales', 'production'],
    max_active_runs=1,
) as dag:

    # 1. 데이터 추출 태스크
    def extract_sales_data(**context):
        """소스 시스템에서 매출 데이터 추출"""
        import pandas as pd
        from google.cloud import storage

        execution_date = context['ds']

        # API에서 데이터 추출 (예시)
        # sales_data = fetch_from_api(execution_date)

        # GCS에 저장
        bucket_name = Variable.get('gcs_bucket')
        blob_path = f'raw/sales/{execution_date}/data.parquet'

        print(f"Extracted sales data for {execution_date}")
        return blob_path

    extract_task = PythonOperator(
        task_id='extract_sales_data',
        python_callable=extract_sales_data,
        provide_context=True,
    )

    # 2. 데이터 검증 태스크 그룹
    with TaskGroup(group_id='data_validation') as validation_group:

        def validate_row_count(**context):
            """행 수 검증"""
            ti = context['ti']
            blob_path = ti.xcom_pull(task_ids='extract_sales_data')
            # 실제 검증 로직
            print(f"Validating row count for {blob_path}")

        def validate_schema(**context):
            """스키마 검증"""
            ti = context['ti']
            blob_path = ti.xcom_pull(task_ids='extract_sales_data')
            print(f"Validating schema for {blob_path}")

        row_count_check = PythonOperator(
            task_id='check_row_count',
            python_callable=validate_row_count,
        )

        schema_check = PythonOperator(
            task_id='check_schema',
            python_callable=validate_schema,
        )

        [row_count_check, schema_check]

    # 3. BigQuery 로드 태스크
    load_to_bigquery = GCSToBigQueryOperator(
        task_id='load_to_bigquery',
        bucket='{{ var.value.gcs_bucket }}',
        source_objects=['raw/sales/{{ ds }}/data.parquet'],
        destination_project_dataset_table='analytics.raw_sales',
        source_format='PARQUET',
        write_disposition='WRITE_APPEND',
        create_disposition='CREATE_IF_NEEDED',
    )

    # 4. 변환 쿼리 실행
    transform_query = """
    INSERT INTO `analytics.sales_daily_summary`
    SELECT
        DATE('{{ ds }}') as report_date,
        product_category,
        SUM(amount) as total_sales,
        COUNT(DISTINCT customer_id) as unique_customers,
        COUNT(*) as transaction_count
    FROM `analytics.raw_sales`
    WHERE DATE(created_at) = '{{ ds }}'
    GROUP BY product_category
    """

    transform_task = BigQueryInsertJobOperator(
        task_id='transform_data',
        configuration={
            'query': {
                'query': transform_query,
                'useLegacySql': False,
            }
        },
    )

    # 5. dbt 실행 (선택적)
    run_dbt = BashOperator(
        task_id='run_dbt_models',
        bash_command='cd /opt/dbt && dbt run --models sales_marts --target prod',
    )

    # 6. 완료 알림
    def send_completion_notification(**context):
        """Slack 알림 전송"""
        from airflow.providers.slack.hooks.slack_webhook import SlackWebhookHook

        slack_hook = SlackWebhookHook(
            slack_webhook_conn_id='slack_webhook'
        )
        slack_hook.send(
            text=f"Daily Sales ETL completed for {context['ds']}"
        )

    notify_task = PythonOperator(
        task_id='send_notification',
        python_callable=send_completion_notification,
        trigger_rule='all_success',
    )

    # 태스크 의존관계 정의
    extract_task >> validation_group >> load_to_bigquery >> transform_task >> run_dbt >> notify_task