Cloud Composer で Storage Transfer Service の動的設定更新ワークフローを構築する

from datetime import datetime, timedelta
from airflow import DAG
from airflow.providers.google.cloud.operators.cloud_storage_transfer_service import (
    CloudDataTransferServiceRunJobOperator,
    CloudDataTransferServiceUpdateJobOperator,
    CloudDataTransferServiceListOperationsOperator
)
from airflow.operators.python import PythonOperator
from airflow.utils.dates import days_ago
import sys
import os

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

# 設定値（注：実際の値に置き換える）
PROJECT_ID = '{プロジェクトID}'
TRANSFER_JOB_ID = 'transferJobs/{転送ジョブID}'

# 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='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

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

    # 転送履歴の取得
    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='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='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
    )

$ gcloud transfer jobs describe transferJobs/{転送ジョブID} | jq '.transferSpec.objectConditions.lastModifiedSince'
null

{
    "projectId": "{プロジェクトID}",
    "transferJob": {
        "transferSpec": {
            "gcsDataSink": {
                "bucketName": "cm_enokawa_work"
            },
            "gcsDataSource": {
                "bucketName": "cm_enokawa_work_source"
            },
            "objectConditions": {
                "lastModifiedSince": "2025-09-01T05:57:18.622936746Z"
            },
            "transferOptions": {
                "metadataOptions": {
                    "acl": "ACL_DESTINATION_BUCKET_DEFAULT",
                    "kmsKey": "KMS_KEY_DESTINATION_BUCKET_DEFAULT",
                    "storageClass": "STORAGE_CLASS_DESTINATION_BUCKET_DEFAULT",
                    "temporaryHold": "TEMPORARY_HOLD_PRESERVE",
                    "timeCreated": "TIME_CREATED_SKIP"
                },
                "overwriteWhen": "DIFFERENT"
            }
        }
    }
}

$ gcloud transfer jobs describe transferJobs/{転送ジョブID} | jq '.transferSpec.objectConditions.lastModifiedSince'
"2025-09-01T05:57:18.622936746Z"