def execute(self, context: Context) -> None:
    if not self.deferrable:
        super().execute(context)   # mode="poke"/"reschedule" 只在這裡起作用
    else:
        if not self.poke(context=context):
            self._defer()          # 交給 triggerer 背景等待，不佔用 worker

import logging
from datetime import UTC, datetime, timedelta

from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor
from airflow.sdk import DAG, Param, Variable, get_current_context, task
from airflow.sdk.definitions.param import ParamsDict
from services.aws.s3_client import S3Client

logger = logging.getLogger(__name__)

default_args = {
    "owner": "Developer",
    "depends_on_past": False,
    "start_date": datetime(2026, 9, 22, tzinfo=UTC),
    "retry_delay": timedelta(minutes=1),
}

with DAG(
    dag_id="s3_key_sensor_test",
    default_args=default_args,
    schedule=None,
    catchup=False,
    max_active_runs=1,
    dag_display_name="S3KeySensor 概念驗證",
    tags=["AWS"],
    doc_md="""\
        ## 測試 `S3KeySensor` 用法

        1. `get_s3_path`：產生要監控的 S3 路徑
        2. `wait_for_key`：等待 `get_s3_path` 產生的 S3 object 出現才會往下走
        3. `fetch_value`：偵測到檔案後用既有的 `S3Client` 讀取內容，回傳長度 (bytes)
        4. `print_value`：接收 `fetch_value` 的回傳值並印出來
    """,
    params=ParamsDict(
        {
            "bucket_name": Param(
                default="",
                type="string",
                title="S3 Bucket",
                description="S3 bucket name",
            ),
            "bucket_key": Param(
                default="",
                type="string",
                title="Bucket Key",
                description="S3 object key",
            ),
        }
    ),
) as dag:

    @task(task_id="get_s3_path", multiple_outputs=True)
    def get_s3_path() -> dict[str, str]:
        """模擬正式流程中路徑是由程式邏輯算出來的，這裡先直接用 Trigger DAG 帶入的 params 代替

        multiple_outputs=True 讓每個 dict key 分別存成獨立的 XCom 鍵值，下游才能用 s3_path["bucket_name"] 這種下標語法取值，否則會拿到 None

        Returns:
            dict[str, str]: 各自對應一個 XCom key
        """
        params = get_current_context().get("params", {})
        return {
            "bucket_name": params["bucket_name"],
            "bucket_key": params["bucket_key"],
        }

    s3_path = get_s3_path()

    wait_for_object = S3KeySensor(
        task_id="wait_for_object",
        bucket_name=s3_path["bucket_name"],  # type: ignore[index]
        bucket_key=s3_path["bucket_key"],  # type: ignore[index]
        deferrable=True,  # 是否交給 triggerer 執行，不佔用 worker slot
        poke_interval=30,  # 每次檢查間隔秒數，正式建議 >= 60
        timeout=3600,  # 逾時秒數
    )

    @task(task_id="fetch_object", retries=0)
    def fetch_object(bucket_name: str, bucket_key: str) -> int:
        content = S3Client(bucket_name=bucket_name).get_object(bucket_key)

        logger.info("[S3KeySensor] 在 S3 %s/%s 偵測到檔案", bucket_name, bucket_key)
        return len(content)

    @task(task_id="print_object_length", retries=0)
    def print_value(content_length: int) -> None:
        print(f"[S3KeySensor] 成功取得 {content_length} bytes")

    fetch_object_task = fetch_object(
        s3_path["bucket_name"],  # type: ignore[index]
        s3_path["bucket_key"],  # type: ignore[index]
    )

    wait_for_object.set_downstream(fetch_object_task)
    print_value(fetch_object_task)  # type: ignore[arg-type]

if __name__ == "__main__":
    dag.cli()