57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
from airflow.plugins_manager import AirflowPlugin
|
|
from airflow.hooks.base import BaseHook
|
|
from airflow.configuration import conf
|
|
import uuid
|
|
import backoff
|
|
|
|
class YTDLPHook(BaseHook):
|
|
def __init__(self, conn_id='ytdlp_default'):
|
|
super().__init__()
|
|
self.conn_id = conn_id
|
|
self.connection = self.get_connection(conn_id)
|
|
self.timeout = conf.getint('ytdlp', 'timeout', fallback=120)
|
|
self.max_retries = conf.getint('ytdlp', 'max_retries', fallback=3)
|
|
|
|
@backoff.on_exception(backoff.expo,
|
|
Exception,
|
|
max_tries=3,
|
|
max_time=300)
|
|
def start_service(self, host, port, service_id, work_dir):
|
|
"""Start token service as a long-running process"""
|
|
import subprocess
|
|
import os
|
|
from pathlib import Path
|
|
|
|
# Get script path relative to Airflow home
|
|
airflow_home = os.getenv('AIRFLOW_HOME', '')
|
|
script_path = Path(airflow_home).parent / 'ytdlp_ops_server.py'
|
|
|
|
# Ensure work directory exists
|
|
os.makedirs(work_dir, exist_ok=True)
|
|
|
|
# Start service process
|
|
cmd = [
|
|
'python', str(script_path),
|
|
'--port', str(port),
|
|
'--host', host,
|
|
'--service-id', service_id,
|
|
'--context-dir', work_dir,
|
|
'--script-dir', str(Path(airflow_home) / 'dags' / 'scripts')
|
|
]
|
|
|
|
self.log.info(f"Starting token service: {' '.join(cmd)}")
|
|
|
|
# Start process detached
|
|
docker_cmd = [
|
|
'docker-compose', '-f', 'docker-compose.yaml',
|
|
'up', '-d', '--build', 'ytdlp-service'
|
|
]
|
|
subprocess.run(docker_cmd, check=True)
|
|
|
|
self.log.info(f"Token service started on {host}:{port}")
|
|
return True
|
|
|
|
class YTDLPPlugin(AirflowPlugin):
|
|
name = 'ytdlp_plugin'
|
|
hooks = [YTDLPHook]
|