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):
    def __init__(self, event_manager):
        super().__init__(event_manager)


    def add_attachment(self, session, message, attachment_data):
        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):
        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):
        attachments = self.get_attachments_by_message_id(message_id)
        return json.dumps(attachments)

    def get_attachment_by_sha512(self, hash_sha512):
        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()