Add missing indexes and mark aliases as created by partner (#2058)

* Add missing indexes and mark aliases as created by partner

* Configure if we should skip the partner aliases or not

app/config.py

@@ -432,6 +432,7 @@ except Exception:
 HIBP_API_KEYS = sl_getenv("HIBP_API_KEYS", list) or []
 HIBP_MAX_ALIAS_CHECK = 10_000
 HIBP_RPM = 100
+HIBP_SKIP_PARTNER_ALIAS = os.environ.get("HIBP_SKIP_PARTNER_ALIAS")
 
 POSTMASTER = os.environ.get("POSTMASTER")
 

app/models.py

@@ -1422,6 +1422,9 @@ def generate_random_alias_email(
 
 class Alias(Base, ModelMixin):
     __tablename__ = "alias"
+
+    FLAG_PARTNER_CREATED = 1 << 0
+
     user_id = sa.Column(
         sa.ForeignKey(User.id, ondelete="cascade"), nullable=False, index=True
     )
@@ -1431,6 +1434,9 @@ class Alias(Base, ModelMixin):
     name = sa.Column(sa.String(128), nullable=True, default=None)
 
     enabled = sa.Column(sa.Boolean(), default=True, nullable=False)
+    flags = sa.Column(
+        sa.BigInteger(), default=0, server_default="0", nullable=False, index=True
+    )
 
     custom_domain_id = sa.Column(
         sa.ForeignKey("custom_domain.id", ondelete="cascade"), nullable=True, index=True
@@ -2586,10 +2592,13 @@ class Job(Base, ModelMixin):
         nullable=False,
         server_default=str(JobState.ready.value),
         default=JobState.ready.value,
+        index=True,
     )
     attempts = sa.Column(sa.Integer, nullable=False, server_default="0", default=0)
     taken_at = sa.Column(ArrowType, nullable=True)
 
+    __table_args__ = (Index("ix_state_run_at_taken_at", state, run_at, taken_at),)
+
     def __repr__(self):
         return f"<Job {self.id} {self.name} {self.payload}>"
 
@@ -2937,7 +2946,9 @@ class RecoveryCode(Base, ModelMixin):
 
 class Notification(Base, ModelMixin):
     __tablename__ = "notification"
-    user_id = sa.Column(sa.ForeignKey(User.id, ondelete="cascade"), nullable=False)
+    user_id = sa.Column(
+        sa.ForeignKey(User.id, ondelete="cascade"), nullable=False, index=True
+    )
     message = sa.Column(sa.Text, nullable=False)
     title = sa.Column(sa.String(512))
 

cron.py

@@ -1059,13 +1059,15 @@ async def check_hibp():
     LOG.d("Preparing list of aliases to check")
     queue = asyncio.Queue()
     max_date = arrow.now().shift(days=-config.HIBP_SCAN_INTERVAL_DAYS)
+    alias_query = Alias.filter(
+        or_(Alias.hibp_last_check.is_(None), Alias.hibp_last_check < max_date),
+        Alias.user_id.notin_(user_ids),
+        Alias.enabled,
+    )
+    if config.HIBP_SKIP_PARTNER_ALIAS:
+        alias_query = alias_query(Alias.flags.op("&")(Alias.FLAG_PARTNER_CREATED) == 0)
     for alias in (
-        Alias.filter(
-            or_(Alias.hibp_last_check.is_(None), Alias.hibp_last_check < max_date),
-            Alias.user_id.notin_(user_ids),
-        )
-        .filter(Alias.enabled)
-        .order_by(nullsfirst(Alias.hibp_last_check.asc()), Alias.id.asc())
+        alias_query.order_by(nullsfirst(Alias.hibp_last_check.asc()), Alias.id.asc())
         .yield_per(500)
         .enable_eagerloads(False)
     ):

migrations/versions/2024_031212_52510a633d6f_.py

@@ -0,0 +1,48 @@
+"""empty message
+
+Revision ID: 52510a633d6f
+Revises: 818b0a956205
+Create Date: 2024-03-12 12:46:24.161644
+
+"""
+import sqlalchemy_utils
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "52510a633d6f"
+down_revision = "818b0a956205"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "alias", sa.Column("flags", sa.BigInteger(), server_default="0", nullable=False)
+    )
+    with op.get_context().autocommit_block():
+        op.create_index(op.f("ix_alias_flags"), "alias", ["flags"], unique=False)
+        op.create_index(op.f("ix_job_state"), "job", ["state"], unique=False)
+        op.create_index(
+            "ix_state_run_at_taken_at",
+            "job",
+            ["state", "run_at", "taken_at"],
+            unique=False,
+        )
+        op.create_index(
+            op.f("ix_notification_user_id"), "notification", ["user_id"], unique=False
+        )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.get_context().autocommit_block():
+        op.drop_index(op.f("ix_notification_user_id"), table_name="notification")
+        op.drop_index("ix_state_run_at_taken_at", table_name="job")
+        op.drop_index(op.f("ix_job_state"), table_name="job")
+        op.drop_index(op.f("ix_alias_flags"), table_name="alias")
+    op.drop_column("alias", "flags")
+    # ### end Alembic commands ###

oneshot/mark_aliases_as_partner_created.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+import argparse
+
+
+from app.log import LOG
+from app.models import Alias, SLDomain
+from app.db import Session
+
+parser = argparse.ArgumentParser(
+    prog="Mark partner created aliases with the PARTNER_CREATED flag",
+)
+args = parser.parse_args()
+
+domains = SLDomain.filter(SLDomain.partner_id.isnot(None)).all()
+
+for domain in domains:
+    LOG.i(f"Checking aliases for domain {domain.domain}")
+    for alias in (
+        Alias.filter(
+            Alias.email.like(f"%{domain.domain}"),
+            Alias.flags.op("&")(Alias.FLAG_PARTNER_CREATED) == 0,
+        )
+        .enable_eagerloads(False)
+        .yield_per(100)
+        .all()
+    ):
+        alias.flags = alias.flags | Alias.FLAG_PARTNER_CREATED
+        LOG.i(f" * Updating {alias.email} to {alias.flags}")
+        Session.commit()