111 lines
3.7 KiB
Python
111 lines
3.7 KiB
Python
import logging
|
|
import os
|
|
from typing import Optional
|
|
|
|
from apscheduler.schedulers.background import BackgroundScheduler
|
|
from apscheduler.triggers.cron import CronTrigger
|
|
from fastapi import FastAPI
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.iptv.stream_manager import StreamManager
|
|
from app.models.db import ChannelDB
|
|
from app.utils.database import get_db_session
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class StreamScheduler:
|
|
"""Scheduler service for periodic stream validation tasks."""
|
|
|
|
def __init__(self, app: Optional[FastAPI] = None):
|
|
"""
|
|
Initialize the scheduler with optional FastAPI app integration.
|
|
|
|
Args:
|
|
app: Optional FastAPI app instance for lifecycle integration
|
|
"""
|
|
self.scheduler = BackgroundScheduler()
|
|
self.app = app
|
|
self.batch_size = int(os.getenv("STREAM_VALIDATION_BATCH_SIZE", "10"))
|
|
self.schedule_time = os.getenv(
|
|
"STREAM_VALIDATION_SCHEDULE", "0 3 * * *"
|
|
) # Default 3 AM daily
|
|
logger.info(f"Scheduler initialized with app: {app is not None}")
|
|
|
|
def validate_streams_batch(self, db_session: Optional[Session] = None) -> None:
|
|
"""
|
|
Validate streams and update their status.
|
|
When batch_size=0, validates all channels.
|
|
|
|
Args:
|
|
db_session: Optional SQLAlchemy session
|
|
"""
|
|
db = db_session if db_session else get_db_session()
|
|
try:
|
|
manager = StreamManager(db)
|
|
|
|
# Get channels to validate
|
|
query = db.query(ChannelDB)
|
|
if self.batch_size > 0:
|
|
query = query.limit(self.batch_size)
|
|
channels = query.all()
|
|
|
|
for channel in channels:
|
|
try:
|
|
logger.info(f"Validating streams for channel {channel.id}")
|
|
manager.validate_and_select_stream(str(channel.id))
|
|
except Exception as e:
|
|
logger.error(f"Error validating channel {channel.id}: {str(e)}")
|
|
continue
|
|
|
|
logger.info(f"Completed stream validation of {len(channels)} channels")
|
|
finally:
|
|
if db_session is None:
|
|
db.close()
|
|
|
|
def start(self) -> None:
|
|
"""Start the scheduler and add jobs."""
|
|
if not self.scheduler.running:
|
|
# Add the scheduled job
|
|
self.scheduler.add_job(
|
|
self.validate_streams_batch,
|
|
trigger=CronTrigger.from_crontab(self.schedule_time),
|
|
id="daily_stream_validation",
|
|
)
|
|
|
|
# Start the scheduler
|
|
self.scheduler.start()
|
|
logger.info(
|
|
f"Stream scheduler started with daily validation job. "
|
|
f"Running: {self.scheduler.running}"
|
|
)
|
|
|
|
# Register shutdown handler if FastAPI app is provided
|
|
if self.app:
|
|
logger.info(
|
|
f"Registering scheduler with FastAPI "
|
|
f"app: {hasattr(self.app, 'state')}"
|
|
)
|
|
|
|
@self.app.on_event("shutdown")
|
|
def shutdown_scheduler():
|
|
self.shutdown()
|
|
|
|
def shutdown(self) -> None:
|
|
"""Shutdown the scheduler gracefully."""
|
|
if self.scheduler.running:
|
|
self.scheduler.shutdown()
|
|
logger.info("Stream scheduler stopped")
|
|
|
|
def trigger_manual_validation(self) -> None:
|
|
"""Trigger manual validation of streams."""
|
|
logger.info("Manually triggering stream validation")
|
|
self.validate_streams_batch()
|
|
|
|
|
|
def init_scheduler(app: FastAPI) -> StreamScheduler:
|
|
"""Initialize and start the scheduler with FastAPI integration."""
|
|
scheduler = StreamScheduler(app)
|
|
scheduler.start()
|
|
return scheduler
|