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()