FreeDATA/freedata_server/message_system_db_attachments.py

from message_system_db_manager import DatabaseManager
from message_system_db_model import MessageAttachment, Attachment, P2PMessage
import json
import hashlib
import os


class DatabaseManagerAttachments(DatabaseManager):
    """Manages database operations for message attachments.

    This class extends the DatabaseManager and provides methods for adding,
    retrieving, and deleting message attachments in the database. It also
    handles orphaned attachments and database sessions.
    """
    def __init__(self, event_manager):
        """Initializes DatabaseManagerAttachments.

        Args:
            event_manager (EventManager): The event manager instance.
        """
        super().__init__(event_manager)


    def add_attachment(self, session, message, attachment_data):
        """Adds an attachment to the database and links it to a message.

        This method adds a new attachment to the database if it doesn't
        already exist, based on its SHA-512 hash. It then creates a link
        between the message and the attachment using the MessageAttachment
        association table. It handles cases where the attachment already
        exists and logs appropriate messages.

        Args:
            session: The database session object.
            message: The message object to link the attachment to.
            attachment_data (dict): A dictionary containing the attachment
                data, including 'name', 'type', 'data', and optionally
                'checksum_crc32'.

        Returns:
            Attachment: The Attachment object that was added or found.
        """
        hash_sha512 = hashlib.sha512(attachment_data['data'].encode()).hexdigest()
        existing_attachment = session.query(Attachment).filter_by(hash_sha512=hash_sha512).first()

        if not existing_attachment:
            attachment = Attachment(
                name=attachment_data['name'],
                data_type=attachment_data['type'],
                data=attachment_data['data'],
                checksum_crc32=attachment_data.get('checksum_crc32', ''),
                hash_sha512=hash_sha512
            )
            session.add(attachment)
            session.flush()  # Ensure the attachment is persisted and has an ID
            self.log(f"Added attachment to database: {attachment.name}")
        else:
            attachment = existing_attachment
            self.log(f"Attachment {attachment.name} already exists in database")

        # Link the message and the attachment through MessageAttachment
        link = MessageAttachment(message=message, attachment=attachment)
        session.add(link)
        self.log(f"Linked message with attachment: {attachment.name}")

        return attachment

    def get_attachments_by_message_id(self, message_id):
        """Retrieves attachments associated with a message ID.

        This method retrieves all attachments linked to a given message ID
        from the database. It returns a list of dictionaries, where each
        dictionary represents an attachment. It handles cases where no
        message or attachments are found and logs any database errors.

        Args:
            message_id: The ID of the message whose attachments are to be retrieved.

        Returns:
            list: A list of dictionaries, each representing an attachment,
            or an empty list if no attachments or message are found.
        """
        session = self.get_thread_scoped_session()
        try:
            # Fetch the message by its ID
            message = session.query(P2PMessage).filter_by(id=message_id).first()
            if message:
                # Navigate through the association to get attachments
                attachments = [ma.attachment.to_dict() for ma in message.message_attachments]
                return attachments
            else:
                return []
        except Exception as e:
            self.log(f"Error fetching attachments for message ID {message_id}: {e}", isWarning=True)
            return []
        finally:
            session.remove()

    def get_attachments_by_message_id_json(self, message_id):
        """Retrieves attachments for a message as a JSON string.

        This method retrieves attachments associated with a given message ID
        and returns them as a JSON-formatted string. It calls
        get_attachments_by_message_id to fetch the attachments and then
        serializes them to JSON.

        Args:
            message_id: The ID of the message.

        Returns:
            str: A JSON string representing the attachments.
        """
        attachments = self.get_attachments_by_message_id(message_id)
        return json.dumps(attachments)

    def get_attachment_by_sha512(self, hash_sha512):
        """Retrieves an attachment by its SHA-512 hash.

        This method queries the database for an attachment matching the
        provided SHA-512 hash. If found, it returns the attachment data as
        a dictionary; otherwise, it returns None. It handles potential
        database errors and logs appropriate messages.

        Args:
            hash_sha512 (str): The SHA-512 hash of the attachment to retrieve.

        Returns:
            dict or None: The attachment data as a dictionary if found,
            None otherwise.
        """
        session = self.get_thread_scoped_session()
        try:
            attachment = session.query(Attachment).filter_by(hash_sha512=hash_sha512).first()
            if attachment:
                return attachment.to_dict()
            else:
                self.log(f"No attachment found with SHA-512 hash: {hash_sha512}")
                return None
        except Exception as e:
            self.log(f"Error fetching attachment with SHA-512 hash {hash_sha512}: {e}", isWarning=True)
            return None
        finally:
            session.remove()

    def delete_attachments_by_message_id(self, message_id):
        """
        Deletes attachment associations for a given message ID.

        For each attachment linked to the message:
          - If the attachment is linked to more than one message, only the association for this message is deleted.
          - If the attachment is linked only to this message, both the association and the attachment record are deleted.

        Parameters:
            message_id (str): The ID of the message whose attachment associations should be deleted.

        Returns:
            bool: True if the deletion was successful, False otherwise.
        """
        session = self.get_thread_scoped_session()
        try:
            # Find all attachment associations for the given message ID.
            links = session.query(MessageAttachment).filter_by(message_id=message_id).all()
            if not links:
                self.log(f"No attachments linked with message ID {message_id} found.")
                return True

            for link in links:
                # Count how many associations exist for this attachment.
                link_count = session.query(MessageAttachment).filter_by(attachment_id=link.attachment_id).count()
                if link_count > 1:
                    # More than one link exists, so only remove the association.
                    session.delete(link)
                    self.log(
                        f"Deleted link for attachment '{link.attachment.name}' from message {message_id} (other links exist).")
                else:
                    # Only one link exists, so delete both the association and the attachment.
                    session.delete(link)
                    session.delete(link.attachment)
                    self.log(f"Deleted attachment '{link.attachment.name}' from message {message_id} (only link).")

            session.commit()
            return True
        except Exception as e:
            session.rollback()
            self.log(f"Error deleting attachments for message ID {message_id}: {e}", isWarning=True)
            return False
        finally:
            session.remove()


    def clean_orphaned_attachments(self):
        """
        Checks for orphaned attachments in the database, i.e. attachments that have no
        MessageAttachment links to any messages. Optionally, deletes these orphaned attachments.

        Parameters:
            cleanup (bool): If True, deletes the orphaned attachments; if False, only returns them.

        Returns:
            If cleanup is False:
                list: A list of dictionaries representing the orphaned attachments.
            If cleanup is True:
                dict: A summary dictionary with the count of deleted attachments.
        """
        session = self.get_thread_scoped_session()
        try:
            orphaned = []
            # Get all attachments in the database.
            attachments = session.query(Attachment).all()
            for attachment in attachments:
                # Count the number of MessageAttachment links for this attachment.
                link_count = session.query(MessageAttachment).filter_by(attachment_id=attachment.id).count()
                if link_count == 0:
                    orphaned.append(attachment)

            for attachment in orphaned:
                self.log(f"Deleting orphaned attachment: {attachment.name}")
                session.delete(attachment)
            self.log(f"Checked for orphaned attachments")
            session.commit()
            return {'status': 'success', 'deleted_count': len(orphaned)}
        except Exception as e:
            session.rollback()
            self.log(f"Error checking orphaned attachments: {e}", isWarning=True)
            return None
        finally:
            session.remove()