yt-dlp-dags/dags/ytdlp_mgmt_queue_check_status.py

# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2024 rl <rl@rlmbp>
#
# Distributed under terms of the MIT license.

"""
Airflow DAG for manually checking the status (type and size) of a specific Redis key used by YTDLP queues.
"""

from airflow import DAG
from airflow.exceptions import AirflowException
from airflow.models.param import Param
from airflow.operators.python import PythonOperator
from airflow.providers.redis.hooks.redis import RedisHook
from airflow.utils.dates import days_ago
from datetime import timedelta
import logging
import redis # Import redis exceptions if needed

# Configure logging
logger = logging.getLogger(__name__)

# Default settings
DEFAULT_REDIS_CONN_ID = 'redis_default'
# Default to a common inbox pattern, user should override with the specific key
DEFAULT_QUEUE_TO_CHECK = 'video_queue_inbox'

# --- Helper Function ---

def _get_redis_client(redis_conn_id):
    """Gets a Redis client connection using RedisHook."""
    try:
        hook = RedisHook(redis_conn_id=redis_conn_id)
        client = hook.get_conn()
        client.ping()
        logger.info(f"Successfully connected to Redis using connection '{redis_conn_id}'.")
        return client
    except redis.exceptions.AuthenticationError:
        logger.error(f"Redis authentication failed for connection '{redis_conn_id}'. Check password.")
        raise AirflowException(f"Redis authentication failed for '{redis_conn_id}'.")
    except Exception as e:
        logger.error(f"Failed to get Redis client for connection '{redis_conn_id}': {e}")
        raise AirflowException(f"Redis connection failed for '{redis_conn_id}': {e}")

# --- Python Callable for Check Status Task ---

def check_status_callable(**context):
    """Checks the length/size of the specified Redis key (queue/hash)."""
    params = context['params']
    redis_conn_id = params['redis_conn_id']
    queue_to_check = params['queue_to_check'] # Specific queue/hash name

    if not queue_to_check:
        raise ValueError("Parameter 'queue_to_check' cannot be empty.")

    logger.info(f"Attempting to check status of Redis key '{queue_to_check}' using connection '{redis_conn_id}'.")
    try:
        redis_client = _get_redis_client(redis_conn_id)
        key_type = redis_client.type(queue_to_check)
        key_type_str = key_type.decode('utf-8') if isinstance(key_type, bytes) else key_type # Decode if needed

        length = 0
        if key_type_str == 'list':
            length = redis_client.llen(queue_to_check)
            logger.info(f"Redis list '{queue_to_check}' has {length} items.")
        elif key_type_str == 'hash':
            length = redis_client.hlen(queue_to_check)
            logger.info(f"Redis hash '{queue_to_check}' has {length} fields.")
        elif key_type_str == 'none':
            logger.info(f"Redis key '{queue_to_check}' does not exist.")
        else:
            # Attempt to get size for other types if possible, e.g., set size
            try:
                length = redis_client.scard(queue_to_check) # Example for set
                logger.info(f"Redis key '{queue_to_check}' (type: {key_type_str}) has size {length}.")
            except:
                 logger.info(f"Redis key '{queue_to_check}' exists but is of unhandled/unsizeable type '{key_type_str}'.")

        # Optionally push length to XCom if needed downstream
        context['task_instance'].xcom_push(key='queue_key_type', value=key_type_str)
        context['task_instance'].xcom_push(key='queue_size', value=length)
        return {'key': queue_to_check, 'type': key_type_str, 'size': length} # Return status info
    except Exception as e:
        logger.error(f"Failed to check status of Redis key '{queue_to_check}': {e}", exc_info=True)
        raise AirflowException(f"Failed to check Redis key status: {e}")

# --- DAG Definition ---
default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(seconds=30),
    'start_date': days_ago(1)
}

with DAG(
    dag_id='ytdlp_mgmt_queue_check_status',
    default_args=default_args,
    schedule_interval=None, # Manually triggered
    catchup=False,
    description='Manually check the type and size of a specific YTDLP Redis queue/key.',
    tags=['ytdlp', 'queue', 'management', 'redis', 'manual', 'status'],
    params={
        'redis_conn_id': Param(DEFAULT_REDIS_CONN_ID, type="string", description="Airflow Redis connection ID."),
        'queue_to_check': Param(
            DEFAULT_QUEUE_TO_CHECK,
            type="string",
            description="Exact name of the Redis key to check (e.g., 'video_queue_inbox_account_xyz', 'video_queue_progress', 'video_queue_result', 'video_queue_fail')."
        ),
    }
) as dag:

    check_status_task = PythonOperator(
        task_id='check_specified_queue_status',
        python_callable=check_status_callable,
        # Params are implicitly passed via context['params']
    )
    check_status_task.doc_md = """
    ### Check Specified Queue/Key Status Task
    Checks the type and size (length for lists, number of fields for hashes) of the Redis key specified by `queue_to_check`.
    Logs the result and pushes `queue_key_type` and `queue_size` to XCom.
    Can check keys like:
    - `_inbox` (Redis List)
    - `_progress` (Redis Hash)
    - `_result` (Redis Hash)
    - `_fail` (Redis Hash)

    *Trigger this task manually via the UI.*
    """