zulip_bot/app/db/zulip_service.py

"""
Service for querying messages from the Zulip database.
"""
from datetime import datetime, timedelta
from sqlalchemy import and_, or_
from app.db import get_db_session
from app.models.zulip import Message, Stream, Recipient, UserProfile, IT_RECIPIENT_IDS

class ZulipDatabaseService:
    """Service for querying messages from the Zulip database."""
    
    @staticmethod
    def get_messages_from_it_channels(days_ago=None, limit=1000, since=None):
        """
        Get recent messages from IT channels.
        
        Args:
            days_ago (int): Number of days to look back (optional)
            limit (int): Maximum number of messages to return
            since (datetime): Get messages after this datetime (optional)
            
        Returns:
            list: List of Message objects
        """
        session = get_db_session()
        
        # Build the query based on parameters
        query = session.query(Message).filter(
            Message.recipient_id.in_(IT_RECIPIENT_IDS)
        )
        
        # Add date filter if specified
        if since:
            query = query.filter(Message.date_sent >= since)
        elif days_ago:
            start_date = datetime.now() - timedelta(days=days_ago)
            query = query.filter(Message.date_sent >= start_date)
        
        # Get results
        messages = query.order_by(Message.id.desc()).limit(limit).all()
        
        return messages
    
    @staticmethod
    def get_messages_newer_than_id(message_id, limit=100):
        """
        Get messages with ID greater than the specified ID.
        
        Args:
            message_id (int): Get messages with ID greater than this
            limit (int): Maximum number of messages to return
            
        Returns:
            list: List of Message objects
        """
        session = get_db_session()
        
        messages = session.query(Message).filter(
            and_(
                Message.recipient_id.in_(IT_RECIPIENT_IDS),
                Message.id > message_id
            )
        ).order_by(Message.id.asc()).limit(limit).all()
        
        return messages
    
    @staticmethod
    def get_message_by_id(message_id):
        """
        Get a specific message by ID.
        
        Args:
            message_id (int): ID of the message to retrieve
            
        Returns:
            Message: Message object or None if not found
        """
        session = get_db_session()
        return session.query(Message).filter(Message.id == message_id).first()
    
    @staticmethod
    def search_messages(search_term, days_ago=365, limit=100):
        """
        Search for messages containing a specific term.
        
        Args:
            search_term (str): Term to search for
            days_ago (int): Number of days to look back
            limit (int): Maximum number of messages to return
            
        Returns:
            list: List of Message objects matching the search
        """
        session = get_db_session()
        start_date = datetime.now() - timedelta(days=days_ago)
        
        # Use the tsquery system if available, otherwise fall back to LIKE
        messages = session.query(Message).filter(
            and_(
                Message.recipient_id.in_(IT_RECIPIENT_IDS),
                Message.date_sent >= start_date,
                or_(
                    Message.content.ilike(f'%{search_term}%'),
                    Message.subject.ilike(f'%{search_term}%')
                )
            )
        ).order_by(Message.date_sent.desc()).limit(limit).all()
        
        return messages
    
    @staticmethod
    def get_channel_name_for_message(message):
        """
        Get the channel name for a message.
        
        Args:
            message (Message): Message object
            
        Returns:
            str: Channel name or "Unknown Channel" if not found
        """
        session = get_db_session()
        
        try:
            if not message or not message.recipient_id:
                return "Unknown Channel"
            
            # First, get the recipient to determine type
            recipient = session.query(Recipient).filter(
                Recipient.id == message.recipient_id
            ).first()
            
            if not recipient:
                return "Unknown Channel"
            
            # Check recipient type (1 = stream, 2 = user, 3 = huddle)
            if recipient.type != 1:
                # For direct messages or huddles
                return "Direct Message" if recipient.type == 2 else "Group Message"
                
            # For stream messages, get the stream name
            stream = session.query(Stream).filter(
                Stream.recipient_id == message.recipient_id
            ).first()
            
            # Return the name or a default value
            return stream.name if stream and stream.name else "Unknown Channel"
        except Exception as e:
            # Log the error but don't crash - return a default value
            print(f"Error getting channel name for message {message.id if message else 'unknown'}: {e}")
            return "Unknown Channel"
    
    @staticmethod
    def get_sender_name_for_message(message):
        """
        Get the sender name for a message.
        
        Args:
            message (Message): Message object
            
        Returns:
            str: Sender full name or 'Unknown User' if not found
        """
        session = get_db_session()
        
        try:
            if not message or not message.sender_id:
                return "Unknown User"
                
            user = session.query(UserProfile).filter(
                UserProfile.id == message.sender_id
            ).first()
            
            return user.full_name if user and user.full_name else "Unknown User"
        except Exception as e:
            # Log the error but don't crash - return a default value
            print(f"Error getting sender name for message {message.id if message else 'unknown'}: {e}")
            return "Unknown User"
    
    @staticmethod
    def count_messages_up_to_id(message_id, since=None):
        """
        Count messages with ID less than or equal to the specified ID.
        
        Args:
            message_id (int): Count messages with ID <= this
            since (datetime): Only count messages after this datetime (optional)
            
        Returns:
            int: Count of messages
        """
        session = get_db_session()
        
        # Build the query
        query = session.query(Message).filter(
            and_(
                Message.recipient_id.in_(IT_RECIPIENT_IDS),
                Message.id <= message_id
            )
        )
        
        # Add date filter if specified
        if since:
            query = query.filter(Message.date_sent >= since)
        
        # Count the messages
        count = query.count()
        
        return count