145 lines
5.2 KiB
Python
145 lines
5.2 KiB
Python
# -*- coding: utf-8 -*-
|
|
#
|
|
# Copyright © 2024 rl
|
|
#
|
|
# Distributed under terms of the MIT license.
|
|
|
|
"""
|
|
Maintenance DAG for managing the lifecycle of ytdlp-ops accounts.
|
|
This DAG is responsible for:
|
|
- Un-banning accounts whose ban duration has expired.
|
|
- Transitioning accounts from RESTING to ACTIVE after their cooldown period.
|
|
- Transitioning accounts from ACTIVE to RESTING after their active duration.
|
|
This logic was previously handled inside the ytdlp-ops-server and has been
|
|
moved here to give the orchestrator full control over account state.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import time
|
|
from datetime import datetime
|
|
|
|
from airflow.decorators import task
|
|
from airflow.models import Variable
|
|
from airflow.models.dag import DAG
|
|
from airflow.utils.dates import days_ago
|
|
|
|
# Import utility functions and Thrift modules
|
|
from utils.redis_utils import _get_redis_client
|
|
from pangramia.yt.tokens_ops import YTTokenOpService
|
|
from thrift.protocol import TBinaryProtocol
|
|
from thrift.transport import TSocket, TTransport
|
|
|
|
# Configure logging
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Default settings from Airflow Variables or hardcoded fallbacks
|
|
DEFAULT_REDIS_CONN_ID = 'redis_default'
|
|
DEFAULT_YT_AUTH_SERVICE_IP = Variable.get("YT_AUTH_SERVICE_IP", default_var="172.17.0.1")
|
|
DEFAULT_YT_AUTH_SERVICE_PORT = Variable.get("YT_AUTH_SERVICE_PORT", default_var=9080)
|
|
|
|
DEFAULT_ARGS = {
|
|
'owner': 'airflow',
|
|
'retries': 1,
|
|
'retry_delay': 30,
|
|
'queue': 'maintenance',
|
|
}
|
|
|
|
|
|
# --- Helper Functions ---
|
|
|
|
def _get_thrift_client(host, port, timeout=60):
|
|
"""Helper to create and connect a Thrift client."""
|
|
transport = TSocket.TSocket(host, port)
|
|
transport.setTimeout(timeout * 1000)
|
|
transport = TTransport.TFramedTransport(transport)
|
|
protocol = TBinaryProtocol.TBinaryProtocol(transport)
|
|
client = YTTokenOpService.Client(protocol)
|
|
transport.open()
|
|
logger.info(f"Connected to Thrift server at {host}:{port}")
|
|
return client, transport
|
|
|
|
|
|
@task
|
|
def manage_account_states():
|
|
"""
|
|
Fetches all account statuses and performs necessary state transitions.
|
|
"""
|
|
host = DEFAULT_YT_AUTH_SERVICE_IP
|
|
port = int(DEFAULT_YT_AUTH_SERVICE_PORT)
|
|
redis_conn_id = DEFAULT_REDIS_CONN_ID
|
|
|
|
client, transport = None, None
|
|
try:
|
|
client, transport = _get_thrift_client(host, port)
|
|
redis_client = _get_redis_client(redis_conn_id)
|
|
|
|
logger.info("Fetching all account statuses from the service...")
|
|
all_accounts = client.getAccountStatus(accountPrefix=None)
|
|
logger.info(f"Found {len(all_accounts)} accounts to process.")
|
|
|
|
accounts_to_unban = []
|
|
accounts_to_activate = []
|
|
accounts_to_rest = []
|
|
|
|
for acc in all_accounts:
|
|
if acc.status == "BANNED (expired)":
|
|
accounts_to_unban.append(acc.accountId)
|
|
elif acc.status == "RESTING (expired)":
|
|
accounts_to_activate.append(acc.accountId)
|
|
elif acc.status == "ACTIVE (should be resting)":
|
|
accounts_to_rest.append(acc.accountId)
|
|
|
|
# --- Perform State Transitions ---
|
|
|
|
# 1. Un-ban accounts via Thrift call
|
|
if accounts_to_unban:
|
|
logger.info(f"Un-banning {len(accounts_to_unban)} accounts: {accounts_to_unban}")
|
|
for acc_id in accounts_to_unban:
|
|
try:
|
|
client.unbanAccount(acc_id, "Automatic un-ban by Airflow maintenance DAG.")
|
|
logger.info(f"Successfully un-banned account '{acc_id}'.")
|
|
except Exception as e:
|
|
logger.error(f"Failed to un-ban account '{acc_id}': {e}")
|
|
|
|
# 2. Activate resting accounts via direct Redis write
|
|
if accounts_to_activate:
|
|
logger.info(f"Activating {len(accounts_to_activate)} accounts: {accounts_to_activate}")
|
|
now_ts = int(time.time())
|
|
with redis_client.pipeline() as pipe:
|
|
for acc_id in accounts_to_activate:
|
|
key = f"account_status:{acc_id}"
|
|
pipe.hset(key, "status", "ACTIVE")
|
|
pipe.hset(key, "active_since_timestamp", now_ts)
|
|
pipe.hset(key, "status_changed_timestamp", now_ts)
|
|
pipe.execute()
|
|
logger.info("Finished activating accounts.")
|
|
|
|
# 3. Rest active accounts via direct Redis write
|
|
if accounts_to_rest:
|
|
logger.info(f"Putting {len(accounts_to_rest)} accounts to rest: {accounts_to_rest}")
|
|
now_ts = int(time.time())
|
|
with redis_client.pipeline() as pipe:
|
|
for acc_id in accounts_to_rest:
|
|
key = f"account_status:{acc_id}"
|
|
pipe.hset(key, "status", "RESTING")
|
|
pipe.hset(key, "status_changed_timestamp", now_ts)
|
|
pipe.execute()
|
|
logger.info("Finished putting accounts to rest.")
|
|
|
|
finally:
|
|
if transport and transport.isOpen():
|
|
transport.close()
|
|
|
|
|
|
with DAG(
|
|
dag_id='ytdlp_ops_account_maintenance',
|
|
default_args=DEFAULT_ARGS,
|
|
schedule='*/5 * * * *', # Run every 5 minutes
|
|
start_date=days_ago(1),
|
|
catchup=False,
|
|
tags=['ytdlp', 'maintenance'],
|
|
doc_md=__doc__,
|
|
) as dag:
|
|
manage_account_states()
|