feat: add periodic task to update domain statuses (#18014)

Author: miketheman

requirements/tests.in

@@ -4,6 +4,7 @@ freezegun
 pretend
 pytest>=3.0.0
 pytest-icdiff
+pytest-mock
 pytest-postgresql>=3.1.3,<8.0.0
 pytest-randomly
 pytest-socket

requirements/tests.txt

@@ -252,6 +252,7 @@ pytest==8.3.5 \
     # via
     #   -r requirements/tests.in
     #   pytest-icdiff
+    #   pytest-mock
     #   pytest-postgresql
     #   pytest-randomly
     #   pytest-socket
@@ -261,6 +262,10 @@ pytest-icdiff==0.9 \
     --hash=sha256:13aede616202e57fcc882568b64589002ef85438046f012ac30a8d959dac8b75 \
     --hash=sha256:efee0da3bd1b24ef2d923751c5c547fbb8df0a46795553fba08ef57c3ca03d82
     # via -r requirements/tests.in
+pytest-mock==3.14.0 \
+    --hash=sha256:0b72c38033392a5f4621342fe11e9219ac11ec9d375f8e2a0c164539e0d70f6f \
+    --hash=sha256:2719255a1efeceadbc056d6bf3df3d1c5015530fb40cf347c0f9afac88410bd0
+    # via -r requirements/tests.in
 pytest-postgresql==7.0.1 \
     --hash=sha256:7723dfbfc57ea6f6f9876c2828e7b36f8b0e60b6cb040b1ddd444a60eed06e0a \
     --hash=sha256:cbc6a67bbad5128b1f00def8cca5cf597020acc79893723f7a9cb60981b6840f

tests/conftest.py

@@ -550,8 +550,10 @@ def search_service():
 
 
 @pytest.fixture
-def domain_status_service():
-    return account_services.NullDomainStatusService()
+def domain_status_service(mocker):
+    service = account_services.NullDomainStatusService()
+    mocker.spy(service, "get_domain_status")
+    return service
 
 
 class QueryRecorder:

tests/unit/accounts/test_services.py

@@ -1691,7 +1691,7 @@ def __init__(self):
             session=session, client_id="some_client_id"
         )
 
-        assert svc.get_domain_status("example.com") == []
+        assert svc.get_domain_status("example.com") is None
         assert session.get.calls == [
             pretend.call(
                 "https://api.domainr.com/v2/status",

tests/unit/accounts/test_tasks.py

@@ -17,7 +17,11 @@
 
 from warehouse.accounts import tasks
 from warehouse.accounts.models import TermsOfServiceEngagement
-from warehouse.accounts.tasks import compute_user_metrics, notify_users_of_tos_update
+from warehouse.accounts.tasks import (
+    batch_update_email_domain_status,
+    compute_user_metrics,
+    notify_users_of_tos_update,
+)
 
 from ...common.db.accounts import EmailFactory, UserFactory
 from ...common.db.packaging import ProjectFactory, ReleaseFactory
@@ -192,3 +196,55 @@ def test_compute_user_metrics(db_request, metrics):
             ],
         ),
     ]
+
+
+def test_update_email_domain_status(db_request, domain_status_service, mocker):
+    """
+    Test that the batch update performs the correct queries and updates
+    """
+    never_checked = EmailFactory.create(
+        email="[email protected]", domain_last_checked=None
+    )
+    over_threshold = EmailFactory.create(
+        email="[email protected]",
+        domain_last_checked=datetime.now(tz=timezone.utc) - timedelta(days=90),
+    )
+    on_threshold = EmailFactory.create(
+        email="[email protected]",
+        domain_last_checked=datetime.now(tz=timezone.utc) - timedelta(days=30),
+    )
+    under_threshold = EmailFactory.create(
+        email="[email protected]",
+        domain_last_checked=datetime.now(tz=timezone.utc) - timedelta(days=1),
+    )
+
+    batch_update_email_domain_status(db_request)
+
+    assert domain_status_service.get_domain_status.call_count == 3
+    domain_status_service.get_domain_status.assert_has_calls(
+        [
+            mocker.call(never_checked.domain),
+            mocker.call(over_threshold.domain),
+            mocker.call(on_threshold.domain),
+        ]
+    )
+
+    assert never_checked.domain_last_status == ["active"]
+    assert over_threshold.domain_last_status == ["active"]
+    assert on_threshold.domain_last_status == ["active"]
+    assert under_threshold.domain_last_status is None  # no default, not updated
+
+
+def test_update_email_domain_status_does_not_update_if_not_needed(
+    db_request, domain_status_service, mocker
+):
+    mocker.patch.object(domain_status_service, "get_domain_status", return_value=None)
+
+    fail_check = EmailFactory.create()
+
+    batch_update_email_domain_status(db_request)
+
+    domain_status_service.get_domain_status.assert_called_once_with(fail_check.domain)
+
+    assert fail_check.domain_last_checked is None
+    assert fail_check.domain_last_status is None

warehouse/accounts/__init__.py

@@ -32,7 +32,11 @@
     TokenServiceFactory,
     database_login_factory,
 )
-from warehouse.accounts.tasks import compute_user_metrics, notify_users_of_tos_update
+from warehouse.accounts.tasks import (
+    batch_update_email_domain_status,
+    compute_user_metrics,
+    notify_users_of_tos_update,
+)
 from warehouse.accounts.utils import UserContext
 from warehouse.admin.flags import AdminFlagValue
 from warehouse.macaroons.security_policy import MacaroonSecurityPolicy
@@ -215,3 +219,5 @@ def includeme(config):
     # Add a periodic task to generate Account metrics
     config.add_periodic_task(crontab(minute="*/20"), compute_user_metrics)
     config.add_periodic_task(crontab(minute="*"), notify_users_of_tos_update)
+    # TODO: After initial backfill, this can be done less frequently
+    config.add_periodic_task(crontab(minute="*/5"), batch_update_email_domain_status)

warehouse/accounts/interfaces.py

@@ -301,7 +301,7 @@ def get_email_breach_count(email: str) -> int | None:
 
 
 class IDomainStatusService(Interface):
-    def get_domain_status(domain: str) -> list[str]:
+    def get_domain_status(domain: str) -> list[str] | None:
         """
         Returns a list of status strings for the given domain.
         """

warehouse/accounts/models.py

@@ -427,6 +427,7 @@ class Email(db.ModelBase):
     # Domain validation information
     domain_last_checked: Mapped[datetime.datetime | None] = mapped_column(
         comment="Last time domain was checked with the domain validation service.",
+        index=True,
     )
     domain_last_status: Mapped[list[str] | None] = mapped_column(
         ARRAY(String),

warehouse/accounts/services.py

@@ -992,7 +992,7 @@ def create_service(cls, _context, request: Request) -> DomainrDomainStatusServic
         domainr_client_id = request.registry.settings.get("domain_status.client_id")
         return cls(session=request.http, client_id=domainr_client_id)
 
-    def get_domain_status(self, domain: str) -> list[str]:
+    def get_domain_status(self, domain: str) -> list[str] | None:
         """
         Check if a domain is available or not.
         See https://domainr.com/docs/api/v2/status
@@ -1006,6 +1006,6 @@ def get_domain_status(self, domain: str) -> list[str]:
             resp.raise_for_status()
         except requests.RequestException as exc:
             logger.warning("Error contacting Domainr: %r", exc)
-            return []
+            return None
 
         return resp.json()["status"][0]["status"].split()

warehouse/accounts/tasks.py

@@ -10,9 +10,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from datetime import datetime, timedelta, timezone
+from __future__ import annotations
 
-from sqlalchemy import func
+import typing
+
+from datetime import UTC, datetime, timedelta, timezone
+
+from sqlalchemy import func, nullsfirst, or_, select
 
 from warehouse import tasks
 from warehouse.accounts.models import (
@@ -22,10 +26,14 @@
     UserTermsOfServiceEngagement,
 )
 from warehouse.accounts.services import IUserService
+from warehouse.accounts.utils import update_email_domain_status
 from warehouse.email import send_user_terms_of_service_updated
 from warehouse.metrics import IMetricsService
 from warehouse.packaging.models import Release
 
+if typing.TYPE_CHECKING:
+    from pyramid.request import Request
+
 
 @tasks.task(ignore_result=True, acks_late=True)
 def notify_users_of_tos_update(request):
@@ -136,3 +144,32 @@ def compute_user_metrics(request):
             "primary:true",
         ],
     )
+
+
+@tasks.task(ignore_result=True, acks_late=True)
+def batch_update_email_domain_status(request: Request) -> None:
+    """
+    Update the email domain status for any domain last checked over 30 days ago.
+
+    30 days is roughly the time between a domain's expiration
+    and when it enters a renewal grace period.
+    Each TLD may express their own grace period, 30 days is an estimate
+    of time before the registrar is likely to sell it.
+    """
+    stmt = (
+        select(Email)
+        .where(
+            # TODO: After completely backfilled, remove the `or_` for None
+            or_(
+                Email.domain_last_checked.is_(None),
+                Email.domain_last_checked < datetime.now(tz=UTC) - timedelta(days=30),
+            )
+        )
+        .order_by(nullsfirst(Email.domain_last_checked.asc()))
+        .limit(10_000)
+    )
+    # Run in batches to avoid too much memory usage, API rate limits
+    stmt = stmt.execution_options(yield_per=1_000)
+
+    for email in request.db.scalars(stmt):
+        update_email_domain_status(email, request)