diff --git a/src/omniload/core/registry.py b/src/omniload/core/registry.py index 48f0bcfa5..5fb7b1dfa 100644 --- a/src/omniload/core/registry.py +++ b/src/omniload/core/registry.py @@ -84,6 +84,7 @@ "https+webdav": "dlt_filesystem.source.fsspec.webdav:WebdavSource", "hubspot": "omniload.source.hubspot.api:HubspotSource", "ibmmq+mqb": "omniload.source.mqbridge.api:MqBridgeSource", + "imap": "omniload.source.imap.api:ImapSource", "indeed": "omniload.source.indeed.api:IndeedSource", "influxdb": "omniload.source.influxdb.api:InfluxDBSource", "intercom": "omniload.source.intercom.api:IntercomSource", diff --git a/src/omniload/source/imap/adapter.py b/src/omniload/source/imap/adapter.py new file mode 100644 index 000000000..6d3663d7d --- /dev/null +++ b/src/omniload/source/imap/adapter.py @@ -0,0 +1,201 @@ +# Copyright 2022-2026 ScaleVector +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Reads messages and attachments from e-mail inbox via IMAP protocol""" + +import imaplib +from copy import deepcopy +from typing import Iterable, List, Optional, Sequence + +import dlt +from dlt.common import logger, pendulum +from dlt.sources import DltResource, TDataItem, TDataItems +from dlt.sources.filesystem import FileItemDict + +from .helpers import ( + extract_attachments, + extract_email_info, + get_message_uids, + get_message_with_internal_date, +) +from .settings import DEFAULT_CHUNK_SIZE, DEFAULT_START_DATE, GMAIL_GROUP + + +@dlt.source +def inbox_source( + host: str = dlt.secrets.value, + email_account: str = dlt.secrets.value, + password: str = dlt.secrets.value, + folder: str = "INBOX", + gmail_group: Optional[str] = GMAIL_GROUP, + start_date: pendulum.DateTime = DEFAULT_START_DATE, + filter_emails: Optional[Sequence[str]] = None, + filter_by_mime_type: Optional[Sequence[str]] = None, + chunksize: int = DEFAULT_CHUNK_SIZE, +) -> Sequence[DltResource]: + """This source collects inbox emails and downloads attachments to the local folder. + + Args: + host (str, optional): The hostname of the IMAP server. Default is 'dlt.secrets.value'. + email_account (str, optional): The email account used to log in to the IMAP server. Default is 'dlt.secrets.value'. + password (str, optional): The password for the email account. Default is 'dlt.secrets.value'. + folder (str, optional): The mailbox folder from which to collect emails. Default is 'INBOX'. + gmail_group (str, optional): The email address of the Google Group to filter emails sent to the group. Default is 'GMAIL_GROUP' from settings. + start_date (pendulum.Date, optional): The start date (with a day resolution) from which to collect emails. Default is 'DEFAULT_START_DATE' from settings. + filter_emails (Sequence[str], optional): A sequence of email addresses used to filter emails based on the 'FROM' field. Default is 'FILTER_EMAILS' from settings. + filter_by_mime_type (Sequence[str], optional): A sequence of MIME types used to filter attachments based on their content type. Default is an empty sequence. + chunksize (int, optional): The number of message UIDs to collect at a time. Default is 'DEFAULT_CHUNK_SIZE' from settings. + + Returns: + Sequence[DltResource]: Returns following resources: uids, messages, attachments + """ + + filter_emails = filter_emails or [] + filter_by_mime_type = filter_by_mime_type or [] + + def _login(client: imaplib.IMAP4_SSL) -> None: + # raise ValueError(f"email_account, password: {email_account}, {password}") + ##scsdc + client.login(email_account, password) + r, dat = client.select(folder, readonly=True) + if r != "OK": + raise client.error(dat[-1]) + + @dlt.resource(name="uids") + def get_messages_uids( + initial_message_num: Optional[ + dlt.sources.incremental[int] + ] = dlt.sources.incremental("message_uid", initial_value=1), + ) -> TDataItem: + """Collects email message UIDs (Unique IDs) from the mailbox. + Args: + initial_message_num (int, optional): Controls incremental loading on UID + + Yields: + TDataItem: A dictionary containing the 'message_uid' of the collected email message. + """ + + if initial_message_num is None or initial_message_num.last_value is None: + raise ValueError("initial_message_num.last_value is undefined") + + last_message_num = initial_message_num.last_value + + with imaplib.IMAP4_SSL(host) as client: + _login(client) + + criteria = [ + f"(SINCE {start_date.strftime('%d-%b-%Y')})", + f"(UID {str(int(last_message_num))}:*)", + ] + + if gmail_group: + logger.info(f"Load all emails for Group: {gmail_group}") + criteria.extend([f"(TO {gmail_group})"]) + + if filter_emails: + logger.info(f"Load emails only from: {filter_emails}") + if len(filter_emails) == 1: + criteria.append(f"(FROM {filter_emails[0]})") + else: + email_filter = " ".join( + [f"FROM {email}" for email in filter_emails] + ) + criteria.append(f"(OR {email_filter})") + + uids = get_message_uids(client, criteria) + if uids: + for i in range(0, len(uids), chunksize): + yield uids[i : i + chunksize] + + @dlt.transformer(name="messages", primary_key="message_uid") + def get_messages( + items: TDataItems, + include_body: bool = True, + ) -> TDataItem: + """Reads email messages from the mailbox based on the provided message UIDs. + + Args: + items (TDataItems): An iterable containing dictionaries with 'message_uid' representing the email message UIDs. + include_body (bool, optional): If True, includes the email body in the result. Default is True. + + Yields: + TDataItem: A dictionary containing the extracted email information from the read email message. + """ + + with imaplib.IMAP4_SSL(host) as client: + _login(client) + + for item in items: + message_uid = str(item["message_uid"]) + msg, internal_date = get_message_with_internal_date(client, message_uid) + result = deepcopy(item) + result["modification_date"] = internal_date + result.update(extract_email_info(msg, include_body=include_body)) + + yield result + + @dlt.transformer( + name="attachments", + primary_key="file_hash", + ) + def get_attachments( + items: TDataItems, + ) -> Iterable[List[FileItemDict]]: + """Downloads attachments from email messages based on the provided message UIDs. + + Args: + items (TDataItems): An iterable containing dictionaries with 'message_uid' representing the email message UIDs. + + Yields: + Iterable[List[FileItem]]: A dictionary containing the attachment FileItem. + """ + + with imaplib.IMAP4_SSL(host) as client: + _login(client) + + files_dict: List[FileItemDict] = [] + + for item in items: + message_uid = str(item["message_uid"]) + msg, internal_date = get_message_with_internal_date(client, message_uid) + attachments = list(extract_attachments(msg, filter_by_mime_type)) + if len(attachments) == 0: + continue + + email_info = extract_email_info(msg) + + for attachment in attachments: + attachment["modification_date"] = internal_date + attachment["file_url"] = ( + f"imap://{email_account}/{message_uid}/{attachment['file_name']}" + ) + + file_dict = FileItemDict(attachment) + file_dict["message"] = dict(email_info) + file_dict["message"].update(item) + + files_dict.append(file_dict) + if len(files_dict) >= chunksize: + yield files_dict + files_dict = [] + + # yield remainder + if files_dict: + yield files_dict + + return ( + get_messages_uids, + get_messages_uids | get_attachments, + get_messages_uids | get_messages, + ) diff --git a/src/omniload/source/imap/api.py b/src/omniload/source/imap/api.py new file mode 100644 index 000000000..163c66a88 --- /dev/null +++ b/src/omniload/source/imap/api.py @@ -0,0 +1,44 @@ +from typing import Any, Dict +from urllib.parse import parse_qs, urlparse + +from dlt.common.time import ensure_pendulum_datetime_utc + +from omniload.error import MissingValueError + + +class ImapSource: + def handles_incrementality(self) -> bool: + return True + + def dlt_source(self, uri: str, table: str, **kwargs): + if kwargs.get("incremental_key"): + raise ValueError( + "IMAP takes care of incrementality on its own, you should not provide incremental_key" + ) + + parsed_uri = urlparse(uri) + params = parse_qs(parsed_uri.query) + host = params.get("host") + username = params.get("username") + password = params.get("password") + start_date = params.get("start_date") + + if host is None: + raise MissingValueError("host", "IMAP") + if username is None: + raise MissingValueError("username", "IMAP") + if password is None: + raise MissingValueError("password", "IMAP") + + from omniload.source.imap.adapter import inbox_source + + kwargs: Dict[str, Any] = {} + if start_date is not None: + kwargs["start_date"] = ensure_pendulum_datetime_utc(start_date[0]) + + return inbox_source( + host=host[0], + email_account=username[0], + password=password[0], + **kwargs, + ) diff --git a/src/omniload/source/imap/helpers.py b/src/omniload/source/imap/helpers.py new file mode 100644 index 000000000..60a7e4bfc --- /dev/null +++ b/src/omniload/source/imap/helpers.py @@ -0,0 +1,212 @@ +# Copyright 2022-2026 ScaleVector +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import email +import hashlib +import imaplib +from email.header import decode_header, make_header +from email.message import Message +from email.utils import parsedate_to_datetime +from time import mktime +from typing import Any, Dict, Iterator, Optional, Sequence, Tuple + +from dlt.common import logger, pendulum +from dlt.sources import TDataItems +from dlt.sources.filesystem import FileItem + + +class ImapFileItem(FileItem): + """A Imap file item.""" + + file_hash: str + + +def decode_header_word(v: Any) -> Any: + if not isinstance(v, str): + return v + try: + v = str(make_header(decode_header(v))) + except Exception: # noqa: S110 + pass + + return v + + +def get_message_uids( + client: imaplib.IMAP4_SSL, criterias: Sequence[str] +) -> Optional[TDataItems]: + """Get the message uids from the imap server. + + Args: + client (imaplib.IMAP4_SSL): The imap client. + criterias (Sequence[str]): The search criterias. + + Returns: + Optional[TDataItems]: The list of message uids. + """ + status, messages = client.uid("search", *criterias) + + if status != "OK": + raise client.error(messages[-1]) + + message_uids = messages[0].split() + + if not message_uids: + logger.warning("No emails found.") + return None + + return [{"message_uid": int(message_uid)} for message_uid in message_uids] + + +def get_internal_date(client: imaplib.IMAP4_SSL, message_uid: str) -> Optional[Any]: + """Get the internal date of the email message. + + Parameters: + client (imaplib.IMAP4_SSL): The imap client. + message_uid (str): The uid of the message. + + Returns: + Optional[Any]: The internal date of the email message. + """ + # client.select() + status, data = client.uid("fetch", message_uid, "(INTERNALDATE)") + date = None + + if status != "OK": + raise client.error(data[-1]) + + timestruct = imaplib.Internaldate2tuple(data[0]) + if timestruct is None: + raise ValueError("timestruct is None") + date = pendulum.from_timestamp(mktime(timestruct)) + return date + + +def extract_email_info(msg: Message, include_body: bool = False) -> Dict[str, Any]: + """Extract the email information from the email message. + + Parameters: + msg (Message): The email message object. + include_body (bool, optional): If true, the body of the email will be included. + + Returns: + Dict[str, Any]: The email information. + """ + email_data = dict(msg) + dt = parsedate_to_datetime(msg["Date"]) + dt_pendulum = pendulum.instance(dt) + email_data["Date"] = dt_pendulum + email_data["content_type"] = msg.get_content_type() + if include_body: + email_data["body"] = get_email_body(msg) + + return { + k: decode_header_word(v) + for k, v in email_data.items() + if not k.startswith(("X-", "ARC-", "DKIM-")) + } + + +def get_message_with_internal_date( + client: imaplib.IMAP4_SSL, message_uid: str +) -> Tuple[Message, pendulum.DateTime]: + """Get the email message and internal date from the imap server. + + Parameters: + client (imaplib.IMAP4_SSL): The imap client. + message_uid (str): The uid of the message. + + Returns: + Tuple[Message, pendulum.DateTime]: The email message object and internal date as pendulum DateTime + """ + status, data = client.uid("fetch", message_uid, "(RFC822 INTERNALDATE)") + + if status == "OK": + try: + raw_email = data[0][1] + except (IndexError, TypeError): + raise Exception(f"Error getting content of email with uid {message_uid}.") + else: + raise client.error(data[-1]) + + msg = email.message_from_bytes(raw_email) + # FIXME: This was referred to as `data[1]` before, which seems to be wrong? + timestruct = imaplib.Internaldate2tuple(data[0][0]) + if timestruct is None: + raise ValueError("timestruct is None") + return msg, pendulum.from_timestamp(mktime(timestruct)) + + +def extract_attachments( + message: Message, filter_by_mime_type: Optional[Sequence[str]] = None +) -> Iterator[ImapFileItem]: + """Extract the attachments from the email message. + + Parameters: + message (Message): The email message object. + filter_by_mime_type (str): The mime type to filter the attachments. + + Returns: + Iterable[ImapFileItem]: The attachments. + """ + + filter_by_mime_type = filter_by_mime_type or [] + + for part in message.walk(): + content_type = part.get_content_type() + content_disposition = part.get_content_disposition() + + # Checks if the content is an attachment + if not content_disposition or content_disposition.lower() != "attachment": + continue + + # Checks if the mime type is in the filter list + if filter_by_mime_type and content_type not in filter_by_mime_type: + continue + + file_name = part.get_filename() + if file_name is None: + raise ValueError("file_name is None") + + file_md = ImapFileItem( # type: ignore + file_name=file_name, + mime_type=content_type, + file_content=part.get_payload(decode=True), # ty: ignore[invalid-argument-type] + ) + file_md["file_hash"] = hashlib.sha256(file_md["file_content"]).hexdigest() + file_md["size_in_bytes"] = len(file_md["file_content"]) + + yield file_md + + +def get_email_body(msg: Message) -> str: + """ + Get the body of the email message. + + Parameters: + msg (Message): The email message object. + + Returns: + str: The email body as a string. + """ + body = "" + if msg.is_multipart(): + for part in msg.walk(): + content_type = part.get_content_type() + if content_type == "text/plain": + body += part.get_payload(decode=True).decode(errors="ignore") # ty: ignore[unresolved-attribute] + else: + body = msg.get_payload(decode=True).decode(errors="ignore") # ty: ignore[unresolved-attribute] + + return body diff --git a/src/omniload/source/imap/settings.py b/src/omniload/source/imap/settings.py new file mode 100644 index 000000000..5d67a9da0 --- /dev/null +++ b/src/omniload/source/imap/settings.py @@ -0,0 +1,19 @@ +# Copyright 2022-2026 ScaleVector +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dlt.common import pendulum + +GMAIL_GROUP = None +DEFAULT_START_DATE = pendulum.datetime(1970, 1, 1) +DEFAULT_CHUNK_SIZE = 100 diff --git a/tests/saas/test_imap.py b/tests/saas/test_imap.py new file mode 100644 index 000000000..923e236cc --- /dev/null +++ b/tests/saas/test_imap.py @@ -0,0 +1,94 @@ +import imaplib +import os +import time +from email.message import Message + +import duckdb +import pytest +from testcontainers.core.container import DockerContainer +from testcontainers.core.wait_strategies import PortWaitStrategy + +from tests.util import invoke_ingest_command + +pytestmark = pytest.mark.integration + + +class DovecotContainer(DockerContainer): + """ + A Testcontainer for Dovecot IMAP. + """ + + DOVECOT_VERSION = os.environ.get("DOVECOT_VERSION", "2.4.4") + + def __init__( + self, + image: str = f"docker.io/dovecot/dovecot:{DOVECOT_VERSION}", + imap_port: int = 31993, + **kwargs, + ) -> None: + super().__init__(image=image, **kwargs) + self.imap_port = imap_port + self.with_env("USER_PASSWORD", "secret") + # TODO: Because the connector will always connect to port 993, + # we need to use port _binding_ here. + self.with_bind_ports(f"{self.imap_port}/tcp", 993) + self.waiting_for(PortWaitStrategy(self.imap_port)) + + +@pytest.fixture +def dovecot(): + """Fixture for providing a Dovecot server.""" + container = DovecotContainer() + container.start() + # TODO: Get rid of `time.sleep`. + time.sleep(1) + try: + host = container.get_container_host_ip() + yield host + finally: + container.stop() + + +@pytest.fixture +def dovecot_with_message(dovecot): + """Fixture for providing a Dovecot server including a single message in `INBOX`.""" + imap = imaplib.IMAP4_SSL(host=dovecot, port=993) + imap.login("hotzenplotz", "secret") + + new_message = Message() + new_message["From"] = "hello@example.org" + new_message["Subject"] = "Test mail." + new_message["Date"] = "Thu, 20 Aug 2026 11:35:19 +0200" + new_message.set_payload("This is the message.") + + # TODO: Alternatively use given mailbox name than just `INBOX`. + # imap.create("testdrive") + imap.append( + "INBOX", + "", + imaplib.Time2Internaldate(time.time()), + str(new_message).encode("utf-8"), + ) + imap.logout() + yield dovecot + + +def test_imap_basic(dovecot_with_message, tmp_path): + """Verify a basic ingest from an IMAP mailbox.""" + + abs_db_path = tmp_path / "test_imap.duckdb" + uri = f"duckdb:///{abs_db_path}" + + result = invoke_ingest_command( + f"imap://?host={dovecot_with_message}&username=hotzenplotz&password=secret", + "", + uri, + "raw.imap", + ) + assert result.exit_code == 0, result.output + + conn = duckdb.connect(abs_db_path) + result = conn.sql("select count(*) from raw.imap").fetchone() + assert result is not None, "Database result is empty" + assert result[0] > 0, "No records found in table raw.imap" + conn.close()