contextbase-plugin-gmail 0.2.6__py3-none-any.whl

plugin_gmail/utils/attachments.py

@@ -0,0 +1,251 @@
+from __future__ import annotations
+
+import hashlib
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from pathlib import PurePosixPath
+
+from pydantic import ValidationError
+from shared_plugins.scratch import replace_scratch_dir_files
+
+from ..models.ctx import AttachmentRow
+from ..models.ingress import GmailAttachmentIngress, GmailMessageAttachmentIngress
+from ..models.translators import Base64UrlDecodeError, decode_base64url_bytes
+
+
+@dataclass(frozen=True)
+class AttachmentMaterializationError(RuntimeError):
+    message_id: str
+    reason: str
+    part_id: str | None = None
+    attachment_id: str | None = None
+
+    def __str__(self) -> str:
+        return self.reason
+
+
+FetchedAttachmentMap = Mapping[tuple[str, str], GmailMessageAttachmentIngress]
+
+
+@dataclass(frozen=True)
+class _PreparedAttachment:
+    part_id: str
+    attachment_id: str | None
+    filename: str | None
+    mime_type: str | None
+    size: int | None
+    content_disposition: str | None
+    content_id: str | None
+    materialized_file_name: str
+    content: bytes
+
+
+def _hash_path_segment(*parts: str) -> str:
+    hash_input = "\n".join(parts)
+    return hashlib.sha256(hash_input.encode("utf-8")).hexdigest()
+
+
+def _build_deterministic_file_name(*, part_id: str, filename: str | None) -> str:
+    digest = _hash_path_segment(part_id, filename or "")
+    ext = ""
+    if filename:
+        suffixes = PurePosixPath(filename.strip()).suffixes
+        ext = "".join(suffixes)
+    return f"{digest}{ext}"
+
+
+def _append_stable_suffix(file_name: str, occurrence: int) -> str:
+    if occurrence <= 1:
+        return file_name
+
+    suffixes = PurePosixPath(file_name).suffixes
+    suffix = "".join(suffixes)
+    stem = file_name[: -len(suffix)] if suffix else file_name
+    return f"{stem}-{occurrence}{suffix}"
+
+
+def _decode_attachment_content(
+    *,
+    message_id: str,
+    attachment_metadata: GmailAttachmentIngress,
+    fetched_map: FetchedAttachmentMap,
+) -> tuple[bytes, GmailMessageAttachmentIngress | None]:
+    part_id = attachment_metadata.part_id
+    attachment_id = attachment_metadata.attachment_id
+
+    if attachment_id is not None:
+        fetched_attachment = fetched_map.get((message_id, attachment_id))
+        if fetched_attachment is None:
+            raise AttachmentMaterializationError(
+                message_id=message_id,
+                part_id=part_id,
+                attachment_id=attachment_id,
+                reason="Fetched attachment missing from batch result map.",
+            )
+
+        try:
+            content = decode_base64url_bytes(fetched_attachment.data)
+        except Base64UrlDecodeError as exc:
+            raise AttachmentMaterializationError(
+                message_id=message_id,
+                part_id=part_id,
+                attachment_id=attachment_id,
+                reason=f"Fetched attachment base64url decode failed: {exc}",
+            ) from exc
+
+        if content is None:
+            raise AttachmentMaterializationError(
+                message_id=message_id,
+                part_id=part_id,
+                attachment_id=attachment_id,
+                reason="Fetched attachment is missing decodable base64url content.",
+            )
+        return content, fetched_attachment
+
+    try:
+        content = decode_base64url_bytes(attachment_metadata.inline_data_b64url)
+    except Base64UrlDecodeError as exc:
+        raise AttachmentMaterializationError(
+            message_id=message_id,
+            part_id=part_id,
+            attachment_id=None,
+            reason=f"Inline attachment base64url decode failed: {exc}",
+        ) from exc
+
+    if content is None:
+        raise AttachmentMaterializationError(
+            message_id=message_id,
+            part_id=part_id,
+            attachment_id=None,
+            reason="Inline attachment metadata is missing decodable base64url content.",
+        )
+    return content, None
+
+
+def _prepare_attachment_for_materialization(
+    *,
+    message_id: str,
+    attachment_metadata: GmailAttachmentIngress,
+    fetched_map: FetchedAttachmentMap,
+    file_name_counts: dict[str, int],
+) -> _PreparedAttachment:
+    part_id = attachment_metadata.part_id
+    attachment_id = attachment_metadata.attachment_id
+    content, fetched_attachment = _decode_attachment_content(
+        message_id=message_id,
+        attachment_metadata=attachment_metadata,
+        fetched_map=fetched_map,
+    )
+
+    size = attachment_metadata.size
+    if size is None and fetched_attachment is not None:
+        size = fetched_attachment.size
+
+    deterministic_file_name = _build_deterministic_file_name(
+        part_id=part_id,
+        filename=attachment_metadata.filename,
+    )
+    next_occurrence = file_name_counts.get(deterministic_file_name, 0) + 1
+    file_name_counts[deterministic_file_name] = next_occurrence
+    materialized_file_name = _append_stable_suffix(
+        deterministic_file_name,
+        next_occurrence,
+    )
+
+    return _PreparedAttachment(
+        part_id=part_id,
+        attachment_id=attachment_id,
+        filename=attachment_metadata.filename,
+        mime_type=attachment_metadata.mime_type,
+        size=size,
+        content_disposition=attachment_metadata.content_disposition,
+        content_id=attachment_metadata.content_id,
+        materialized_file_name=materialized_file_name,
+        content=content,
+    )
+
+
+def _build_attachment_row(
+    *,
+    binding_id: str,
+    message_id: str,
+    prepared_attachment: _PreparedAttachment,
+    file_path: str,
+) -> AttachmentRow:
+    candidate_row = {
+        "ctx_binding_id": binding_id,
+        "ctx_source_updated_at": None,
+        "message_id": message_id,
+        "part_id": prepared_attachment.part_id,
+        "attachment_id": prepared_attachment.attachment_id,
+        "filename": prepared_attachment.filename,
+        "mime_type": prepared_attachment.mime_type,
+        "size": prepared_attachment.size,
+        "content_disposition": prepared_attachment.content_disposition,
+        "content_id": prepared_attachment.content_id,
+        "file_path": file_path,
+    }
+
+    try:
+        return AttachmentRow.model_validate(candidate_row)
+    except ValidationError as exc:
+        raise AttachmentMaterializationError(
+            message_id=message_id,
+            part_id=prepared_attachment.part_id,
+            attachment_id=prepared_attachment.attachment_id,
+            reason=f"Attachment row validation failed: {exc}",
+        ) from exc
+
+
+def build_attachment_rows_for_message(
+    *,
+    binding_id: str,
+    message_id: str,
+    attachments: Sequence[GmailAttachmentIngress],
+    fetched_map: FetchedAttachmentMap,
+) -> list[AttachmentRow]:
+    if len(attachments) == 0:
+        return []
+
+    safe_message_id = _hash_path_segment(message_id)
+    file_name_counts: dict[str, int] = {}
+    prepared_attachments = [
+        _prepare_attachment_for_materialization(
+            message_id=message_id,
+            attachment_metadata=attachment_metadata,
+            fetched_map=fetched_map,
+            file_name_counts=file_name_counts,
+        )
+        for attachment_metadata in attachments
+    ]
+    file_payloads = {
+        prepared_attachment.materialized_file_name: prepared_attachment.content
+        for prepared_attachment in prepared_attachments
+    }
+
+    try:
+        path_map = replace_scratch_dir_files(
+            binding_id=binding_id,
+            relative_dir=f"attachments/{safe_message_id}",
+            files=file_payloads,
+        )
+    except Exception as exc:
+        raise AttachmentMaterializationError(
+            message_id=message_id,
+            reason=f"Attachment file write failed: {exc}",
+        ) from exc
+
+    rebuilt_rows: list[AttachmentRow] = []
+    for prepared_attachment in prepared_attachments:
+        file_path = path_map[prepared_attachment.materialized_file_name]
+
+        rebuilt_rows.append(
+            _build_attachment_row(
+                binding_id=binding_id,
+                message_id=message_id,
+                prepared_attachment=prepared_attachment,
+                file_path=file_path,
+            )
+        )
+
+    return rebuilt_rows