# Version: 2025-08-20-02
# This file contains custom hooks for the Airflow environment.
from airflow import settings
from airflow.models.dagrun import DagRun
from airflow.utils.session import provide_session


@provide_session
def task_instance_mutation_hook(ti, session=None):
    if ti.dag_id == 'ytdlp_ops_worker_per_url':
        # Query the DagRun from the DB using run_id to reliably get the conf.
        # The ti.dag_run attribute is not always populated when the hook is called.
        dag_run = session.query(DagRun).filter(DagRun.run_id == ti.run_id).first()
        conf = dag_run.conf if dag_run else {}
        worker_queue = conf.get('worker_queue')

        if worker_queue:
            print(f"MUTATION HOOK: For dag '{ti.dag_id}', pinning task '{ti.task_id}' (run_id: {ti.run_id}) to queue '{worker_queue}'.")
            ti.queue = worker_queue
        else:
            print(f"MUTATION HOOK: For dag '{ti.dag_id}', no 'worker_queue' in conf for run_id '{ti.run_id}'. Falling back to 'queue-dl'.")
            ti.queue = 'queue-dl'

# Register the hook only in appropriate contexts
# This hook can cause issues with the Triggerer, which does not have a `dag_run` context
# when it runs its own maintenance tasks.
if not settings.CONFIG.get('core', 'executor').lower().startswith('debug'):
    settings.task_instance_mutation_hook = task_instance_mutation_hook