109 lines
4.4 KiB
Python
109 lines
4.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
# vim:fenc=utf-8
|
|
"""
|
|
Airflow DAG for manually clearing (deleting) 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'
|
|
# Provide a placeholder default, user MUST specify the queue to clear
|
|
DEFAULT_QUEUE_TO_CLEAR = 'PLEASE_SPECIFY_QUEUE_TO_CLEAR'
|
|
|
|
# --- 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 Clear Task ---
|
|
|
|
def clear_queue_callable(**context):
|
|
"""Clears (deletes) the specified Redis key (queue/hash)."""
|
|
params = context['params']
|
|
redis_conn_id = params['redis_conn_id']
|
|
queue_to_clear = params['queue_to_clear'] # Specific queue/hash name
|
|
|
|
if not queue_to_clear or queue_to_clear == DEFAULT_QUEUE_TO_CLEAR:
|
|
raise ValueError("Parameter 'queue_to_clear' must be specified and cannot be the default placeholder.")
|
|
|
|
logger.info(f"Attempting to clear Redis key '{queue_to_clear}' using connection '{redis_conn_id}'.")
|
|
try:
|
|
redis_client = _get_redis_client(redis_conn_id)
|
|
deleted_count = redis_client.delete(queue_to_clear)
|
|
if deleted_count > 0:
|
|
logger.info(f"Successfully cleared Redis key '{queue_to_clear}'.")
|
|
else:
|
|
logger.info(f"Redis key '{queue_to_clear}' did not exist or was already empty.")
|
|
except Exception as e:
|
|
logger.error(f"Failed to clear Redis key '{queue_to_clear}': {e}", exc_info=True)
|
|
raise AirflowException(f"Failed to clear Redis key: {e}")
|
|
|
|
# --- DAG Definition ---
|
|
default_args = {
|
|
'owner': 'airflow',
|
|
'depends_on_past': False,
|
|
'email_on_failure': False,
|
|
'email_on_retry': False,
|
|
'retries': 0, # No retries for manual clear operation
|
|
'start_date': days_ago(1)
|
|
}
|
|
|
|
with DAG(
|
|
dag_id='ytdlp_mgmt_queue_clear',
|
|
default_args=default_args,
|
|
schedule_interval=None, # Manually triggered
|
|
catchup=False,
|
|
description='Manually clear/delete a specific YTDLP Redis queue/key (inbox, progress, result, fail). Use with caution!',
|
|
tags=['ytdlp', 'queue', 'management', 'redis', 'manual', 'clear'],
|
|
params={
|
|
'redis_conn_id': Param(DEFAULT_REDIS_CONN_ID, type="string", description="Airflow Redis connection ID."),
|
|
'queue_to_clear': Param(
|
|
DEFAULT_QUEUE_TO_CLEAR,
|
|
type="string",
|
|
description="Exact name of the Redis key to clear (e.g., 'video_queue_inbox_account_xyz', 'video_queue_progress', 'video_queue_result', 'video_queue_fail')."
|
|
),
|
|
}
|
|
) as dag:
|
|
|
|
clear_queue_task = PythonOperator(
|
|
task_id='clear_specified_queue',
|
|
python_callable=clear_queue_callable,
|
|
# Params are implicitly passed via context['params']
|
|
)
|
|
clear_queue_task.doc_md = """
|
|
### Clear Specified Queue/Key Task
|
|
Deletes the Redis key specified by the `queue_to_clear` parameter.
|
|
This can target any key, including:
|
|
- `_inbox` (Redis List): Contains URLs waiting to be processed.
|
|
- `_progress` (Redis Hash): Contains URLs currently being processed.
|
|
- `_result` (Redis Hash): Contains details of successfully processed URLs.
|
|
- `_fail` (Redis Hash): Contains details of failed URLs.
|
|
|
|
**Warning:** This operation is destructive and cannot be undone. Ensure you specify the correct key name.
|
|
*Trigger this task manually via the UI.*
|
|
"""
|