"""Referral-code generation, lookup, and signup-time linkage. D.1 lays down the bookkeeping only — actual credit application happens in D.3 when the Paddle webhook fires. The flow: 1. /login renders an "invited" banner when the URL carries `?ref=`. 2. The code travels through the email-OTP flow inside the pending cookie so it survives the GET /login → POST /login → /verify hops. 3. When the new user's row is first created (POST /login on an unknown email), `referred_by_user_id` is set and a `Referral` row is written. 4. On the new user's first paid subscription (D.3), we read the `Referral` row to apply discounts to both parties. The code itself is 8 characters from an unambiguous alphabet so users can read it off a phone screen or dictate it over the phone. """ from __future__ import annotations import secrets from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.db import utcnow from app.logging import get_logger from app.models import Referral, User log = get_logger("referral") # Unambiguous alphabet — no 0/O, no 1/I/L. 32 chars → 8 positions ≈ 1e12 # combinations, plenty for our scale, and a unique-constraint catches # collisions if we ever generate the same one twice. _ALPHABET = "ABCDEFGHJKMNPQRSTUVWXYZ23456789" _CODE_LEN = 8 def generate_code() -> str: """Cryptographically random 8-char code from the unambiguous alphabet.""" return "".join(secrets.choice(_ALPHABET) for _ in range(_CODE_LEN)) def normalise_code(raw: str | None) -> str | None: """Trim, uppercase, strip non-alphabet characters. Used on inbound `?ref=` params so users can paste with spaces / lowercase. Returns None if the result isn't a plausible code.""" if not raw: return None cleaned = "".join(c for c in raw.upper() if c in _ALPHABET) if len(cleaned) != _CODE_LEN: return None return cleaned async def assign_code_if_missing(session: AsyncSession, user: User) -> User: """Generate + persist a referral code on `user` if they don't have one yet. Retries on the (very rare) collision. The `user` argument is the User attached to the auth-dependency session, which has since been closed — so it is detached from our `session`. We re-fetch it here before mutating so SQLAlchemy doesn't refuse with 'not persistent within this Session'. """ if user.referral_code: return user db_user = await session.get(User, user.id) if db_user is None: raise RuntimeError(f"referral_service: user {user.id} vanished mid-request") if db_user.referral_code: # Raced with another request — accept their code. return db_user for _ in range(8): code = generate_code() existing = (await session.execute( select(User.id).where(User.referral_code == code) )).scalar_one_or_none() if existing is None: db_user.referral_code = code await session.commit() log.info("referral.code_assigned", user_id=db_user.id, code=code) return db_user # 8 collisions in a row would be a statistical event we'd want to # know about. raise RuntimeError("referral_service: exhausted code-collision retries") async def lookup_referrer(session: AsyncSession, code: str | None) -> User | None: """Return the User whose `referral_code` matches, or None. Normalises the input via `normalise_code` so URL-paste variations all resolve.""" code = normalise_code(code) if not code: return None return (await session.execute( select(User).where(User.referral_code == code) )).scalar_one_or_none() async def link_new_user( session: AsyncSession, new_user: User, referrer: User | None, ) -> Referral | None: """Record a referral if the supplied referrer is valid. Idempotent (safe to call multiple times for the same new user — the unique constraint on `referred_user_id` makes duplicate inserts a no-op). Self-referral is silently rejected. """ if referrer is None or new_user.id is None or referrer.id == new_user.id: return None if new_user.referred_by_user_id is not None: # Already linked; this user can't be referred twice. return None new_user.referred_by_user_id = referrer.id ref = Referral( referrer_user_id=referrer.id, referred_user_id=new_user.id, created_at=utcnow(), ) session.add(ref) await session.commit() await session.refresh(new_user) await session.refresh(ref) log.info( "referral.linked", referrer_id=referrer.id, referred_id=new_user.id, ) return ref