Source code for operators.postgres.utils.mangle_users

"""
Mangle Users Operator

This module provides an operator for mangling user emails in a Postgres database.
Mangling prefixes emails and nulls passwords, effectively disabling login access
while preserving the user records for referential integrity.

Usage:
    from operators.postgres.utils.mangle_users import MangleUsersOperator

    mangle_prod_users = MangleUsersOperator(
        task_id="mangle_prod_users",
        postgres_conn_id="boss_grapi_stage",
        tenant="arkansas",
        exclude_emails=["admin@ready.net", "test@ready.net"],
        prefix="synced",
    )
"""

from __future__ import annotations

from typing import Sequence

from airflow.models import BaseOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.utils.context import Context


################################################################################
## Constants
################################################################################

DEFAULT_PREFIX = "synced"
DEFAULT_DOMAIN = "ready.net"
STAGE_CONN_PATTERN = "stage"

SQL_MANGLE_USERS = """
UPDATE {tenant}."User"
SET
    email = CONCAT('{prefix}+', REGEXP_REPLACE(email, '@|\\+', '_', 'g'), '@{domain}'),
    password = NULL
WHERE email NOT LIKE '{prefix}+%%'
AND email NOT LIKE '%%@{domain}'
AND email != ALL(%(exclude_emails)s);
"""


################################################################################
## Operator
################################################################################




class MangleUsersOperator(BaseOperator):
    """
    Mangle user emails by prefixing and nulling passwords.

    This disables login access while preserving user records for FK integrity.
    Emails are transformed: user@example.com -> synced+user_example.com@ready.net

    :param postgres_conn_id: Airflow connection ID for the Postgres database.
    :param tenant: The tenant schema name.
    :param exclude_emails: List of email addresses to exclude from mangling.
    :param prefix: Prefix to prepend to mangled emails (default: "synced").
    :param domain: Domain to append to mangled emails (default: "ready.net").
    """

    ui_color = "#DEFACE"
    template_fields: Sequence[str] = ("postgres_conn_id", "tenant", "exclude_emails")

    def __init__(
        self,
        *,
        postgres_conn_id: str,
        tenant: str,
        exclude_emails: list[str] | None = None,
        prefix: str = DEFAULT_PREFIX,
        domain: str = DEFAULT_DOMAIN,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.postgres_conn_id = postgres_conn_id
        self.tenant = tenant
        self.exclude_emails = exclude_emails or []
        self.prefix = prefix
        self.domain = domain

    ############################################################################
    ## Execute
    ############################################################################



    def execute(self, context: Context) -> int:
        if STAGE_CONN_PATTERN not in self.postgres_conn_id.lower():
            raise ValueError(
                f"Safety check failed: '{self.postgres_conn_id}' does not appear to be a stage connection. "
                "This operator should only modify stage databases."
            )

        hook = PostgresHook(postgres_conn_id=self.postgres_conn_id)

        sql = SQL_MANGLE_USERS.format(
            tenant=self.tenant,
            prefix=self.prefix,
            domain=self.domain,
        )

        exclude_list = self.exclude_emails if self.exclude_emails else []

        self.log.info("Mangling users in %s.User (excluding %d emails)", self.tenant, len(exclude_list))
        self.log.debug("SQL: %s", sql)

        hook.run(sql, parameters={"exclude_emails": exclude_list})

        self.log.info("Mangled users in %s.User", self.tenant)
        return 0