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

plugin_gmail/sources/attachments.py

@@ -0,0 +1,307 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from typing import Any
+
+import dlt
+from dlt.destinations.sql_client import SqlClientBase
+from pydantic import ValidationError
+from shared_plugins.dlt import destination_has_table
+from shared_plugins.models import format_validation_error
+from shared_plugins.naming import (
+    dlt_resource_name,
+    dlt_source_name,
+    plugin_id_from_module,
+)
+from shared_plugins.resources import ctx_dlt_resource
+from shared_plugins.values import non_empty_string
+
+from ..models.ctx import AttachmentCandidateProjection, AttachmentRow
+from ..models.ingress import GmailMessageAttachmentIngress
+from ..models.translators import attachment_candidate_rows_to_ctx_models
+from ..utils.attachments import (
+    AttachmentMaterializationError,
+    FetchedAttachmentMap,
+    build_attachment_rows_for_message,
+)
+from ..utils.client import (
+    MESSAGES_ATTACHMENTS_GET_BATCH_MAX_SUBREQUESTS,
+    GmailApiClient,
+)
+
+PLUGIN_ID = plugin_id_from_module(__file__)
+JOB = "attachment_content"
+ATTACHMENT_CONTENT_CANDIDATE_LIMIT = 200
+LOGGER = logging.getLogger(__name__)
+
+AttachmentCandidate = AttachmentCandidateProjection
+
+
+def parse_attachment_candidate(row: Mapping[str, Any]) -> AttachmentCandidateProjection:
+    return next(attachment_candidate_rows_to_ctx_models([row]))
+
+
+def _validate_limit(limit: int) -> None:
+    if isinstance(limit, bool) or not isinstance(limit, int) or limit < 1:
+        raise ValueError("limit must be an integer >= 1")
+
+
+def build_attachment_candidate_query(*, limit: int) -> str:
+    _validate_limit(limit)
+
+    return """
+SELECT
+  m.id,
+  m.attachment_count,
+  m.attachments,
+  COALESCE(a.attachment_row_count, 0) AS existing_attachment_count
+FROM messages AS m
+LEFT JOIN (
+  SELECT
+    message_id,
+    COUNT(*)::bigint AS attachment_row_count
+  FROM attachments
+  WHERE _ctx_binding_id = %s
+  GROUP BY message_id
+) AS a
+  ON a.message_id = m.id
+WHERE m._ctx_binding_id = %s
+  AND m.attachment_count > 0
+  AND m.attachment_count <> COALESCE(a.attachment_row_count, 0)
+ORDER BY m._ctx_source_updated_at DESC NULLS LAST, m.id ASC
+LIMIT %s
+""".strip()
+
+
+def build_attachment_candidate_query_bootstrap(*, limit: int) -> str:
+    _validate_limit(limit)
+
+    return """
+SELECT
+  m.id,
+  m.attachment_count,
+  m.attachments,
+  0::bigint AS existing_attachment_count
+FROM messages AS m
+WHERE m._ctx_binding_id = %s
+  AND m.attachment_count > 0
+ORDER BY m._ctx_source_updated_at DESC NULLS LAST, m.id ASC
+LIMIT %s
+""".strip()
+
+
+def iter_attachment_candidate_rows(
+    sql_client: SqlClientBase[Any],
+    *,
+    binding_id: str,
+    limit: int,
+) -> Iterator[dict[str, Any]]:
+    if destination_has_table(sql_client, "attachments"):
+        query = build_attachment_candidate_query(limit=limit)
+        yield from _iter_query_rows(sql_client, query, binding_id, binding_id, limit)
+    else:
+        LOGGER.info(
+            "gmail.attachments.bootstrap binding_id=%s reason=attachments table does not exist yet; all candidates treated as new",
+            binding_id,
+        )
+        query = build_attachment_candidate_query_bootstrap(limit=limit)
+        yield from _iter_query_rows(sql_client, query, binding_id, limit)
+
+
+def _parse_attachment_candidate_row(
+    *,
+    binding_id: str,
+    raw_candidate: Mapping[str, Any],
+) -> AttachmentCandidateProjection | None:
+    message_id = non_empty_string(raw_candidate.get("id")) or "-"
+    try:
+        return parse_attachment_candidate(raw_candidate)
+    except ValidationError as exc:
+        LOGGER.error(
+            "gmail.attachments.materialization_failed binding_id=%s message_id=%s part_id=- attachment_id=- reason=Candidate message row failed validation: %s",
+            binding_id,
+            message_id,
+            format_validation_error(exc),
+        )
+        return None
+
+
+def _rebuild_attachment_rows(
+    *,
+    binding_id: str,
+    candidate: AttachmentCandidateProjection,
+    fetched_map: FetchedAttachmentMap,
+) -> list[AttachmentRow] | None:
+    if candidate.attachment_count > 0 and len(candidate.attachments) == 0:
+        LOGGER.error(
+            "gmail.attachments.materialization_failed binding_id=%s message_id=%s part_id=- attachment_id=- reason=Candidate message has attachment_count>0 with empty attachments metadata",
+            binding_id,
+            candidate.message_id,
+        )
+        return None
+
+    try:
+        message_rows = build_attachment_rows_for_message(
+            binding_id=binding_id,
+            message_id=candidate.message_id,
+            attachments=candidate.attachments,
+            fetched_map=fetched_map,
+        )
+    except AttachmentMaterializationError as exc:
+        LOGGER.error(
+            "gmail.attachments.materialization_failed binding_id=%s message_id=%s part_id=%s attachment_id=%s reason=%s",
+            binding_id,
+            candidate.message_id,
+            exc.part_id or "-",
+            exc.attachment_id or "-",
+            exc.reason,
+        )
+        return None
+
+    LOGGER.info(
+        "gmail.attachments.message_rebuilt binding_id=%s message_id=%s expected_count=%d existing_count=%d rebuilt_count=%d",
+        binding_id,
+        candidate.message_id,
+        candidate.attachment_count,
+        candidate.existing_attachment_count,
+        len(message_rows),
+    )
+    return message_rows
+
+
+def _fetch_keys_for_candidate(
+    candidate: AttachmentCandidateProjection,
+) -> list[tuple[str, str]]:
+    return [
+        (candidate.message_id, attachment.attachment_id)
+        for attachment in candidate.attachments
+        if attachment.attachment_id is not None
+    ]
+
+
+def _group_candidates_by_fetch_budget(
+    candidates: Iterable[AttachmentCandidateProjection],
+    *,
+    max_fetches_per_group: int,
+) -> Iterator[list[AttachmentCandidateProjection]]:
+    if max_fetches_per_group < 1:
+        raise ValueError("max_fetches_per_group must be >= 1")
+
+    group: list[AttachmentCandidateProjection] = []
+    group_fetch_count = 0
+    for candidate in candidates:
+        fetch_count = len(_fetch_keys_for_candidate(candidate))
+        if group and group_fetch_count + fetch_count > max_fetches_per_group:
+            yield group
+            group = []
+            group_fetch_count = 0
+        group.append(candidate)
+        group_fetch_count += fetch_count
+    if group:
+        yield group
+
+
+def _fetch_attachments_for_group(
+    client: GmailApiClient,
+    group: Sequence[AttachmentCandidateProjection],
+) -> dict[tuple[str, str], GmailMessageAttachmentIngress]:
+    all_keys: list[tuple[str, str]] = []
+    for candidate in group:
+        all_keys.extend(_fetch_keys_for_candidate(candidate))
+
+    fetched_map: dict[tuple[str, str], GmailMessageAttachmentIngress] = {}
+    offset = 0
+    for batch in client.iter_attachment_batches(all_keys):
+        batch_keys = all_keys[offset : offset + len(batch)]
+        fetched_map.update(zip(batch_keys, batch))
+        offset += len(batch)
+    return fetched_map
+
+
+@dlt.source(name=dlt_source_name(PLUGIN_ID, JOB))
+def gmail_attachment_content_source(
+    binding_id: str,
+    *,
+    client: GmailApiClient,
+) -> tuple[Any, ...]:
+    @ctx_dlt_resource(
+        name=dlt_resource_name("attachments"),
+        write_disposition={"disposition": "merge", "strategy": "delete-insert"},
+        primary_key=("_ctx_binding_id", "message_id", "part_id"),
+        merge_key=("_ctx_binding_id", "message_id"),
+        columns={
+            "content_id": {"data_type": "text"},
+            "file_path": {"data_type": "text"},
+        },
+    )
+    def attachments_resource() -> Iterator[AttachmentRow]:
+        candidate_count = 0
+        successful_messages = 0
+        failed_messages = 0
+        rows_written = 0
+
+        with dlt.current.pipeline().sql_client() as sql_client:
+            raw_candidates = list(
+                iter_attachment_candidate_rows(
+                    sql_client,
+                    binding_id=binding_id,
+                    limit=ATTACHMENT_CONTENT_CANDIDATE_LIMIT,
+                )
+            )
+
+        candidates: list[AttachmentCandidateProjection] = []
+        for raw_candidate in raw_candidates:
+            candidate_count += 1
+            candidate = _parse_attachment_candidate_row(
+                binding_id=binding_id,
+                raw_candidate=raw_candidate,
+            )
+            if candidate is None:
+                failed_messages += 1
+                continue
+            candidates.append(candidate)
+
+        for group in _group_candidates_by_fetch_budget(
+            candidates,
+            max_fetches_per_group=MESSAGES_ATTACHMENTS_GET_BATCH_MAX_SUBREQUESTS,
+        ):
+            fetched_map = _fetch_attachments_for_group(client, group)
+            for candidate in group:
+                message_rows = _rebuild_attachment_rows(
+                    binding_id=binding_id,
+                    candidate=candidate,
+                    fetched_map=fetched_map,
+                )
+                if message_rows is None:
+                    failed_messages += 1
+                    continue
+
+                successful_messages += 1
+                rows_written += len(message_rows)
+                yield from message_rows
+
+        LOGGER.info(
+            "gmail.attachments.run_complete binding_id=%s candidates=%d rebuilt_messages=%d failed_messages=%d rows_written=%d",
+            binding_id,
+            candidate_count,
+            successful_messages,
+            failed_messages,
+            rows_written,
+        )
+
+    return (attachments_resource,)
+
+
+def _iter_query_rows(
+    sql_client: SqlClientBase[Any],
+    query: str,
+    *params: Any,
+) -> Iterator[dict[str, Any]]:
+    with sql_client.execute_query(query, *params) as cursor:
+        if cursor.description is None:
+            return
+
+        columns = [column[0] for column in cursor.description]
+        for row in cursor.fetchall():
+            yield dict(zip(columns, row))

plugin_gmail/sources/backfill.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import logging
+import time
+from collections.abc import Iterator
+from typing import Any
+
+import dlt
+from shared_plugins.naming import (
+    dlt_resource_name,
+    dlt_source_name,
+    plugin_id_from_module,
+)
+from shared_plugins.resources import ctx_dlt_resource
+
+from ..models.ctx import MESSAGES_COLUMNS, LabelRow, MessageRow, ProfileRow
+from ..models.translators import (
+    labels_to_ctx_models,
+    messages_to_ctx_models,
+    profiles_to_ctx_models,
+)
+from ..utils.client import MESSAGES_GET_BATCH_MAX_SUBREQUESTS, GmailApiClient
+
+logger = logging.getLogger(__name__)
+
+PLUGIN_ID = plugin_id_from_module(__file__)
+JOB = "backfill"
+BACKFILL_MESSAGES_MAX_RESULTS = 50_000
+
+
+@dlt.source(name=dlt_source_name(PLUGIN_ID, JOB))
+def gmail_backfill_source(
+    binding_id: str,
+    *,
+    client: GmailApiClient,
+) -> tuple[Any, ...]:
+    """Gmail bootstrap over profile, labels, and messages for one binding."""
+
+    @ctx_dlt_resource(
+        name=dlt_resource_name("profile"),
+        write_disposition="merge",
+        primary_key=("_ctx_binding_id",),
+    )
+    def profile_resource() -> Iterator[ProfileRow]:
+        yield from profiles_to_ctx_models(binding_id, [client.get_profile()])
+
+    @ctx_dlt_resource(
+        name=dlt_resource_name("labels"),
+        write_disposition="merge",
+        primary_key=("_ctx_binding_id", "id"),
+    )
+    def labels_resource() -> Iterator[LabelRow]:
+        yield from labels_to_ctx_models(binding_id, client.iter_labels())
+
+    @ctx_dlt_resource(
+        name=dlt_resource_name("messages"),
+        write_disposition="merge",
+        primary_key=("_ctx_binding_id", "id"),
+        columns=MESSAGES_COLUMNS,
+    )
+    def messages_resource() -> Iterator[MessageRow]:
+        count = 0
+        page_count = 0
+        batch_count = 0
+        t0 = time.monotonic()
+        logger.info(
+            "messages_resource: starting message page enumeration (limit=%d)",
+            BACKFILL_MESSAGES_MAX_RESULTS,
+        )
+        for message_ids in client.iter_message_id_pages():
+            remaining = BACKFILL_MESSAGES_MAX_RESULTS - count
+            if remaining <= 0:
+                logger.info(
+                    "messages_resource: reached message limit=%d, stopping",
+                    BACKFILL_MESSAGES_MAX_RESULTS,
+                )
+                break
+
+            message_ids_to_process = message_ids[:remaining]
+            page_count += 1
+            elapsed = time.monotonic() - t0
+            logger.info(
+                "messages_resource: page %d enumerated %d IDs, processing %d IDs (%.1fs elapsed)",
+                page_count,
+                len(message_ids),
+                len(message_ids_to_process),
+                elapsed,
+            )
+
+            for message_batch in client.iter_message_batches(
+                message_ids_to_process,
+                batch_size=MESSAGES_GET_BATCH_MAX_SUBREQUESTS,
+            ):
+                batch_count += 1
+                count += len(message_batch)
+                yield from messages_to_ctx_models(binding_id, message_batch)
+
+                if batch_count % 10 == 0:
+                    elapsed = time.monotonic() - t0
+                    logger.info(
+                        "messages_resource: fetched %d messages in %d batches across %d pages (%.1fs elapsed)",
+                        count,
+                        batch_count,
+                        page_count,
+                        elapsed,
+                    )
+
+            if count >= BACKFILL_MESSAGES_MAX_RESULTS:
+                logger.info(
+                    "messages_resource: reached message limit=%d after page %d",
+                    BACKFILL_MESSAGES_MAX_RESULTS,
+                    page_count,
+                )
+                break
+
+        elapsed = time.monotonic() - t0
+        logger.info(
+            "messages_resource: finished %d messages across %d pages in %d batches (%.1fs)",
+            count,
+            page_count,
+            batch_count,
+            elapsed,
+        )
+
+    return (
+        profile_resource,
+        labels_resource,
+        messages_resource,
+    )

plugin_gmail/sources/history.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import Any
+
+import dlt
+from dlt.destinations.sql_client import SqlClientBase
+from shared_plugins.google_client.http_errors import extract_http_status_code
+from shared_plugins.dlt import destination_has_table
+from shared_plugins.naming import (
+    dlt_resource_name,
+    dlt_source_name,
+    plugin_id_from_module,
+)
+from shared_plugins.resources import ctx_dlt_resource, ctx_dlt_transformer
+from shared_plugins.values import require_non_negative_int
+
+from ..models.ctx import (
+    MESSAGES_COLUMNS,
+    HistoryEventRow,
+    LabelRow,
+    MessageRow,
+    ProfileRow,
+)
+from ..models.translators import (
+    extract_changed_ids,
+    history_events_to_ctx_models,
+    labels_to_ctx_models,
+    messages_to_ctx_models,
+    profiles_to_ctx_models,
+    utc_now_iso,
+)
+from ..utils.client import GmailApiClient
+
+HISTORY_CURSOR_KEY = "history_cursor"
+HISTORY_LAST_SYNCED_AT_KEY = "history_last_synced_at"
+HISTORY_SEEDED_AT_KEY = "history_seeded_at"
+PLUGIN_ID = plugin_id_from_module(__file__)
+JOB = "history"
+
+
+def read_max_messages_history_id(
+    sql_client: SqlClientBase[Any],
+    *,
+    binding_id: str,
+) -> int | None:
+    if not destination_has_table(sql_client, "messages"):
+        return None
+
+    with sql_client.execute_query(
+        """
+SELECT MAX(history_id)
+FROM messages
+WHERE _ctx_binding_id = %s
+""".strip(),
+        binding_id,
+    ) as cursor:
+        row = cursor.fetchone()
+        if row is None:
+            return None
+        max_history_id = row[0]
+
+    if max_history_id is None:
+        return None
+
+    return require_non_negative_int(max_history_id)
+
+
+@dlt.source(name=dlt_source_name(PLUGIN_ID, JOB))
+def gmail_history_source(
+    binding_id: str,
+    *,
+    client: GmailApiClient,
+) -> tuple[Any, ...]:
+    """Incremental Gmail history polling job with explicit cursor state."""
+
+    seen_message_ids: set[str] = set()
+
+    @ctx_dlt_resource(
+        name=dlt_resource_name("profile"),
+        write_disposition="merge",
+        primary_key=("_ctx_binding_id",),
+    )
+    def profile_snapshot_resource() -> Iterator[ProfileRow]:
+        yield from profiles_to_ctx_models(binding_id, [client.get_profile()])
+
+    @ctx_dlt_resource(
+        name=dlt_resource_name("labels"),
+        write_disposition="merge",
+        primary_key=("_ctx_binding_id", "id"),
+    )
+    def labels_snapshot_resource() -> Iterator[LabelRow]:
+        yield from labels_to_ctx_models(binding_id, client.iter_labels())
+
+    @ctx_dlt_resource(
+        name=dlt_resource_name("history_events"),
+        write_disposition="append",
+        primary_key=("_ctx_binding_id", "id"),
+    )
+    def history_events_resource() -> Iterator[HistoryEventRow]:
+        source_state = dlt.current.source_state()
+        cursor = source_state.get(HISTORY_CURSOR_KEY)
+
+        if not cursor:
+            with dlt.current.pipeline().sql_client() as sql_client:
+                seeded_cursor = read_max_messages_history_id(
+                    sql_client,
+                    binding_id=binding_id,
+                )
+
+            if seeded_cursor is None:
+                return
+
+            source_state[HISTORY_CURSOR_KEY] = str(seeded_cursor)
+            source_state[HISTORY_SEEDED_AT_KEY] = utc_now_iso()
+            return
+
+        start_cursor = str(require_non_negative_int(cursor))
+        latest_cursor = start_cursor
+        for page in client.iter_history_pages(start_history_id=start_cursor):
+            latest_cursor = page.history_id
+            yield from history_events_to_ctx_models(binding_id, page.history)
+
+        source_state[HISTORY_CURSOR_KEY] = latest_cursor
+        source_state[HISTORY_LAST_SYNCED_AT_KEY] = utc_now_iso()
+
+    @ctx_dlt_transformer(
+        data_from=history_events_resource,
+        name=dlt_resource_name("messages"),
+        write_disposition="merge",
+        primary_key=("_ctx_binding_id", "id"),
+        columns=MESSAGES_COLUMNS,
+    )
+    def messages_from_history(
+        event_record: HistoryEventRow | dict[str, Any],
+    ) -> Iterator[MessageRow]:
+        event_row = (
+            event_record
+            if isinstance(event_record, HistoryEventRow)
+            else HistoryEventRow.model_validate(event_record)
+        )
+        message_ids, _ = extract_changed_ids(event_row)
+        for message_id in sorted(message_ids):
+            if message_id in seen_message_ids:
+                continue
+            seen_message_ids.add(message_id)
+            try:
+                message = client.get_message(message_id)
+            except Exception as exc:
+                if extract_http_status_code(exc) == 404:
+                    continue
+                raise
+            yield from messages_to_ctx_models(binding_id, [message])
+
+    return (
+        profile_snapshot_resource,
+        labels_snapshot_resource,
+        history_events_resource,
+        messages_from_history,
+    )