Started (incomplete) implementation of stream verification scheduler and endpoints
All checks were successful
AWS Deploy on Push / build (push) Successful in 5m18s
All checks were successful
AWS Deploy on Push / build (push) Successful in 5m18s
This commit is contained in:
110
app/iptv/scheduler.py
Normal file
110
app/iptv/scheduler.py
Normal 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
151
app/iptv/stream_manager.py
Normal 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__()
|
||||
Reference in New Issue
Block a user