contextbase-plugin-imessage-local 0.2.9__py3-none-any.whl

plugin_imessage_local/models/translators.py

@@ -0,0 +1,384 @@
+from __future__ import annotations
+
+import base64
+from collections.abc import Iterable, Iterator, Mapping
+from datetime import date, datetime
+from typing import Any
+
+from ..attributed_body import decode_attributed_body
+from .ctx import (
+    AttachmentRow,
+    ChatHandleJoinRow,
+    ChatMessageJoinRow,
+    ChatRow,
+    HandleRow,
+    MessageAttachmentJoinRow,
+    MessageRow,
+)
+from .ingress import (
+    IMessageAttachment,
+    IMessageChat,
+    IMessageChatHandleJoin,
+    IMessageChatMessageJoin,
+    IMessageHandle,
+    IMessageMessage,
+    IMessageMessageAttachmentJoin,
+)
+
+
+def _json_safe(value: Any) -> Any:
+    """Convert a decoded-attribute structure to JSON-serializable values.
+
+    ``bytes`` inside attribute dicts (e.g. DataDetected bplist blobs) become
+    base64 strings so the whole tree fits into a JSONB column without
+    silently dropping any bytes.
+    """
+    if value is None or isinstance(value, (str, int, float, bool)):
+        return value
+    if isinstance(value, bytes):
+        return base64.b64encode(value).decode("ascii")
+    if isinstance(value, (datetime, date)):
+        return value.isoformat()
+    if isinstance(value, Mapping):
+        return {str(k): _json_safe(v) for k, v in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_json_safe(v) for v in value]
+    return str(value)
+
+
+def handles_to_ctx_models(
+    *,
+    binding_id: str,
+    rows: Iterable[IMessageHandle],
+) -> Iterator[HandleRow]:
+    for handle in rows:
+        yield HandleRow(
+            ctx_binding_id=binding_id,
+            ctx_source_updated_at=None,
+            id=handle.id,
+            service=handle.service,
+            country=handle.country,
+            uncanonicalized_id=handle.uncanonicalized_id,
+            person_centric_id=handle.person_centric_id,
+        )
+
+
+def chats_to_ctx_models(
+    *,
+    binding_id: str,
+    rows: Iterable[IMessageChat],
+) -> Iterator[ChatRow]:
+    for chat in rows:
+        yield ChatRow(
+            ctx_binding_id=binding_id,
+            ctx_source_updated_at=None,
+            guid=chat.guid,
+            style=chat.style,
+            state=chat.state,
+            account_id=chat.account_id,
+            properties=chat.properties,
+            chat_identifier=chat.chat_identifier,
+            service_name=chat.service_name,
+            room_name=chat.room_name,
+            account_login=chat.account_login,
+            is_archived=chat.is_archived,
+            last_addressed_handle=chat.last_addressed_handle,
+            display_name=chat.display_name,
+            group_id=chat.group_id,
+            is_filtered=chat.is_filtered,
+            successful_query=chat.successful_query,
+            engram_id=chat.engram_id,
+            server_change_token=chat.server_change_token,
+            ck_sync_state=chat.ck_sync_state,
+            original_group_id=chat.original_group_id,
+            last_read_message_timestamp=chat.last_read_message_timestamp,
+            sr_server_change_token=chat.sr_server_change_token,
+            sr_ck_sync_state=chat.sr_ck_sync_state,
+            cloudkit_record_id=chat.cloudkit_record_id,
+            sr_cloudkit_record_id=chat.sr_cloudkit_record_id,
+            last_addressed_sim_id=chat.last_addressed_sim_id,
+            is_blackholed=chat.is_blackholed,
+            syndication_date=chat.syndication_date,
+            syndication_type=chat.syndication_type,
+            is_recovered=chat.is_recovered,
+            is_deleting_incoming_messages=chat.is_deleting_incoming_messages,
+        )
+
+
+def attachments_to_ctx_models(
+    *,
+    binding_id: str,
+    rows: Iterable[IMessageAttachment],
+) -> Iterator[AttachmentRow]:
+    for attachment in rows:
+        yield AttachmentRow(
+            ctx_binding_id=binding_id,
+            ctx_source_updated_at=attachment.created_date,
+            guid=attachment.guid,
+            created_date=attachment.created_date,
+            start_date=attachment.start_date,
+            filename=attachment.filename,
+            uti=attachment.uti,
+            mime_type=attachment.mime_type,
+            transfer_state=attachment.transfer_state,
+            is_outgoing=attachment.is_outgoing,
+            user_info=attachment.user_info,
+            transfer_name=attachment.transfer_name,
+            total_bytes=attachment.total_bytes,
+            is_sticker=attachment.is_sticker,
+            sticker_user_info=attachment.sticker_user_info,
+            attribution_info=attachment.attribution_info,
+            hide_attachment=attachment.hide_attachment,
+            ck_sync_state=attachment.ck_sync_state,
+            ck_server_change_token_blob=attachment.ck_server_change_token_blob,
+            ck_record_id=attachment.ck_record_id,
+            original_guid=attachment.original_guid,
+            sr_ck_sync_state=attachment.sr_ck_sync_state,
+            sr_ck_server_change_token_blob=attachment.sr_ck_server_change_token_blob,
+            sr_ck_record_id=attachment.sr_ck_record_id,
+            is_commsafety_sensitive=attachment.is_commsafety_sensitive,
+            emoji_image_content_identifier=attachment.emoji_image_content_identifier,
+            emoji_image_short_description=attachment.emoji_image_short_description,
+            preview_generation_state=attachment.preview_generation_state,
+        )
+
+
+def messages_to_ctx_models(
+    *,
+    binding_id: str,
+    rows: Iterable[IMessageMessage],
+    handle_lookup: Mapping[int, tuple[str, str]],
+) -> Iterator[MessageRow]:
+    for msg in rows:
+        decoded = (
+            decode_attributed_body(msg.attributed_body) if msg.attributed_body else None
+        )
+        if decoded is None:
+            decoded_text: str | None = None
+            attribute_runs: list[dict[str, Any]] | None = None
+        else:
+            decoded_text = decoded.text
+            attribute_runs = [
+                {
+                    "location": run.location,
+                    "length": run.length,
+                    "attributes": _json_safe(dict(run.attributes)),
+                }
+                for run in decoded.runs
+            ]
+
+        text = decoded_text if decoded_text is not None else msg.text
+
+        handle_natural_id, handle_natural_service = handle_lookup.get(
+            msg.handle_id, (None, None)
+        )
+        other_handle_natural_id, other_handle_natural_service = handle_lookup.get(
+            msg.other_handle, (None, None)
+        )
+
+        source_updated_at = max(
+            (
+                timestamp
+                for timestamp in (msg.date, msg.date_edited, msg.date_retracted)
+                if timestamp is not None
+            ),
+            default=None,
+        )
+
+        yield MessageRow(
+            ctx_binding_id=binding_id,
+            ctx_source_updated_at=source_updated_at,
+            guid=msg.guid,
+            text=text,
+            replace=msg.replace,
+            service_center=msg.service_center,
+            handle_id=msg.handle_id,
+            handle_natural_id=handle_natural_id,
+            handle_natural_service=handle_natural_service,
+            subject=msg.subject,
+            country=msg.country,
+            attributed_body=msg.attributed_body,
+            attribute_runs=attribute_runs,
+            version=msg.version,
+            type=msg.type,
+            service=msg.service,
+            account=msg.account,
+            account_guid=msg.account_guid,
+            error=msg.error,
+            date=msg.date,
+            date_read=msg.date_read,
+            date_delivered=msg.date_delivered,
+            is_delivered=msg.is_delivered,
+            is_finished=msg.is_finished,
+            is_emote=msg.is_emote,
+            is_from_me=msg.is_from_me,
+            is_empty=msg.is_empty,
+            is_delayed=msg.is_delayed,
+            is_auto_reply=msg.is_auto_reply,
+            is_prepared=msg.is_prepared,
+            is_read=msg.is_read,
+            is_system_message=msg.is_system_message,
+            is_sent=msg.is_sent,
+            has_dd_results=msg.has_dd_results,
+            is_service_message=msg.is_service_message,
+            is_forward=msg.is_forward,
+            was_downgraded=msg.was_downgraded,
+            is_archive=msg.is_archive,
+            was_data_detected=msg.was_data_detected,
+            was_deduplicated=msg.was_deduplicated,
+            is_audio_message=msg.is_audio_message,
+            is_played=msg.is_played,
+            date_played=msg.date_played,
+            item_type=msg.item_type,
+            other_handle=msg.other_handle,
+            other_handle_natural_id=other_handle_natural_id,
+            other_handle_natural_service=other_handle_natural_service,
+            group_title=msg.group_title,
+            group_action_type=msg.group_action_type,
+            share_status=msg.share_status,
+            share_direction=msg.share_direction,
+            is_expirable=msg.is_expirable,
+            expire_state=msg.expire_state,
+            message_action_type=msg.message_action_type,
+            message_source=msg.message_source,
+            associated_message_guid=msg.associated_message_guid,
+            associated_message_type=msg.associated_message_type,
+            balloon_bundle_id=msg.balloon_bundle_id,
+            payload_data=msg.payload_data,
+            expressive_send_style_id=msg.expressive_send_style_id,
+            associated_message_range_location=msg.associated_message_range_location,
+            associated_message_range_length=msg.associated_message_range_length,
+            time_expressive_send_played=msg.time_expressive_send_played,
+            message_summary_info=msg.message_summary_info,
+            destination_caller_id=msg.destination_caller_id,
+            is_corrupt=msg.is_corrupt,
+            reply_to_guid=msg.reply_to_guid,
+            sort_id=msg.sort_id,
+            is_spam=msg.is_spam,
+            has_unseen_mention=msg.has_unseen_mention,
+            thread_originator_guid=msg.thread_originator_guid,
+            thread_originator_part=msg.thread_originator_part,
+            syndication_ranges=msg.syndication_ranges,
+            was_delivered_quietly=msg.was_delivered_quietly,
+            did_notify_recipient=msg.did_notify_recipient,
+            synced_syndication_ranges=msg.synced_syndication_ranges,
+            date_retracted=msg.date_retracted,
+            date_edited=msg.date_edited,
+            was_detonated=msg.was_detonated,
+            part_count=msg.part_count,
+            is_stewie=msg.is_stewie,
+            is_kt_verified=msg.is_kt_verified,
+            is_sos=msg.is_sos,
+            is_critical=msg.is_critical,
+            bia_reference_id=msg.bia_reference_id,
+            fallback_hash=msg.fallback_hash,
+            associated_message_emoji=msg.associated_message_emoji,
+            is_pending_satellite_send=msg.is_pending_satellite_send,
+            needs_relay=msg.needs_relay,
+            schedule_type=msg.schedule_type,
+            schedule_state=msg.schedule_state,
+            sent_or_received_off_grid=msg.sent_or_received_off_grid,
+            date_recovered=msg.date_recovered,
+        )
+
+
+def chat_handle_joins_to_ctx_models(
+    *,
+    binding_id: str,
+    rows: Iterable[IMessageChatHandleJoin],
+    chat_lookup: Mapping[int, str],
+    handle_lookup: Mapping[int, tuple[str, str]],
+) -> Iterator[ChatHandleJoinRow]:
+    """Resolve chat rowid -> guid and handle rowid -> (id, service).
+
+    Raises if a join row references a missing chat or handle — the source
+    enforces FK integrity, so an unresolvable row indicates either
+    corruption or a pre-load race. Fail loudly per the "errors must be
+    loud" rule.
+    """
+    for row in rows:
+        chat_guid = chat_lookup.get(row.chat_id)
+        if chat_guid is None:
+            raise RuntimeError(
+                f"chat_handle_join row references missing chat_id={row.chat_id}"
+            )
+        handle_pair = handle_lookup.get(row.handle_id)
+        if handle_pair is None:
+            raise RuntimeError(
+                f"chat_handle_join row references missing handle_id={row.handle_id}"
+            )
+        handle_natural_id, handle_natural_service = handle_pair
+        yield ChatHandleJoinRow(
+            ctx_binding_id=binding_id,
+            ctx_source_updated_at=None,
+            chat_id=row.chat_id,
+            chat_guid=chat_guid,
+            handle_id=row.handle_id,
+            handle_natural_id=handle_natural_id,
+            handle_natural_service=handle_natural_service,
+        )
+
+
+def chat_message_joins_to_ctx_models(
+    *,
+    binding_id: str,
+    rows: Iterable[IMessageChatMessageJoin],
+    chat_lookup: Mapping[int, str],
+    message_lookup: Mapping[int, str],
+) -> Iterator[ChatMessageJoinRow]:
+    """Resolve chat rowid -> guid and message rowid -> guid.
+
+    Fails loudly on unresolvable rowids — the source declares FKs and the
+    triggers CASCADE on delete, so missing refs indicate corruption.
+    """
+    for row in rows:
+        chat_guid = chat_lookup.get(row.chat_id)
+        if chat_guid is None:
+            raise RuntimeError(
+                f"chat_message_join row references missing chat_id={row.chat_id}"
+            )
+        message_guid = message_lookup.get(row.message_id)
+        if message_guid is None:
+            raise RuntimeError(
+                f"chat_message_join row references missing message_id={row.message_id}"
+            )
+        yield ChatMessageJoinRow(
+            ctx_binding_id=binding_id,
+            ctx_source_updated_at=row.message_date,
+            chat_id=row.chat_id,
+            chat_guid=chat_guid,
+            message_id=row.message_id,
+            message_guid=message_guid,
+            message_date=row.message_date,
+        )
+
+
+def message_attachment_joins_to_ctx_models(
+    *,
+    binding_id: str,
+    rows: Iterable[IMessageMessageAttachmentJoin],
+    message_lookup: Mapping[int, str],
+    attachment_lookup: Mapping[int, str],
+) -> Iterator[MessageAttachmentJoinRow]:
+    """Resolve message rowid -> guid and attachment rowid -> guid."""
+    for row in rows:
+        message_guid = message_lookup.get(row.message_id)
+        if message_guid is None:
+            raise RuntimeError(
+                f"message_attachment_join row references missing "
+                f"message_id={row.message_id}"
+            )
+        attachment_guid = attachment_lookup.get(row.attachment_id)
+        if attachment_guid is None:
+            raise RuntimeError(
+                f"message_attachment_join row references missing "
+                f"attachment_id={row.attachment_id}"
+            )
+        yield MessageAttachmentJoinRow(
+            ctx_binding_id=binding_id,
+            ctx_source_updated_at=None,
+            message_id=row.message_id,
+            message_guid=message_guid,
+            attachment_id=row.attachment_id,
+            attachment_guid=attachment_guid,
+        )

plugin_imessage_local/plugin.json

@@ -0,0 +1,7 @@
+{
+	"auth": {
+		"type": "none"
+	},
+	"mode": "dagster",
+	"plugin_id": "imessage_local"
+}

plugin_imessage_local/sources/snapshot.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+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 shared_plugins.sqlite import sqlite_snapshot
+from sqlalchemy import Engine, create_engine, select
+from sqlalchemy.orm import Session
+
+from ..binding_config import IMessageLocalBindingConfig
+from ..models.ctx import (
+    AttachmentRow,
+    ChatHandleJoinRow,
+    ChatMessageJoinRow,
+    ChatRow,
+    HandleRow,
+    MessageAttachmentJoinRow,
+    MessageRow,
+)
+from ..models.ingress import (
+    IMessageAttachment,
+    IMessageChat,
+    IMessageChatHandleJoin,
+    IMessageChatMessageJoin,
+    IMessageHandle,
+    IMessageMessage,
+    IMessageMessageAttachmentJoin,
+)
+from ..models.translators import (
+    attachments_to_ctx_models,
+    chat_handle_joins_to_ctx_models,
+    chat_message_joins_to_ctx_models,
+    chats_to_ctx_models,
+    handles_to_ctx_models,
+    message_attachment_joins_to_ctx_models,
+    messages_to_ctx_models,
+)
+
+PLUGIN_ID = plugin_id_from_module(__file__)
+JOB = "snapshot"
+LOGGER = logging.getLogger(__name__)
+MERGE_WRITE_DISPOSITION = {"disposition": "merge", "strategy": "delete-insert"}
+MERGE_KEY = ("_ctx_binding_id",)
+
+
+@contextmanager
+def _db_snapshot(db_path: Path) -> Iterator[Engine]:
+    if not db_path.is_file():
+        raise RuntimeError(f"Missing iMessage database: '{db_path}'.")
+    with sqlite_snapshot(db_path) as snapshot_path:
+        engine = create_engine(f"sqlite:///{snapshot_path}")
+        try:
+            yield engine
+        finally:
+            engine.dispose()
+
+
+@dlt.source(name=dlt_source_name(PLUGIN_ID, JOB))
+def imessage_local_snapshot_source(
+    binding_id: str,
+    cfg: IMessageLocalBindingConfig,
+) -> tuple[Any, ...]:
+    chat_db = cfg.chat_db
+
+    @ctx_dlt_resource(
+        name=dlt_resource_name("handle"),
+        write_disposition=MERGE_WRITE_DISPOSITION,
+        merge_key=MERGE_KEY,
+        primary_key=("_ctx_binding_id", "id", "service"),
+    )
+    def handle_resource() -> Iterator[HandleRow]:
+        with _db_snapshot(chat_db) as engine:
+            with Session(engine) as session:
+                yield from handles_to_ctx_models(
+                    binding_id=binding_id,
+                    rows=session.scalars(select(IMessageHandle)),
+                )
+
+    @ctx_dlt_resource(
+        name=dlt_resource_name("chat"),
+        write_disposition=MERGE_WRITE_DISPOSITION,
+        merge_key=MERGE_KEY,
+        primary_key=("_ctx_binding_id", "guid"),
+    )
+    def chat_resource() -> Iterator[ChatRow]:
+        with _db_snapshot(chat_db) as engine:
+            with Session(engine) as session:
+                yield from chats_to_ctx_models(
+                    binding_id=binding_id,
+                    rows=session.scalars(select(IMessageChat)),
+                )
+
+    @ctx_dlt_resource(
+        name=dlt_resource_name("attachment"),
+        write_disposition=MERGE_WRITE_DISPOSITION,
+        merge_key=MERGE_KEY,
+        primary_key=("_ctx_binding_id", "guid"),
+    )
+    def attachment_resource() -> Iterator[AttachmentRow]:
+        with _db_snapshot(chat_db) as engine:
+            with Session(engine) as session:
+                yield from attachments_to_ctx_models(
+                    binding_id=binding_id,
+                    rows=session.scalars(select(IMessageAttachment)),
+                )
+
+    @ctx_dlt_resource(
+        name=dlt_resource_name("message"),
+        write_disposition=MERGE_WRITE_DISPOSITION,
+        merge_key=MERGE_KEY,
+        primary_key=("_ctx_binding_id", "guid"),
+    )
+    def message_resource() -> Iterator[MessageRow]:
+        with _db_snapshot(chat_db) as engine:
+            with Session(engine) as session:
+                handle_lookup = {
+                    rowid: (id_, service)
+                    for rowid, id_, service in session.execute(
+                        select(
+                            IMessageHandle.rowid,
+                            IMessageHandle.id,
+                            IMessageHandle.service,
+                        )
+                    )
+                }
+                yield from messages_to_ctx_models(
+                    binding_id=binding_id,
+                    rows=session.scalars(select(IMessageMessage)),
+                    handle_lookup=handle_lookup,
+                )
+
+    @ctx_dlt_resource(
+        name=dlt_resource_name("chat_handle_join"),
+        write_disposition=MERGE_WRITE_DISPOSITION,
+        merge_key=MERGE_KEY,
+        primary_key=(
+            "_ctx_binding_id",
+            "chat_guid",
+            "handle_natural_id",
+            "handle_natural_service",
+        ),
+    )
+    def chat_handle_join_resource() -> Iterator[ChatHandleJoinRow]:
+        with _db_snapshot(chat_db) as engine:
+            with Session(engine) as session:
+                chat_lookup = dict(
+                    session.execute(select(IMessageChat.rowid, IMessageChat.guid)).all()
+                )
+                handle_lookup = {
+                    rowid: (id_, service)
+                    for rowid, id_, service in session.execute(
+                        select(
+                            IMessageHandle.rowid,
+                            IMessageHandle.id,
+                            IMessageHandle.service,
+                        )
+                    )
+                }
+                yield from chat_handle_joins_to_ctx_models(
+                    binding_id=binding_id,
+                    rows=session.scalars(select(IMessageChatHandleJoin)),
+                    chat_lookup=chat_lookup,
+                    handle_lookup=handle_lookup,
+                )
+
+    @ctx_dlt_resource(
+        name=dlt_resource_name("chat_message_join"),
+        write_disposition=MERGE_WRITE_DISPOSITION,
+        merge_key=MERGE_KEY,
+        primary_key=("_ctx_binding_id", "chat_guid", "message_guid"),
+    )
+    def chat_message_join_resource() -> Iterator[ChatMessageJoinRow]:
+        with _db_snapshot(chat_db) as engine:
+            with Session(engine) as session:
+                chat_lookup = dict(
+                    session.execute(select(IMessageChat.rowid, IMessageChat.guid)).all()
+                )
+                message_lookup = dict(
+                    session.execute(
+                        select(IMessageMessage.rowid, IMessageMessage.guid)
+                    ).all()
+                )
+                yield from chat_message_joins_to_ctx_models(
+                    binding_id=binding_id,
+                    rows=session.scalars(select(IMessageChatMessageJoin)),
+                    chat_lookup=chat_lookup,
+                    message_lookup=message_lookup,
+                )
+
+    @ctx_dlt_resource(
+        name=dlt_resource_name("message_attachment_join"),
+        write_disposition=MERGE_WRITE_DISPOSITION,
+        merge_key=MERGE_KEY,
+        primary_key=("_ctx_binding_id", "message_guid", "attachment_guid"),
+    )
+    def message_attachment_join_resource() -> Iterator[MessageAttachmentJoinRow]:
+        with _db_snapshot(chat_db) as engine:
+            with Session(engine) as session:
+                message_lookup = dict(
+                    session.execute(
+                        select(IMessageMessage.rowid, IMessageMessage.guid)
+                    ).all()
+                )
+                attachment_lookup = dict(
+                    session.execute(
+                        select(IMessageAttachment.rowid, IMessageAttachment.guid)
+                    ).all()
+                )
+                yield from message_attachment_joins_to_ctx_models(
+                    binding_id=binding_id,
+                    rows=session.scalars(select(IMessageMessageAttachmentJoin)),
+                    message_lookup=message_lookup,
+                    attachment_lookup=attachment_lookup,
+                )
+
+    return (
+        handle_resource,
+        chat_resource,
+        attachment_resource,
+        message_resource,
+        chat_handle_join_resource,
+        chat_message_join_resource,
+        message_attachment_join_resource,
+    )