fix(workflows): keep one mail from label per email sending domain

posthog/api/integration.py

@@ -2705,7 +2705,7 @@ def _get_email_integration(self) -> EmailIntegration:
         instance = self.get_object()
         if instance.kind != "email":
             raise ValidationError("This endpoint is only supported for email integrations")
-        return EmailIntegration(instance)
+        return EmailIntegration(instance, acting_user=cast(User, self.request.user))
 
     @validated_request(
         query_serializer=DomainConnectCheckQuerySerializer,
@@ -2748,7 +2748,7 @@ def domain_connect_apply_url(self, request: ValidatedRequest, *args: Any, **kwar
         if context == "email":
             integration_id = data["integration_id"]
             try:
-                resolved = resolve_email_context(integration_id, self.team_id)
+                resolved = resolve_email_context(integration_id, self.team_id, acting_user=cast(User, request.user))
             except Integration.DoesNotExist:
                 raise NotFound("No email sender with this integration_id exists in this project.")
             except ValueError as e:

posthog/api/test/test_integration.py

@@ -818,6 +818,47 @@ def test_email_verify_updates_all_other_integrations_with_same_domain(
         assert integration2.config["verified"]
         assert not integrationOtherDomain.config["verified"]
 
+    @patch("products.workflows.backend.facade.api.update_ses_mail_from_subdomain")
+    def test_editing_the_label_shares_it_when_the_flag_is_on_for_the_requesting_user(
+        self, _mock_update_mail_from_subdomain, client: HttpClient
+    ):
+        admin = User.objects.create_and_join(
+            self.organization, "admin@example.com", "test", level=OrganizationMembership.Level.ADMIN
+        )
+        edited, sibling = (
+            Integration.objects.create(
+                team=self.team,
+                kind="email",
+                integration_id=email,
+                config={"email": email, "domain": "example.com", "mail_from_subdomain": "feedback", "provider": "ses"},
+            )
+            for email in ["edited@example.com", "sibling@example.com"]
+        )
+        client.force_login(admin)
+
+        with patch(
+            "posthoganalytics.feature_enabled",
+            side_effect=lambda key, distinct_id, **_kwargs: (
+                key == "workflows-email-domain-agent-setup" and distinct_id == admin.distinct_id
+            ),
+        ):
+            response = client.patch(
+                f"/api/environments/{self.team.pk}/integrations/{edited.pk}/email",
+                {
+                    "config": {
+                        "email": "edited@example.com",
+                        "name": "Edited",
+                        "provider": "ses",
+                        "mail_from_subdomain": "bounce",
+                    }
+                },
+                content_type="application/json",
+            )
+
+        assert response.status_code == status.HTTP_200_OK, response.json()
+        sibling.refresh_from_db()
+        assert sibling.config["mail_from_subdomain"] == "bounce"
+
 
 class TestDatabricksIntegration:
     @pytest.fixture(autouse=True)

posthog/domain_connect.py

@@ -9,6 +9,7 @@
 
 import base64
 import logging
+from typing import TYPE_CHECKING
 from urllib.parse import urlencode
 
 from django.conf import settings
@@ -26,6 +27,9 @@
 
 from posthog.dataclasses import frozen
 
+if TYPE_CHECKING:
+    from posthog.models.user import User
+
 logger = logging.getLogger(__name__)
 
 # Maps _domainconnect TXT record values (provider endpoints) to display names.
@@ -234,7 +238,7 @@ class DomainConnectContext:
 EMAIL_TEMPLATE_GROUPS_WITHOUT_DMARC: tuple[str, ...] = ("verification", "dkim", "spf", "mailfrom")
 
 
-def resolve_email_context(integration_id: int, team_id: int) -> DomainConnectContext:
+def resolve_email_context(integration_id: int, team_id: int, acting_user: "User | None" = None) -> DomainConnectContext:
     """Resolve Domain Connect parameters for an email integration.
 
     Triggers SES verification to get current tokens, then extracts the
@@ -244,12 +248,12 @@ def resolve_email_context(integration_id: int, team_id: int) -> DomainConnectCon
 
     instance = Integration.objects.get(id=integration_id, team_id=team_id, kind="email")
 
-    email_integration = EmailIntegration(instance)
+    email_integration = EmailIntegration(instance, acting_user=acting_user)
     verification_result = email_integration.verify()
 
     dns_records = verification_result.get("dnsRecords", [])
     domain_parts = extract_root_domain_and_host(instance.config.get("domain", ""))
-    mail_from_subdomain = instance.config.get("mail_from_subdomain", "feedback")
+    mail_from_subdomain = email_integration.mail_from_subdomain
 
     verify_token = ""
     dkim_tokens: list[str] = []

posthog/models/integration/email.py

@@ -1,9 +1,12 @@
 """Native email-sending integration (SES / maildev) and its cleanup signal."""
 
+from collections.abc import Iterable, Iterator
+from contextlib import AbstractContextManager, contextmanager, nullcontext
+from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
 from django.conf import settings
-from django.db import models, transaction
+from django.db import connection, models, transaction
 from django.dispatch import receiver
 
 from disposable_email_domains import blocklist as disposable_email_domains_list
@@ -12,6 +15,7 @@
 
 from posthog.models.team.team import Team
 from posthog.models.user import User
+from posthog.ph_client import feature_enabled_or_false
 from posthog.plugins.plugin_server_api import reload_integrations_on_workers
 
 from . import model
@@ -20,22 +24,69 @@
     from products.workflows.backend.facade.contracts import EmailDomainVerification
 
 
+DEFAULT_MAIL_FROM_SUBDOMAIN = "feedback"
+EMAIL_DOMAIN_AGENT_SETUP_FLAG = "workflows-email-domain-agent-setup"
+
+
 class EmailIntegration:
     integration: model.Integration
 
-    def __init__(self, integration: model.Integration) -> None:
+    def __init__(self, integration: model.Integration, acting_user: User | None = None) -> None:
         if integration.kind != "email":
             raise Exception("EmailIntegration init called with Integration with wrong 'kind'")
         self.integration = integration
+        self.acting_user = acting_user
+
+    @staticmethod
+    def _shares_domain_label_for(team: Team, acting_user: User | None) -> bool:
+        return feature_enabled_or_false(
+            EMAIL_DOMAIN_AGENT_SETUP_FLAG,
+            str(acting_user.distinct_id) if acting_user else str(team.uuid),
+            groups={"organization": str(team.organization_id), "project": str(team.uuid)},
+            group_properties={"organization": {"id": str(team.organization_id)}, "project": {"id": str(team.uuid)}},
+            send_feature_flag_events=False,
+        )
+
+    @cached_property
+    def _shares_domain_label(self) -> bool:
+        return self._shares_domain_label_for(self.integration.team, self.acting_user)
+
+    @property
+    def mail_from_subdomain(self) -> str:
+        if self._shares_domain_label:
+            return self._shared_mail_from_subdomain(self.integration.config)
+        return self.integration.config.get("mail_from_subdomain", DEFAULT_MAIL_FROM_SUBDOMAIN)
+
+    @staticmethod
+    def _shared_mail_from_subdomain(config: dict[str, Any]) -> str:
+        return config.get("mail_from_subdomain") or DEFAULT_MAIL_FROM_SUBDOMAIN
 
     @classmethod
     def create_native_integration(
-        cls, config: dict, team_id: int, organization_id: str, created_by: User | None = None
+        cls, config: dict[str, Any], team_id: int, organization_id: str, created_by: User | None = None
+    ) -> model.Integration:
+        shares_domain_label = cls._shares_domain_label_for(Team.objects.get(id=team_id), created_by)
+        domain_access: AbstractContextManager[None] = (
+            cls._exclusive_domain_access(cls._email_domain(config["email"])) if shares_domain_label else nullcontext()
+        )
+        with domain_access:
+            return cls._create_native_integration(
+                config, team_id, organization_id, created_by, shares_domain_label=shares_domain_label
+            )
+
+    @classmethod
+    def _create_native_integration(
+        cls,
+        config: dict[str, Any],
+        team_id: int,
+        organization_id: str,
+        created_by: User | None,
+        *,
+        shares_domain_label: bool,
     ) -> model.Integration:
         email_address: str = config["email"].lower()
         name: str = config["name"]
-        domain: str = email_address.split("@")[1]
-        mail_from_subdomain: str = config.get("mail_from_subdomain", "feedback")
+        domain: str = cls._email_domain(email_address)
         provider: str = config.get("provider", "ses")
 
         if domain in free_email_domains_list or domain in disposable_email_domains_list:
@@ -50,6 +101,12 @@ def create_native_integration(
                     f"An email integration with domain {domain} already exists in another organization. Try a different domain or contact support if you believe this is a mistake."
                 )
 
+        mail_from_subdomain = (
+            cls._domain_wide_mail_from_subdomain(domain, config.get("mail_from_subdomain"), same_domain_integrations)
+            if shares_domain_label
+            else config.get("mail_from_subdomain", DEFAULT_MAIL_FROM_SUBDOMAIN)
+        )
+
         # Create domain in the appropriate provider
         if provider == "ses":
             from products.workflows.backend.facade.api import (
@@ -91,14 +148,60 @@ def create_native_integration(
 
         return integration
 
-    def update_native_integration(self, config: dict, team_id: int) -> model.Integration:
+    @staticmethod
+    def _email_domain(email_address: str) -> str:
+        return email_address.lower().split("@")[1]
+
+    @staticmethod
+    @contextmanager
+    def _exclusive_domain_access(domain: str) -> Iterator[None]:
+        with transaction.atomic():
+            with connection.cursor() as cursor:
+                cursor.execute("SELECT pg_advisory_xact_lock(hashtextextended(%s, 0))", [f"email-domain:{domain}"])
+            yield
+
+    @contextmanager
+    def _domain_guard(self) -> Iterator[None]:
+        if not self._shares_domain_label:
+            yield
+            return
+        with self._exclusive_domain_access(self.integration.config["domain"]):
+            self.integration.refresh_from_db(fields=["config"])
+            yield
+
+    @staticmethod
+    def _domain_wide_mail_from_subdomain(
+        domain: str, requested_subdomain: str | None, same_domain_integrations: Iterable[model.Integration]
+    ) -> str:
+        domain_subdomains = {
+            EmailIntegration._shared_mail_from_subdomain(integration.config) for integration in same_domain_integrations
+        }
+        if not domain_subdomains:
+            return requested_subdomain or DEFAULT_MAIL_FROM_SUBDOMAIN
+        if len(domain_subdomains) > 1:
+            raise ValidationError(
+                f"Senders on {domain} use different MAIL FROM subdomains. "
+                f"Open an existing sender on {domain} in any project of this organization and save the subdomain you want. "
+                "Then add this sender."
+            )
+        domain_subdomain = domain_subdomains.pop()
+        if requested_subdomain and requested_subdomain != domain_subdomain:
+            raise ValidationError(
+                f"{domain} already uses the MAIL FROM subdomain '{domain_subdomain}'. "
+                f"Use '{domain_subdomain}' for this sender. To change it for every sender on {domain}, edit an existing sender."
+            )
+        return domain_subdomain
+
+    def update_native_integration(self, config: dict[str, Any], team_id: int) -> model.Integration:
+        with self._domain_guard():
+            return self._update_native_integration(config)
+
+    def _update_native_integration(self, config: dict[str, Any]) -> model.Integration:
         provider = self.integration.config.get("provider")
         domain = self.integration.config.get("domain")
         # Only name and mail_from_subdomain can be updated
         name: str = config.get("name", self.integration.config.get("name"))
-        mail_from_subdomain: str = config.get(
-            "mail_from_subdomain", self.integration.config.get("mail_from_subdomain", "feedback")
-        )
+        mail_from_subdomain = self._edited_mail_from_subdomain(config.get("mail_from_subdomain"))
 
         # Update domain in the appropriate provider
         if provider == "ses":
@@ -119,17 +222,42 @@ def update_native_integration(self, config: dict, team_id: int) -> model.Integra
             }
         )
         self.integration.save()
+        if self._shares_domain_label:
+            self._share_mail_from_subdomain_with_domain_senders(domain, mail_from_subdomain)
 
         return self.integration
 
     def reject_address_change(self, email_address: str) -> None:
         if email_address.lower() != self.integration.config.get("email", "").lower():
             raise ValidationError(f"The sender address cannot change. Create a new sender for {email_address} instead.")
 
+    def _edited_mail_from_subdomain(self, requested_subdomain: str | None) -> str:
+        if self._shares_domain_label:
+            return requested_subdomain or self.mail_from_subdomain
+        return self.mail_from_subdomain if requested_subdomain is None else requested_subdomain
+
+    def _share_mail_from_subdomain_with_domain_senders(self, domain: str, mail_from_subdomain: str) -> None:
+        domain_senders = (
+            model.Integration.objects.select_for_update(of=("self",))
+            .filter(
+                kind="email",
+                config__domain=domain,
+                team__organization_id=self.integration.team.organization_id,
+            )
+            .exclude(pk=self.integration.pk)
+        )
+        for sender in domain_senders:
+            sender.config["mail_from_subdomain"] = mail_from_subdomain
+            sender.save(update_fields=["config"])
+
     def verify(self) -> "EmailDomainVerification":
+        with self._domain_guard():
+            return self._verify()
+
+    def _verify(self) -> "EmailDomainVerification":
         domain = self.integration.config.get("domain")
         provider = self.integration.config.get("provider", "ses")
-        mail_from_subdomain = self.integration.config.get("mail_from_subdomain", "feedback")
+        mail_from_subdomain = self.mail_from_subdomain
 
         verification_result: EmailDomainVerification
         # Use the appropriate provider for verification
@@ -165,8 +293,9 @@ def verify(self) -> "EmailDomainVerification":
                 integration.config["verified"] = True
                 integration.save()
 
-            reload_integrations_on_workers(
-                self.integration.team_id, [integration.id for integration in all_integrations_for_domain]
+            verified_integration_ids = [integration.id for integration in all_integrations_for_domain]
+            transaction.on_commit(
+                lambda: reload_integrations_on_workers(self.integration.team_id, verified_integration_ids)
             )
 
         return verification_result