Repository: PrefectHQ/prefect

Issue:
Migrate `prefect-email` to core



Discussion:

Resolving patch:
diff --git a/src/integrations/prefect-email/prefect_email/__init__.py b/src/integrations/prefect-email/prefect_email/__init__.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-email/prefect_email/__init__.py
@@ -0,0 +1,7 @@
+from . import _version
+from .credentials import EmailServerCredentials, SMTPServer, SMTPType
+from .message import email_send_message
+
+__version__ = _version.__version__
+
+__all__ = ["EmailServerCredentials", "SMTPServer", "SMTPType", "email_send_message"]
diff --git a/src/integrations/prefect-email/prefect_email/credentials.py b/src/integrations/prefect-email/prefect_email/credentials.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-email/prefect_email/credentials.py
@@ -0,0 +1,198 @@
+"""Credential classes used to perform authenticated interactions with email services"""
+
+import ssl
+from enum import Enum
+from functools import partial
+from smtplib import SMTP, SMTP_SSL
+from typing import Optional, Union
+
+from pydantic import VERSION as PYDANTIC_VERSION
+
+from prefect.blocks.core import Block
+
+if PYDANTIC_VERSION.startswith("2."):
+    from pydantic.v1 import Field, SecretStr, validator
+else:
+    from pydantic import Field, SecretStr, validator
+
+
+class SMTPType(Enum):
+    """
+    Protocols used to secure email transmissions.
+    """
+
+    SSL = 465
+    STARTTLS = 587
+    INSECURE = 25
+
+
+class SMTPServer(Enum):
+    """
+    Server used to send email.
+    """
+
+    AOL = "smtp.aol.com"
+    ATT = "smtp.mail.att.net"
+    COMCAST = "smtp.comcast.net"
+    ICLOUD = "smtp.mail.me.com"
+    GMAIL = "smtp.gmail.com"
+    OUTLOOK = "smtp-mail.outlook.com"
+    YAHOO = "smtp.mail.yahoo.com"
+
+
+def _cast_to_enum(obj: Union[str, SMTPType], enum: Enum, restrict: bool = False):
+    """
+    Casts string to an enum member, if valid; else returns the input
+    obj, or raises a ValueError if restrict.
+
+    Args:
+        obj: A member's name of the enum.
+        enum: An Enum class.
+        restrict: Whether to only allow passing members from the enum.
+
+    Returns:
+        A member of the enum.
+    """
+    if isinstance(obj, enum):
+        # if already an enum member, continue
+        return obj
+    valid_enums = list(enum.__members__)
+
+    # capitalize and try to match an enum member
+    if obj.upper() not in valid_enums:
+        if restrict:
+            raise ValueError(f"Must be one of {valid_enums}; got {obj!r}")
+        else:
+            # primarily used for SMTPServer so that users
+            # can pass in their custom servers not listed
+            # as one of the SMTPServer Enum values.
+            return obj
+    else:
+        return getattr(enum, obj.upper())
+
+
+class EmailServerCredentials(Block):
+    """
+    Block used to manage generic email server authentication.
+    It is recommended you use a
+    [Google App Password](https://support.google.com/accounts/answer/185833)
+    if you use Gmail.
+
+    Attributes:
+        username: The username to use for authentication to the server.
+            Unnecessary if SMTP login is not required.
+        password: The password to use for authentication to the server.
+            Unnecessary if SMTP login is not required.
+        smtp_server: Either the hostname of the SMTP server, or one of the
+            keys from the built-in SMTPServer Enum members, like "gmail".
+        smtp_type: Either "SSL", "STARTTLS", or "INSECURE".
+        smtp_port: If provided, overrides the smtp_type's default port number.
+
+    Example:
+        Load stored email server credentials:
+        ```python
+        from prefect_email import EmailServerCredentials
+        email_credentials_block = EmailServerCredentials.load("BLOCK_NAME")
+        ```
+    """  # noqa E501
+
+    _block_type_name = "Email Server Credentials"
+    _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/82bc6ed16ca42a2252a5512c72233a253b8a58eb-250x250.png"  # noqa
+    _documentation_url = "https://prefecthq.github.io/prefect-email/credentials/#prefect_email.credentials.EmailServerCredentials"  # noqa
+
+    username: Optional[str] = Field(
+        default=None,
+        description=(
+            "The username to use for authentication to the server. "
+            "Unnecessary if SMTP login is not required."
+        ),
+    )
+    password: SecretStr = Field(
+        default_factory=partial(SecretStr, ""),
+        description=(
+            "The password to use for authentication to the server. "
+            "Unnecessary if SMTP login is not required."
+        ),
+    )
+    smtp_server: Union[SMTPServer, str] = Field(
+        default=SMTPServer.GMAIL,
+        description=(
+            "Either the hostname of the SMTP server, or one of the "
+            "keys from the built-in SMTPServer Enum members, like 'gmail'."
+        ),
+        title="SMTP Server",
+    )
+    smtp_type: Union[SMTPType, str] = Field(
+        default=SMTPType.SSL,
+        description=("Either 'SSL', 'STARTTLS', or 'INSECURE'."),
+        title="SMTP Type",
+    )
+    smtp_port: Optional[int] = Field(
+        default=None,
+        description=("If provided, overrides the smtp_type's default port number."),
+        title="SMTP Port",
+    )
+
+    @validator("smtp_server", pre=True)
+    def _cast_smtp_server(cls, value):
+        """
+        Cast the smtp_server to an SMTPServer Enum member, if valid.
+        """
+        return _cast_to_enum(value, SMTPServer)
+
+    @validator("smtp_type", pre=True)
+    def _cast_smtp_type(cls, value):
+        """
+        Cast the smtp_type to an SMTPType Enum member, if valid.
+        """
+        if isinstance(value, int):
+            return SMTPType(value)
+        return _cast_to_enum(value, SMTPType, restrict=True)
+
+    def get_server(self) -> SMTP:
+        """
+        Gets an authenticated SMTP server.
+
+        Returns:
+            SMTP: An authenticated SMTP server.
+
+        Example:
+            Gets a GMail SMTP server through defaults.
+            ```python
+            from prefect import flow
+            from prefect_email import EmailServerCredentials
+
+            @flow
+            def example_get_server_flow():
+                email_server_credentials = EmailServerCredentials(
+                    username="username@gmail.com",
+                    password="password",
+                )
+                server = email_server_credentials.get_server()
+                return server
+
+            example_get_server_flow()
+            ```
+        """
+        smtp_server = self.smtp_server
+        if isinstance(smtp_server, SMTPServer):
+            smtp_server = smtp_server.value
+
+        smtp_type = self.smtp_type
+        smtp_port = self.smtp_port
+        if smtp_port is None:
+            smtp_port = smtp_type.value
+
+        if smtp_type == SMTPType.INSECURE:
+            server = SMTP(smtp_server, smtp_port)
+        else:
+            context = ssl.create_default_context()
+            if smtp_type == SMTPType.SSL:
+                server = SMTP_SSL(smtp_server, smtp_port, context=context)
+            elif smtp_type == SMTPType.STARTTLS:
+                server = SMTP(smtp_server, smtp_port)
+                server.starttls(context=context)
+            if self.username is not None:
+                server.login(self.username, self.password.get_secret_value())
+
+        return server
diff --git a/src/integrations/prefect-email/prefect_email/message.py b/src/integrations/prefect-email/prefect_email/message.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-email/prefect_email/message.py
@@ -0,0 +1,121 @@
+"""
+Tasks for interacting with email message services
+"""
+# Do NOT rename this module to `email.py` as it
+# will conflict with the standard library!
+
+import os
+from email import encoders
+from email.mime.base import MIMEBase
+from email.mime.multipart import MIMEMultipart
+from email.mime.text import MIMEText
+from functools import partial
+from typing import TYPE_CHECKING, List, Optional, Union
+
+from anyio import to_thread
+
+from prefect import task
+
+if TYPE_CHECKING:
+    from prefect_email import EmailServerCredentials
+
+
+@task
+async def email_send_message(
+    subject: str,
+    msg: str,
+    email_server_credentials: "EmailServerCredentials",
+    msg_plain: Optional[str] = None,
+    email_from: Optional[str] = None,
+    email_to: Optional[Union[str, List[str]]] = None,
+    email_to_cc: Optional[Union[str, List[str]]] = None,
+    email_to_bcc: Optional[Union[str, List[str]]] = None,
+    attachments: Optional[List[str]] = None,
+):
+    """
+    Sends an email message from an authenticated email service over SMTP.
+    Sending messages containing HTML code is supported - the default MIME
+    type is set to the text/html.
+
+    Args:
+        subject: The subject line of the email.
+        msg: The contents of the email, added as html; can be used in
+            combination with msg_plain.
+        msg_plain: The contents of the email as plain text,
+            can be used in combination with msg.
+        email_to: The email addresses to send the message to, separated by commas.
+            If a list is provided, will join the items, separated by commas.
+        email_to_cc: Additional email addresses to send the message to as cc,
+            separated by commas. If a list is provided, will join the items,
+            separated by commas.
+        email_to_bcc: Additional email addresses to send the message to as bcc,
+            separated by commas. If a list is provided, will join the items,
+            separated by commas.
+        attachments: Names of files that should be sent as attachment.
+
+    Returns:
+        MimeText: The MIME Multipart message of the email.
+
+    Example:
+        Sends a notification email to someone@gmail.com.
+        ```python
+        from prefect import flow
+        from prefect_email import EmailServerCredentials, email_send_message
+
+        @flow
+        def example_email_send_message_flow():
+            email_server_credentials = EmailServerCredentials(
+                username="username@email.com",
+                password="password",
+            )
+            subject = email_send_message(
+                email_server_credentials=email_server_credentials,
+                subject="Example Flow Notification",
+                msg="This proves email_send_message works!",
+                email_to="someone@email.com",
+            )
+            return subject
+
+        example_email_send_message_flow()
+        ```
+    """
+    message = MIMEMultipart()
+    message["Subject"] = subject
+    message["From"] = email_from or email_server_credentials.username
+
+    email_to_dict = {"To": email_to, "Cc": email_to_cc, "Bcc": email_to_bcc}
+    if all(val is None for val in email_to_dict.values()):
+        raise ValueError(
+            "One of email_to, email_to_cc, or email_to_bcc must be specified"
+        )
+
+    for key, val in email_to_dict.items():
+        if isinstance(val, list):
+            val = ", ".join(val)
+        message[key] = val
+
+    # First add the message in plain text, then the HTML version;
+    # email clients try to render the last part first
+    if msg_plain:
+        message.attach(MIMEText(msg_plain, "plain"))
+    if msg:
+        message.attach(MIMEText(msg, "html"))
+
+    for filepath in attachments or []:
+        with open(filepath, "rb") as attachment:
+            part = MIMEBase("application", "octet-stream")
+            part.set_payload(attachment.read())
+
+        encoders.encode_base64(part)
+        filename = os.path.basename(filepath)
+        part.add_header(
+            "Content-Disposition",
+            f"attachment; filename= {filename}",
+        )
+        message.attach(part)
+
+    with email_server_credentials.get_server() as server:
+        partial_send_message = partial(server.send_message, message)
+        await to_thread.run_sync(partial_send_message)
+
+    return message


Test patch:
diff --git a/src/integrations/prefect-azure/tests/test_version.py b/src/integrations/prefect-azure/tests/test_version.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-azure/tests/test_version.py
@@ -0,0 +1,9 @@
+from packaging.version import Version
+
+
+def test_version():
+    from prefect_azure import __version__
+
+    assert isinstance(__version__, str)
+    assert Version(__version__)
+    assert __version__.startswith("0.")
diff --git a/src/integrations/prefect-email/tests/attachment.txt b/src/integrations/prefect-email/tests/attachment.txt
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-email/tests/attachment.txt
@@ -0,0 +1 @@
+This is a test attachment!
\ No newline at end of file
diff --git a/src/integrations/prefect-email/tests/conftest.py b/src/integrations/prefect-email/tests/conftest.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-email/tests/conftest.py
@@ -0,0 +1,50 @@
+from unittest.mock import MagicMock
+
+import pytest
+
+from prefect.testing.utilities import prefect_test_harness
+
+
+@pytest.fixture(autouse=True)
+def prefect_db():
+    with prefect_test_harness():
+        yield
+
+
+class EmailServerMethodsMock:
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        return False
+
+    def send_message(self, message):
+        return message
+
+
+@pytest.fixture
+def email_server_credentials():
+    email_server_credentials = MagicMock(username="someone@email.com")
+    email_server_credentials.get_server.side_effect = lambda: EmailServerMethodsMock()
+    return email_server_credentials
+
+
+class SMTPMock(MagicMock):
+    def __init__(self, server, port, context=None):
+        super().__init__()
+        self.server = server
+        self.port = port
+        self.context = context
+
+    def login(self, username, password):
+        self.username = username
+        self.password = password
+
+    def starttls(self, context=None):
+        self.context = context
+
+
+@pytest.fixture
+def smtp(monkeypatch):
+    monkeypatch.setattr("prefect_email.credentials.SMTP", SMTPMock)
+    monkeypatch.setattr("prefect_email.credentials.SMTP_SSL", SMTPMock)
diff --git a/src/integrations/prefect-email/tests/test_block_standards.py b/src/integrations/prefect-email/tests/test_block_standards.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-email/tests/test_block_standards.py
@@ -0,0 +1,23 @@
+import pytest
+
+from prefect.blocks.core import Block
+from prefect.testing.standard_test_suites import BlockStandardTestSuite
+from prefect.utilities.dispatch import get_registry_for_type
+from prefect.utilities.importtools import to_qualified_name
+
+
+def find_module_blocks():
+    blocks = get_registry_for_type(Block)
+    module_blocks = [
+        block
+        for block in blocks.values()
+        if to_qualified_name(block).startswith("prefect_email")
+    ]
+    return module_blocks
+
+
+@pytest.mark.parametrize("block", find_module_blocks())
+class TestAllBlocksAdhereToStandards(BlockStandardTestSuite):
+    @pytest.fixture
+    def block(self, block):
+        return block
diff --git a/src/integrations/prefect-email/tests/test_credentials.py b/src/integrations/prefect-email/tests/test_credentials.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-email/tests/test_credentials.py
@@ -0,0 +1,93 @@
+import pytest
+from prefect_email.credentials import (
+    EmailServerCredentials,
+    SMTPServer,
+    SMTPType,
+    _cast_to_enum,
+)
+
+
+@pytest.mark.parametrize("obj", ["gmail", "Gmail", "GMAIL", SMTPServer.GMAIL])
+def test_cast_to_enum(obj):
+    assert _cast_to_enum(obj, SMTPServer, restrict=False) == SMTPServer.GMAIL
+
+
+@pytest.mark.parametrize("obj", ["smtp.prefect.io", "smtp.gmail.com"])
+def test_cast_to_enum_no_restrict_server(obj):
+    assert _cast_to_enum(obj, SMTPServer, restrict=False) == obj
+
+
+@pytest.mark.parametrize("obj", ["ssl", "Ssl", "SSL", SMTPType.SSL])
+def test_cast_to_enum_restrict_type(obj):
+    assert _cast_to_enum(obj, SMTPType, restrict=True) == SMTPType.SSL
+
+
+def test_cast_to_enum_restrict_error():
+    with pytest.raises(ValueError):
+        _cast_to_enum("Invalid", SMTPType, restrict=True)
+
+
+@pytest.mark.parametrize(
+    "smtp_server", ["gmail", "GMAIL", "smtp.gmail.com", "SMTP.GMAIL.COM"]
+)
+@pytest.mark.parametrize(
+    "smtp_type, smtp_port",
+    [
+        ("SSL", 465),
+        ("STARTTLS", 587),
+        ("ssl", 465),
+        ("StartTLS", 587),
+        ("insecure", 25),
+    ],
+)
+def test_email_server_credentials_get_server(smtp_server, smtp_type, smtp_port, smtp):
+    server = EmailServerCredentials(
+        username="username",
+        password="password",
+        smtp_server=smtp_server,
+        smtp_type=smtp_type,
+    ).get_server()
+    if smtp_type.lower() != "insecure":
+        assert server.username == "username"
+        assert server.password == "password"
+    assert server.server.lower() == "smtp.gmail.com"
+    assert server.port == smtp_port
+
+
+def test_email_server_credentials_get_server_error(smtp):
+    with pytest.raises(ValueError):
+        EmailServerCredentials(
+            username="username", password="password", smtp_type="INVALID"
+        ).get_server()
+
+
+def test_email_server_credentials_override_smtp_port(smtp):
+    server = EmailServerCredentials(
+        username="username",
+        password="password",
+        smtp_server=SMTPServer.GMAIL,
+        smtp_type=SMTPType.SSL,
+        smtp_port=1234,
+    ).get_server()
+    assert server.port == 1234
+
+
+def test_email_server_credentials_defaults(smtp):
+    server = EmailServerCredentials().get_server()
+    assert server.server.lower() == "smtp.gmail.com"
+    assert server.port == 465
+
+
+@pytest.mark.parametrize("smtp_type", [SMTPType.STARTTLS, "STARTTLS", 587])
+def test_email_service_credentials_roundtrip_smtp_type_enum(smtp, smtp_type):
+    email_server_credentials = EmailServerCredentials(
+        smtp_server="us-smtp-outbound-1.mimecast.com",
+        smtp_type=smtp_type,
+        username="username",
+        password="password",
+    )
+    email_server_credentials.save("email-credentials", overwrite=True)
+    credentials = EmailServerCredentials.load("email-credentials")
+    assert credentials.smtp_type == SMTPType.STARTTLS
+    server = credentials.get_server()
+    assert server.port == 587
diff --git a/src/integrations/prefect-email/tests/test_message.py b/src/integrations/prefect-email/tests/test_message.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-email/tests/test_message.py
@@ -0,0 +1,75 @@
+import base64
+import pathlib
+
+import pytest
+from prefect_email.message import email_send_message
+
+from prefect import flow
+
+EMAIL_TO = [
+    "someone@email.com",
+    "someone@email.com, someone_else@email.com",
+    ["some_1@email.com", "some_2@email.com"],
+    None,
+]
+EMAIL_TO_CC = [
+    "cc_someone@email.com",
+    "cc_someone@email.com, cc_someone_else@email.com",
+    ["cc_some_1@email.com", "cc_some_2@email.com"],
+    None,
+]
+EMAIL_TO_BCC = [
+    "bcc_someone@email.com",
+    "bcc_someone@email.com, bcc_someone_else@email.com",
+    ["bcc_some_1@email.com", "bcc_some_2@email.com"],
+    None,
+]
+
+
+@pytest.mark.parametrize("email_to", EMAIL_TO)
+@pytest.mark.parametrize("email_to_cc", EMAIL_TO_CC)
+@pytest.mark.parametrize("email_to_bcc", EMAIL_TO_BCC)
+async def test_email_send_message(
+    email_to, email_to_cc, email_to_bcc, email_server_credentials
+):
+    subject = "Example Flow Notification"
+    msg_plain = "This proves msg plain is attached first!"
+    msg = "<h1>This proves msg is attached second!</h1>"
+
+    attachment = pathlib.Path(__file__).parent.absolute() / "attachment.txt"
+    with open(attachment, "rb") as f:
+        attachment_text = f.read()
+
+    @flow
+    async def test_flow():
+        message = await email_send_message(
+            email_server_credentials=email_server_credentials,
+            subject=subject,
+            msg=msg,
+            msg_plain=msg_plain,
+            attachments=[attachment],
+            email_to=email_to,
+            email_to_cc=email_to_cc,
+            email_to_bcc=email_to_bcc,
+        )
+        return message
+
+    email_to_dict = {"To": email_to, "Cc": email_to_cc, "Bcc": email_to_bcc}
+
+    if all(val is None for val in email_to_dict.values()):
+        with pytest.raises(ValueError):
+            await test_flow()
+        return
+
+    message = await test_flow()
+    assert message["Subject"] == subject
+    assert message["From"] == email_server_credentials.username
+    assert message.get_payload()[0].get_payload() == msg_plain
+    assert message.get_payload()[1].get_payload() == msg
+    attachment = message.get_payload()[2].get_payload()
+    assert base64.b64decode(attachment) == attachment_text
+
+    for key, val in email_to_dict.items():
+        if isinstance(val, list):
+            val = ", ".join(val)
+        assert message[key] == val
diff --git a/src/integrations/prefect-email/tests/test_version.py b/src/integrations/prefect-email/tests/test_version.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-email/tests/test_version.py
@@ -0,0 +1,9 @@
+from packaging.version import Version
+
+
+def test_version():
+    from prefect_email import __version__
+
+    assert isinstance(__version__, str)
+    assert Version(__version__)
+    assert __version__.startswith("0.")