Started (incomplete) implementation of stream verification scheduler and endpoints
All checks were successful
AWS Deploy on Push / build (push) Successful in 5m18s

This commit is contained in:
2025-06-17 17:12:39 -05:00
parent abb467749b
commit a42d4c30a6
14 changed files with 1066 additions and 33 deletions

110
app/iptv/scheduler.py Normal file
View File

@@ -0,0 +1,110 @@
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

151
app/iptv/stream_manager.py Normal file
View File

@@ -0,0 +1,151 @@
import logging
import random
from typing import Optional
from sqlalchemy.orm import Session
from app.models.db import ChannelURL
from app.utils.check_streams import StreamValidator
from app.utils.database import get_db_session
logger = logging.getLogger(__name__)
class StreamManager:
"""Service for managing and validating channel streams."""
def __init__(self, db_session: Optional[Session] = None):
"""
Initialize StreamManager with optional database session.
Args:
db_session: Optional SQLAlchemy session. If None, will create a new one.
"""
self.db = db_session if db_session else get_db_session()
self.validator = StreamValidator()
def get_streams_for_channel(self, channel_id: str) -> list[ChannelURL]:
"""
Get all streams for a channel ordered by priority (lowest first),
with same-priority streams randomized.
Args:
channel_id: UUID of the channel to get streams for
Returns:
List of ChannelURL objects ordered by priority
"""
try:
# Get all streams for channel ordered by priority
streams = (
self.db.query(ChannelURL)
.filter(ChannelURL.channel_id == channel_id)
.order_by(ChannelURL.priority_id)
.all()
)
# Group streams by priority and randomize same-priority streams
grouped = {}
for stream in streams:
if stream.priority_id not in grouped:
grouped[stream.priority_id] = []
grouped[stream.priority_id].append(stream)
# Randomize same-priority streams and flatten
randomized_streams = []
for priority in sorted(grouped.keys()):
random.shuffle(grouped[priority])
randomized_streams.extend(grouped[priority])
return randomized_streams
except Exception as e:
logger.error(f"Error getting streams for channel {channel_id}: {str(e)}")
raise
def validate_and_select_stream(self, channel_id: str) -> Optional[str]:
"""
Find and validate a working stream for the given channel.
Args:
channel_id: UUID of the channel to find a stream for
Returns:
URL of the first working stream found, or None if none found
"""
try:
streams = self.get_streams_for_channel(channel_id)
if not streams:
logger.warning(f"No streams found for channel {channel_id}")
return None
working_stream = None
for stream in streams:
logger.info(f"Validating stream {stream.url} for channel {channel_id}")
is_valid, _ = self.validator.validate_stream(stream.url)
if is_valid:
working_stream = stream
break
if working_stream:
self._update_stream_status(working_stream, streams)
return working_stream.url
else:
logger.warning(f"No valid streams found for channel {channel_id}")
return None
except Exception as e:
logger.error(f"Error validating streams for channel {channel_id}: {str(e)}")
raise
def _update_stream_status(
self, working_stream: ChannelURL, all_streams: list[ChannelURL]
) -> None:
"""
Update in_use status for streams (True for working stream, False for others).
Args:
working_stream: The stream that was validated as working
all_streams: All streams for the channel
"""
try:
for stream in all_streams:
stream.in_use = stream.id == working_stream.id
self.db.commit()
logger.info(
f"Updated stream status - set in_use=True for {working_stream.url}"
)
except Exception as e:
self.db.rollback()
logger.error(f"Error updating stream status: {str(e)}")
raise
def __del__(self):
"""Close database session when StreamManager is destroyed."""
if hasattr(self, "db"):
self.db.close()
def get_working_stream(
channel_id: str, db_session: Optional[Session] = None
) -> Optional[str]:
"""
Convenience function to get a working stream for a channel.
Args:
channel_id: UUID of the channel to get a stream for
db_session: Optional SQLAlchemy session
Returns:
URL of the first working stream found, or None if none found
"""
manager = StreamManager(db_session)
try:
return manager.validate_and_select_stream(channel_id)
finally:
if db_session is None: # Only close if we created the session
manager.__del__()