Cloud Composer で Storage Transfer Service 転送後のGCSファイル検証と自動リトライを実装する

			
			from datetime import timedelta, datetime
from airflow import DAG
from airflow.providers.google.cloud.operators.cloud_storage_transfer_service import (
    CloudDataTransferServiceRunJobOperator,
    CloudDataTransferServiceUpdateJobOperator,
    CloudDataTransferServiceListOperationsOperator
)
from airflow.providers.google.cloud.operators.gcs import GCSListObjectsOperator
from airflow.operators.python import PythonOperator, BranchPythonOperator
from airflow.operators.trigger_dagrun import TriggerDagRunOperator
from airflow.operators.dummy import DummyOperator
from airflow.utils.task_group import TaskGroup
from airflow.utils.dates import days_ago
from airflow.exceptions import AirflowException

# カスタムセンサーをインポート
from custom_sts_sensor import CustomStorageTransferOperationSensor

# 設定値（注：実際の値に置き換える）
PROJECT_ID = '{プロジェクトID}'
TRANSFER_JOB_ID = 'transferJobs/{転送ジョブID}'
GCS_TARGET_BUCKET = '{GCSバケット名}'
GCS_TARGET_PREFIX = 'data/processed/'
TARGET_FILE_EXTENSIONS = ['.parquet', '.csv']
MAX_RETRY_COUNT = 3

# DAGの基本設定
default_args = {
    'start_date': days_ago(1),
    'retries': 0,  # リトライなし
}

def prepare_update_body(**context):
    """転送履歴から最終成功開始時刻を抽出しジョブ更新用設定を生成"""
    from airflow.providers.google.cloud.hooks.cloud_storage_transfer_service import (
        CloudDataTransferServiceHook
    )

    # 前タスクから転送履歴を取得
    operations = context['task_instance'].xcom_pull(task_ids='sts_transfer_group.list_operations')

    # フォールバック時間（履歴が無い場合は24時間前を使用）
    fallback_time = (datetime.utcnow() - timedelta(hours=24)).isoformat() + 'Z'

    # 最終成功時刻を抽出
    if not operations:
        last_success_time = fallback_time
    else:
        # 成功したオペレーションのみを抽出
        success_operations = [
            op for op in operations
            if op.get('metadata', {}).get('status') == 'SUCCESS'
        ]

        if not success_operations:
            last_success_time = fallback_time
        else:
            # startTimeでソートして最新の成功開始時刻を取得
            success_operations.sort(
                key=lambda x: x.get('metadata', {}).get('startTime', ''),
                reverse=True
            )
            last_success_time = success_operations[0].get('metadata', {}).get('startTime')

    # 既存の転送ジョブ設定を取得
    hook = CloudDataTransferServiceHook(gcp_conn_id='google_cloud_default')
    transfer_job = hook.get_transfer_job(
        job_name=TRANSFER_JOB_ID,
        project_id=PROJECT_ID
    )

    # 更新用転送ジョブ設定を作成（変更するフィールドのみ）
    updated_transfer_job = {
        'transferSpec': {
            **transfer_job.get('transferSpec', {}),
            'objectConditions': {
                **transfer_job.get('transferSpec', {}).get('objectConditions', {}),
                'lastModifiedSince': last_success_time
            }
        }
    }

    # 更新用リクエストボディを作成
    update_body = {
        "projectId": PROJECT_ID,
        "transferJob": updated_transfer_job
    }

    return update_body

def check_gcs_files_and_branch(**context):
    """GCSファイル存在チェック + リトライ分岐判定"""

    # GCSファイル一覧を取得
    file_list = context['task_instance'].xcom_pull(
        task_ids='gcs_verification_group.list_gcs_files'
    ) or []

    # 特定拡張子のファイルをフィルタリング
    target_files = [
        f for f in file_list
        if any(f.endswith(ext) for ext in TARGET_FILE_EXTENSIONS)
    ]

    # DAG実行時のconfからリトライ回数を取得
    dag_run_conf = context['dag_run'].conf or {}
    retry_count = dag_run_conf.get('retry_count', 0)

    if target_files:
        # ファイル検証成功
        return 'gcs_verification_group.success_end'
    else:
        # ファイルが見つからない場合
        if retry_count >= MAX_RETRY_COUNT:
            raise AirflowException(
                f"GCSファイル検証失敗: 対象ファイルが見つかりません (リトライ上限到達: {retry_count}/{MAX_RETRY_COUNT})"
            )
        else:
            # リトライ回数を保存
            context['task_instance'].xcom_push(key='retry_count', value=retry_count + 1)
            return 'gcs_verification_group.prepare_retry_conf'

def prepare_retry_conf(**context):
    """DAG再実行時の設定を準備"""
    import time

    # 前タスクからリトライ回数を取得
    retry_count = context['task_instance'].xcom_pull(
        task_ids='gcs_verification_group.check_files_and_branch',
        key='retry_count'
    )

    # リトライ回数に応じた待機時間（1回目: 60秒、2回目: 120秒、3回目: 180秒）
    wait_seconds = retry_count * 60
    time.sleep(wait_seconds)

    # confとして渡す設定
    conf = {
        'retry_count': retry_count,
        'previous_run_id': context['dag_run'].run_id,
        'triggered_at': datetime.utcnow().isoformat()
    }

    return conf

# DAG定義
with DAG(
    dag_id='sts_gcs_verification_with_retry',
    default_args=default_args,
    schedule_interval=None,  # 手動トリガー
    catchup=False,
    render_template_as_native_obj=True,  # XComデータを適切に処理
) as dag:

    # STS転送グループ
    with TaskGroup('sts_transfer_group') as sts_group:

        # 転送履歴の取得
        list_operations = CloudDataTransferServiceListOperationsOperator(
            task_id='list_operations',
            project_id=PROJECT_ID,
            request_filter={'job_names': [TRANSFER_JOB_ID]}
        )

        # 更新設定の準備
        prepare_update = PythonOperator(
            task_id='prepare_update_body',
            python_callable=prepare_update_body,
            provide_context=True
        )

        # 転送ジョブ設定の更新
        update_transfer_job = CloudDataTransferServiceUpdateJobOperator(
            task_id='update_transfer_job',
            job_name=TRANSFER_JOB_ID,
            body="{{ task_instance.xcom_pull(task_ids='sts_transfer_group.prepare_update_body') }}",
            project_id=PROJECT_ID
        )

        # 転送ジョブの実行
        run_transfer = CloudDataTransferServiceRunJobOperator(
            task_id='run_transfer',
            job_name=TRANSFER_JOB_ID,
            project_id=PROJECT_ID
        )

        # 転送完了待機
        wait_for_transfer = CustomStorageTransferOperationSensor(
            task_id='wait_for_transfer',
            operation_name="{{ task_instance.xcom_pull(task_ids='sts_transfer_group.run_transfer')['name'] }}",
            project_id=PROJECT_ID,
            poke_interval=60,        # 60秒ごとにステータスを確認
            timeout=60 * 10,         # 最大10分間待機
            mode='reschedule'        # リソース効率のためrescheduleモードを使用
        )

        # タスク依存関係
        list_operations >> prepare_update >> update_transfer_job >> run_transfer >> wait_for_transfer

    # GCS検証グループ
    with TaskGroup('gcs_verification_group') as verify_group:

        # GCSファイル一覧を取得
        list_gcs_files = GCSListObjectsOperator(
            task_id='list_gcs_files',
            bucket=GCS_TARGET_BUCKET,
            prefix=GCS_TARGET_PREFIX
        )

        # ファイル存在チェック + 分岐判定
        check_files = BranchPythonOperator(
            task_id='check_files_and_branch',
            python_callable=check_gcs_files_and_branch
        )

        # 成功終了
        success_end = DummyOperator(
            task_id='success_end',
            trigger_rule='none_failed_min_one_success'
        )

        # リトライ設定準備
        prepare_retry = PythonOperator(
            task_id='prepare_retry_conf',
            python_callable=prepare_retry_conf,
            trigger_rule='none_failed_min_one_success'
        )

        # DAG再実行
        retry_transfer = TriggerDagRunOperator(
            task_id='retry_transfer',
            trigger_dag_id=dag.dag_id,
            conf="{{ task_instance.xcom_pull(task_ids='gcs_verification_group.prepare_retry_conf') }}",
            wait_for_completion=False,
            reset_dag_run=True,
            trigger_rule='none_failed_min_one_success'
        )

        # タスク依存関係
        list_gcs_files >> check_files >> [success_end, prepare_retry]
        prepare_retry >> retry_transfer

    # DAG全体のフロー
    sts_group >> verify_group

		

			
			# ❌ TaskGroup名を省略した場合、タスクが見つからずエラーになる
return 'success_end'

# ✅ TaskGroup名を含む完全修飾名で指定する必要がある
return 'gcs_verification_group.success_end'

		

			
			# リトライ回数の取得
dag_run_conf = context['dag_run'].conf or {}
retry_count = dag_run_conf.get('retry_count', 0)

# 次回実行への引き継ぎ
retry_transfer = TriggerDagRunOperator(
    trigger_dag_id=dag.dag_id,
    conf={'retry_count': retry_count + 1}
)

		

			
			実行1: manual__2025-10-13T09:21:36+00:00 (retry_count=0)
├── sts_transfer_group ✅
├── gcs_verification_group
│   └── retry_transfer ✅ → 新DAG実行をトリガー

実行2: retry_2025101309xxxx (retry_count=1) ← 新しい実行
├── sts_transfer_group ✅
├── gcs_verification_group
│   └── success_end ✅

		

			
			# 成功終了
success_end = DummyOperator(
    task_id='success_end',
    trigger_rule='none_failed_min_one_success'  # ← 必須
)

# リトライ設定準備
prepare_retry = PythonOperator(
    task_id='prepare_retry_conf',
    python_callable=prepare_retry_conf,
    trigger_rule='none_failed_min_one_success'  # ← 必須
)

		

			
			# リトライ回数に応じた待機時間（1回目: 60秒、2回目: 120秒、3回目: 180秒）
wait_seconds = retry_count * 60
time.sleep(wait_seconds)