From e96de79665e2709978abc2aa0e867bad5ed619f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adri=C3=A0=20Casaj=C3=BAs?= Date: Wed, 13 Mar 2024 14:30:17 +0100 Subject: [PATCH] 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 | 1 + app/models.py | 13 ++++- cron.py | 14 +++--- .../versions/2024_031212_52510a633d6f_.py | 48 +++++++++++++++++++ oneshot/mark_aliases_as_partner_created.py | 29 +++++++++++ 5 files changed, 98 insertions(+), 7 deletions(-) create mode 100644 migrations/versions/2024_031212_52510a633d6f_.py create mode 100644 oneshot/mark_aliases_as_partner_created.py diff --git a/app/config.py b/app/config.py index 5f33d0c5..c55c555f 100644 --- a/app/config.py +++ b/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") diff --git a/app/models.py b/app/models.py index 06f585df..472b98a8 100644 --- a/app/models.py +++ b/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"" @@ -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)) diff --git a/cron.py b/cron.py index 503046e8..b3446961 100644 --- a/cron.py +++ b/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) ): diff --git a/migrations/versions/2024_031212_52510a633d6f_.py b/migrations/versions/2024_031212_52510a633d6f_.py new file mode 100644 index 00000000..3ebbebaf --- /dev/null +++ b/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 ### diff --git a/oneshot/mark_aliases_as_partner_created.py b/oneshot/mark_aliases_as_partner_created.py new file mode 100644 index 00000000..cf99db91 --- /dev/null +++ b/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()