152 lines
4.8 KiB
Python
152 lines
4.8 KiB
Python
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__()
|