mfinder/db/requests_sql.py

import datetime
import asyncio
from sqlalchemy import Column, Integer, BigInteger, TEXT, DateTime, inspect
from sqlalchemy.exc import OperationalError, PendingRollbackError
from mfinder.db.settings_sql import BASE, SESSION

class MovieRequest(BASE):
    __tablename__ = "movie_requests"
    id = Column(Integer, primary_key=True, autoincrement=True)
    user_id = Column(BigInteger, nullable=False)
    user_name = Column(TEXT, nullable=True)
    query = Column(TEXT, nullable=False)
    status = Column(TEXT, default="Pending")  # Pending, Uploaded, Rejected
    channel_msg_id = Column(Integer, nullable=True)
    created_at = Column(DateTime, default=datetime.datetime.utcnow)

    def __init__(self, user_id, user_name, query, status="Pending", channel_msg_id=None):
        self.user_id = user_id
        self.user_name = user_name
        self.query = query
        self.status = status
        self.channel_msg_id = channel_msg_id
        self.created_at = datetime.datetime.utcnow()

# Ensure table exists in DB automatically on module load
try:
    session = SESSION()
    bind = session.get_bind()
    inspector = inspect(bind)
    if not inspector.has_table("movie_requests"):
        MovieRequest.__table__.create(bind=bind, checkfirst=True)
except Exception as _e:
    pass

async def add_movie_request(user_id: int, user_name: str, query: str) -> MovieRequest:
    """Adds a new movie request to the database."""
    session = SESSION()
    try:
        req = MovieRequest(user_id=user_id, user_name=user_name, query=query)
        session.add(req)
        session.commit()
        session.refresh(req)
        return req
    except Exception as e:
        try:
            session.rollback()
        except Exception:
            pass
        try:
            SESSION.remove()
        except Exception:
            pass
        raise e
    finally:
        session.close()

async def update_request_status(req_id: int, status: str, channel_msg_id: int = None) -> MovieRequest:
    """Updates the status and optional channel message ID of a movie request."""
    session = SESSION()
    try:
        req = session.query(MovieRequest).filter_by(id=req_id).first()
        if req:
            req.status = status
            if channel_msg_id is not None:
                req.channel_msg_id = channel_msg_id
            session.commit()
            session.refresh(req)
        return req
    except Exception as e:
        try:
            session.rollback()
        except Exception:
            pass
        try:
            SESSION.remove()
        except Exception:
            pass
        raise e
    finally:
        session.close()

async def get_movie_request(req_id: int) -> MovieRequest:
    """Retrieves a movie request by ID."""
    for attempt in range(3):
        try:
            session = SESSION()
            return session.query(MovieRequest).filter_by(id=req_id).first()
        except (OperationalError, PendingRollbackError) as e:
            from mfinder import LOGGER
            LOGGER.warning("Database connection error in get_movie_request (attempt %d): %s. Reconnecting...", attempt + 1, str(e))
            try:
                SESSION.rollback()
            except Exception:
                pass
            try:
                SESSION.remove()
            except Exception:
                pass
            if attempt == 2:
                return None
            await asyncio.sleep(0.5)
        except Exception as e:
            try:
                SESSION.rollback()
            except Exception:
                pass
            try:
                SESSION.remove()
            except Exception:
                pass
            return None
        finally:
            SESSION.close()

async def get_pending_requests_count() -> int:
    """Returns total pending request count."""
    for attempt in range(3):
        try:
            session = SESSION()
            return session.query(MovieRequest).filter_by(status="Pending").count()
        except (OperationalError, PendingRollbackError) as e:
            from mfinder import LOGGER
            LOGGER.warning("Database connection error in get_pending_requests_count (attempt %d): %s. Reconnecting...", attempt + 1, str(e))
            try:
                SESSION.rollback()
            except Exception:
                pass
            try:
                SESSION.remove()
            except Exception:
                pass
            if attempt == 2:
                return 0
            await asyncio.sleep(0.5)
        except Exception as e:
            try:
                SESSION.rollback()
            except Exception:
                pass
            try:
                SESSION.remove()
            except Exception:
                pass
            return 0
        finally:
            SESSION.close()

async def get_all_pending_requests() -> list:
    """Retrieves all pending movie requests."""
    session = SESSION()
    try:
        return session.query(MovieRequest).filter_by(status="Pending").all()
    except Exception as e:
        try:
            session.rollback()
        except Exception:
            pass
        return []
    finally:
        session.close()