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

plugin_gmail/utils/client.py

@@ -0,0 +1,494 @@
+from __future__ import annotations
+
+import importlib
+import logging
+import time
+from dataclasses import dataclass
+from typing import Any, Callable, Iterator, TypeVar
+
+from googleapiclient.errors import HttpError
+
+from shared_plugins.exceptions import PluginConfigurationError, PluginCursorExpiredError
+from shared_plugins.google_client.batch_retry import (
+    BatchRetryExhaustedError,
+    BatchRetryPolicy,
+    BatchTerminalSubrequestError,
+    execute_batch_with_failed_subset_retries,
+)
+from shared_plugins.google_client.http_errors import extract_http_status_code
+from shared_plugins.models import IngressModel
+
+from ..models.ingress import (
+    GmailHistoryListResponseIngress,
+    GmailHistoryRecordIngress,
+    GmailLabelIngress,
+    GmailLabelsListResponseIngress,
+    GmailMessageAttachmentIngress,
+    GmailMessageIngress,
+    GmailMessagesListResponseIngress,
+    GmailProfileIngress,
+    GmailThreadIngress,
+    GmailThreadsListResponseIngress,
+)
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T", bound=IngressModel)
+USER_ID = "me"
+LIST_MAX_RESULTS = 500
+HISTORY_MAX_RESULTS = 500
+API_NUM_RETRIES = 5
+MESSAGES_GET_BATCH_MAX_SUBREQUESTS = 20
+MESSAGES_GET_BATCH_RETRY_POLICY = BatchRetryPolicy()
+MESSAGES_ATTACHMENTS_GET_BATCH_MAX_SUBREQUESTS = 20
+MESSAGES_ATTACHMENTS_GET_BATCH_RETRY_POLICY = BatchRetryPolicy()
+
+
+@dataclass(frozen=True)
+class HistoryPage:
+    history: list[GmailHistoryRecordIngress]
+    history_id: str
+
+
+GmailServiceFactory = Callable[[], Any]
+
+
+def load_gmail_service_factory(path: str) -> GmailServiceFactory:
+    module_name, _, attr = path.rpartition(".")
+    if not module_name or not attr:
+        raise PluginConfigurationError(
+            f"Invalid GMAIL_SERVICE_FACTORY '{path}'. Expected dotted path like 'pkg.module.make_service'."
+        )
+    module = importlib.import_module(module_name)
+    factory = getattr(module, attr, None)
+    if not callable(factory):
+        raise PluginConfigurationError(
+            f"GMAIL_SERVICE_FACTORY '{path}' does not resolve to a callable"
+        )
+    return factory
+
+
+class GmailApiClient:
+    def __init__(self, *, service: Any) -> None:
+        if service is None:
+            raise PluginConfigurationError(
+                "No Gmail service provided. Pass an authenticated service via service=... "
+                "or resolve one from GMAIL_SERVICE_FACTORY in your orchestrator."
+            )
+        self._service = service
+
+    def _execute_model(self, request: Any, model_type: type[T]) -> T:
+        payload = request.execute(num_retries=API_NUM_RETRIES)
+        return model_type.model_validate(payload)
+
+    def get_profile(self) -> GmailProfileIngress:
+        request = self._service.users().getProfile(userId=USER_ID)
+        return self._execute_model(request, GmailProfileIngress)
+
+    def iter_labels(self) -> Iterator[GmailLabelIngress]:
+        request = self._service.users().labels().list(userId=USER_ID)
+        payload = self._execute_model(request, GmailLabelsListResponseIngress)
+        for label in payload.labels:
+            yield label
+
+    def iter_message_ids(
+        self,
+        *,
+        query: str | None = None,
+        label_ids: list[str] | None = None,
+        max_results: int | None = None,
+    ) -> Iterator[str]:
+        for message_ids in self.iter_message_id_pages(
+            query=query,
+            label_ids=label_ids,
+            max_results=max_results,
+        ):
+            for message_id in message_ids:
+                yield message_id
+
+    def iter_message_id_pages(
+        self,
+        *,
+        query: str | None = None,
+        label_ids: list[str] | None = None,
+        max_results: int | None = None,
+    ) -> Iterator[list[str]]:
+        page_token: str | None = None
+        page_size = max_results or LIST_MAX_RESULTS
+
+        while True:
+            kwargs: dict[str, Any] = {
+                "userId": USER_ID,
+                "maxResults": page_size,
+            }
+            if query:
+                kwargs["q"] = query
+            if label_ids:
+                kwargs["labelIds"] = label_ids
+            if page_token:
+                kwargs["pageToken"] = page_token
+
+            request = self._service.users().messages().list(**kwargs)
+            payload = self._execute_model(request, GmailMessagesListResponseIngress)
+
+            yield [item.id for item in payload.messages]
+
+            page_token = payload.next_page_token
+            if not page_token:
+                break
+
+    def _build_message_get_request(self, message_id: str, *, fmt: str) -> Any:
+        return (
+            self._service.users()
+            .messages()
+            .get(
+                userId=USER_ID,
+                id=message_id,
+                format=fmt,
+            )
+        )
+
+    def get_message(
+        self,
+        message_id: str,
+        *,
+        fmt: str = "full",
+    ) -> GmailMessageIngress:
+        request = self._build_message_get_request(message_id, fmt=fmt)
+        return self._execute_model(request, GmailMessageIngress)
+
+    def get_attachment(
+        self,
+        message_id: str,
+        attachment_id: str,
+    ) -> GmailMessageAttachmentIngress:
+        request = (
+            self._service.users()
+            .messages()
+            .attachments()
+            .get(
+                userId=USER_ID,
+                messageId=message_id,
+                id=attachment_id,
+            )
+        )
+        payload = request.execute(num_retries=API_NUM_RETRIES)
+        size: int | None = None
+        data: str | None = None
+        if isinstance(payload, dict):
+            raw_size = payload.get("size")
+            if isinstance(raw_size, int):
+                size = raw_size
+
+            raw_data = payload.get("data")
+            if isinstance(raw_data, str):
+                data = raw_data
+
+        return GmailMessageAttachmentIngress.model_validate(
+            {
+                "attachment_id": attachment_id,
+                "size": size,
+                "data": data,
+            }
+        )
+
+    def get_attachments_batch(
+        self,
+        keys: list[tuple[str, str]],
+    ) -> list[GmailMessageAttachmentIngress]:
+        if not keys:
+            return []
+        if len(keys) > MESSAGES_ATTACHMENTS_GET_BATCH_MAX_SUBREQUESTS:
+            raise ValueError(
+                "keys exceeds Gmail batch limit: "
+                f"{len(keys)} > {MESSAGES_ATTACHMENTS_GET_BATCH_MAX_SUBREQUESTS}"
+            )
+
+        request_id_to_key: dict[str, tuple[str, str]] = {}
+        request_factories: dict[str, Callable[[], Any]] = {}
+        for index, (message_id, attachment_id) in enumerate(keys):
+            request_id = str(index)
+            request_id_to_key[request_id] = (message_id, attachment_id)
+            request_factories[request_id] = lambda m_id=message_id, a_id=attachment_id: (
+                self._service.users()
+                .messages()
+                .attachments()
+                .get(userId=USER_ID, messageId=m_id, id=a_id)
+            )
+
+        try:
+            responses = execute_batch_with_failed_subset_retries(
+                request_factories=request_factories,
+                new_batch=lambda callback: self._service.new_batch_http_request(
+                    callback=callback
+                ),
+                policy=MESSAGES_ATTACHMENTS_GET_BATCH_RETRY_POLICY,
+            )
+        except BatchTerminalSubrequestError as exc:
+            failed_key = request_id_to_key.get(exc.request_id)
+            raise RuntimeError(
+                "Gmail batch attachments.get failed with a non-retryable subrequest "
+                f"for key={failed_key}: {exc.exception}"
+            ) from exc.exception
+        except BatchRetryExhaustedError as exc:
+            failed_keys = [
+                str(request_id_to_key.get(request_id, request_id))
+                for request_id in exc.failed_request_ids
+            ]
+            sample_keys = ", ".join(failed_keys[:5])
+            raise RuntimeError(
+                "Gmail batch attachments.get failed after retries for "
+                f"{len(failed_keys)} of {len(keys)} subrequests. "
+                f"Sample failed keys: {sample_keys}"
+            ) from exc
+
+        attachments: list[GmailMessageAttachmentIngress] = []
+        for index, (_message_id, attachment_id) in enumerate(keys):
+            request_id = str(index)
+            if request_id not in responses:
+                raise RuntimeError(
+                    "Gmail batch attachments.get missing response payload for "
+                    f"key={request_id_to_key[request_id]}"
+                )
+            payload = responses[request_id]
+            size: int | None = None
+            data: str | None = None
+            if isinstance(payload, dict):
+                raw_size = payload.get("size")
+                if isinstance(raw_size, int):
+                    size = raw_size
+                raw_data = payload.get("data")
+                if isinstance(raw_data, str):
+                    data = raw_data
+            attachments.append(
+                GmailMessageAttachmentIngress.model_validate(
+                    {
+                        "attachment_id": attachment_id,
+                        "size": size,
+                        "data": data,
+                    }
+                )
+            )
+        return attachments
+
+    def iter_attachment_batches(
+        self,
+        keys: list[tuple[str, str]],
+        *,
+        batch_size: int = MESSAGES_ATTACHMENTS_GET_BATCH_MAX_SUBREQUESTS,
+    ) -> Iterator[list[GmailMessageAttachmentIngress]]:
+        if batch_size < 1:
+            raise ValueError("batch_size must be >= 1")
+        if batch_size > MESSAGES_ATTACHMENTS_GET_BATCH_MAX_SUBREQUESTS:
+            raise ValueError(
+                "batch_size exceeds Gmail batch limit: "
+                f"{batch_size} > {MESSAGES_ATTACHMENTS_GET_BATCH_MAX_SUBREQUESTS}"
+            )
+
+        prev_batch_end: float | None = None
+        for offset in range(0, len(keys), batch_size):
+            batch_keys = keys[offset : offset + batch_size]
+            batch_start = time.monotonic()
+            delta_since_prev = (
+                batch_start - prev_batch_end if prev_batch_end is not None else 0.0
+            )
+            logger.info(
+                "gmail attachments.get batch starting: size=%d offset=%d delta_since_prev=%.3fs",
+                len(batch_keys),
+                offset,
+                delta_since_prev,
+            )
+            batch = self.get_attachments_batch(batch_keys)
+            batch_end = time.monotonic()
+            duration = batch_end - batch_start
+            logger.info(
+                "gmail attachments.get batch complete: returned=%d duration=%.3fs rate=%.1f atts/s",
+                len(batch),
+                duration,
+                len(batch) / max(duration, 1e-9),
+            )
+            prev_batch_end = batch_end
+            yield batch
+
+    def get_messages_batch(
+        self,
+        message_ids: list[str],
+        *,
+        fmt: str = "full",
+    ) -> list[GmailMessageIngress]:
+        if not message_ids:
+            return []
+        if len(message_ids) > MESSAGES_GET_BATCH_MAX_SUBREQUESTS:
+            raise ValueError(
+                "message_ids exceeds Gmail batch limit: "
+                f"{len(message_ids)} > {MESSAGES_GET_BATCH_MAX_SUBREQUESTS}"
+            )
+
+        request_id_to_message_id: dict[str, str] = {}
+        request_factories: dict[str, Callable[[], Any]] = {}
+        for index, message_id in enumerate(message_ids):
+            request_id = str(index)
+            request_id_to_message_id[request_id] = message_id
+            request_factories[request_id] = lambda m_id=message_id: (
+                self._build_message_get_request(m_id, fmt=fmt)
+            )
+
+        try:
+            responses = execute_batch_with_failed_subset_retries(
+                request_factories=request_factories,
+                new_batch=lambda callback: self._service.new_batch_http_request(
+                    callback=callback
+                ),
+                policy=MESSAGES_GET_BATCH_RETRY_POLICY,
+            )
+        except BatchTerminalSubrequestError as exc:
+            failed_message_id = request_id_to_message_id.get(
+                exc.request_id, exc.request_id
+            )
+            raise RuntimeError(
+                "Gmail batch messages.get failed with a non-retryable subrequest "
+                f"for message_id={failed_message_id}: {exc.exception}"
+            ) from exc.exception
+        except BatchRetryExhaustedError as exc:
+            failed_ids = [
+                request_id_to_message_id.get(request_id, request_id)
+                for request_id in exc.failed_request_ids
+            ]
+            sample_ids = ", ".join(failed_ids[:5])
+            raise RuntimeError(
+                "Gmail batch messages.get failed after retries for "
+                f"{len(failed_ids)} of {len(message_ids)} subrequests. "
+                f"Sample failed message IDs: {sample_ids}"
+            ) from exc
+
+        messages: list[GmailMessageIngress] = []
+        for index, message_id in enumerate(message_ids):
+            request_id = str(index)
+            if request_id not in responses:
+                raise RuntimeError(
+                    "Gmail batch messages.get missing response payload for "
+                    f"message_id={message_id}"
+                )
+            messages.append(GmailMessageIngress.model_validate(responses[request_id]))
+        return messages
+
+    def iter_message_batches(
+        self,
+        message_ids: list[str],
+        *,
+        batch_size: int = MESSAGES_GET_BATCH_MAX_SUBREQUESTS,
+        fmt: str = "full",
+    ) -> Iterator[list[GmailMessageIngress]]:
+        if batch_size < 1:
+            raise ValueError("batch_size must be >= 1")
+        if batch_size > MESSAGES_GET_BATCH_MAX_SUBREQUESTS:
+            raise ValueError(
+                "batch_size exceeds Gmail batch limit: "
+                f"{batch_size} > {MESSAGES_GET_BATCH_MAX_SUBREQUESTS}"
+            )
+
+        prev_batch_end: float | None = None
+        for offset in range(0, len(message_ids), batch_size):
+            batch_ids = message_ids[offset : offset + batch_size]
+            batch_start = time.monotonic()
+            delta_since_prev = (
+                batch_start - prev_batch_end if prev_batch_end is not None else 0.0
+            )
+            logger.info(
+                "gmail messages.get batch starting: size=%d offset=%d delta_since_prev=%.3fs",
+                len(batch_ids),
+                offset,
+                delta_since_prev,
+            )
+            batch = self.get_messages_batch(batch_ids, fmt=fmt)
+            batch_end = time.monotonic()
+            duration = batch_end - batch_start
+            logger.info(
+                "gmail messages.get batch complete: returned=%d duration=%.3fs rate=%.1f msgs/s",
+                len(batch),
+                duration,
+                len(batch) / max(duration, 1e-9),
+            )
+            prev_batch_end = batch_end
+            yield batch
+
+    def iter_thread_ids(
+        self,
+        *,
+        query: str | None = None,
+        label_ids: list[str] | None = None,
+        max_results: int | None = None,
+    ) -> Iterator[str]:
+        page_token: str | None = None
+        page_size = max_results or LIST_MAX_RESULTS
+
+        while True:
+            kwargs: dict[str, Any] = {
+                "userId": USER_ID,
+                "maxResults": page_size,
+            }
+            if query:
+                kwargs["q"] = query
+            if label_ids:
+                kwargs["labelIds"] = label_ids
+            if page_token:
+                kwargs["pageToken"] = page_token
+
+            request = self._service.users().threads().list(**kwargs)
+            payload = self._execute_model(request, GmailThreadsListResponseIngress)
+
+            for item in payload.threads:
+                yield item.id
+
+            page_token = payload.next_page_token
+            if not page_token:
+                break
+
+    def get_thread(self, thread_id: str, *, fmt: str = "full") -> GmailThreadIngress:
+        request = (
+            self._service.users()
+            .threads()
+            .get(
+                userId=USER_ID,
+                id=thread_id,
+                format=fmt,
+            )
+        )
+        return self._execute_model(request, GmailThreadIngress)
+
+    def iter_history_pages(
+        self,
+        *,
+        start_history_id: str,
+        max_results: int | None = None,
+    ) -> Iterator[HistoryPage]:
+        page_token: str | None = None
+        page_size = max_results or HISTORY_MAX_RESULTS
+
+        while True:
+            kwargs: dict[str, Any] = {
+                "userId": USER_ID,
+                "startHistoryId": str(start_history_id),
+                "maxResults": page_size,
+            }
+            if page_token:
+                kwargs["pageToken"] = page_token
+
+            try:
+                request = self._service.users().history().list(**kwargs)
+                payload = self._execute_model(request, GmailHistoryListResponseIngress)
+            except HttpError as exc:
+                if extract_http_status_code(exc) == 404:
+                    raise PluginCursorExpiredError(
+                        "Gmail history cursor is stale. Reset history source state and rerun history sync to reseed from the current profile historyId. "
+                        "For strict completeness after expiration, run backfill as well."
+                    ) from exc
+                raise
+
+            yield HistoryPage(
+                history=payload.history, history_id=str(payload.history_id)
+            )
+
+            page_token = payload.next_page_token
+            if not page_token:
+                break