Use random string in ref link instead of tg ID

- Enable old links with tg id by default in .env with LEGACY_REFS=true

db/dal/user_dal.py

@@ -1,4 +1,6 @@
 import logging
+import secrets
+import string
 from typing import Optional, List, Dict, Any, Tuple
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.future import select
@@ -18,6 +20,53 @@ from ..models import (
     AdAttribution,
 )
 
+REFERRAL_CODE_ALPHABET = string.ascii_uppercase + string.digits
+REFERRAL_CODE_LENGTH = 9
+MAX_REFERRAL_CODE_ATTEMPTS = 25
+
+
+def _generate_referral_code_candidate() -> str:
+    return "".join(
+        secrets.choice(REFERRAL_CODE_ALPHABET) for _ in range(REFERRAL_CODE_LENGTH)
+    )
+
+
+async def _referral_code_exists(session: AsyncSession, code: str) -> bool:
+    stmt = select(User.user_id).where(User.referral_code == code)
+    result = await session.execute(stmt)
+    return result.scalar_one_or_none() is not None
+
+
+async def generate_unique_referral_code(session: AsyncSession) -> str:
+    """
+    Generate a unique referral code consisting of uppercase alphanumeric characters.
+    Retries until a free code is found or raises RuntimeError after exceeding attempts.
+    """
+    for _ in range(MAX_REFERRAL_CODE_ATTEMPTS):
+        candidate = _generate_referral_code_candidate()
+        if not await _referral_code_exists(session, candidate):
+            return candidate
+    raise RuntimeError("Failed to generate a unique referral code after several attempts.")
+
+
+async def ensure_referral_code(session: AsyncSession, user: User) -> str:
+    """
+    Ensure the provided user has a referral code, generating and persisting it if missing.
+    Returns the existing or newly generated code.
+    """
+    if user.referral_code:
+        normalized = user.referral_code.strip().upper()
+        if normalized != user.referral_code:
+            user.referral_code = normalized
+            await session.flush()
+            await session.refresh(user)
+        return user.referral_code
+
+    user.referral_code = await generate_unique_referral_code(session)
+    await session.flush()
+    await session.refresh(user)
+    return user.referral_code
+
 
 async def get_user_by_id(session: AsyncSession, user_id: int) -> Optional[User]:
     stmt = select(User).where(User.user_id == user_id)
@@ -52,6 +101,11 @@ async def create_user(session: AsyncSession, user_data: Dict[str, Any]) -> Tuple
     if "registration_date" not in user_data:
         user_data["registration_date"] = datetime.now(timezone.utc)
 
+    if not user_data.get("referral_code"):
+        user_data["referral_code"] = await generate_unique_referral_code(session)
+    else:
+        user_data["referral_code"] = user_data["referral_code"].strip().upper()
+
     # Use PostgreSQL upsert to avoid IntegrityError on concurrent inserts
     stmt = (
         pg_insert(User)
@@ -80,6 +134,15 @@ async def create_user(session: AsyncSession, user_data: Dict[str, Any]) -> Tuple
     return user, created
 
 
+async def get_user_by_referral_code(session: AsyncSession, referral_code: str) -> Optional[User]:
+    normalized = referral_code.strip().upper()
+    if not normalized:
+        return None
+    stmt = select(User).where(User.referral_code == normalized)
+    result = await session.execute(stmt)
+    return result.scalar_one_or_none()
+
+
 async def update_user(
     session: AsyncSession, user_id: int, update_data: Dict[str, Any]
 ) -> Optional[User]:

db/migrator.py

@@ -48,12 +48,86 @@ def _migration_0001_add_channel_subscription_fields(connection: Connection) -> N
         connection.execute(text(stmt))
 
 
+def _migration_0002_add_referral_code(connection: Connection) -> None:
+    inspector = inspect(connection)
+    columns: Set[str] = {col["name"] for col in inspector.get_columns("users")}
+
+    if "referral_code" not in columns:
+        connection.execute(
+            text("ALTER TABLE users ADD COLUMN referral_code VARCHAR(16)")
+        )
+
+    connection.execute(
+        text(
+            """
+            WITH generated_codes AS (
+                SELECT
+                    user_id,
+                    UPPER(
+                        SUBSTRING(
+                            md5(
+                                user_id::text
+                                || clock_timestamp()::text
+                                || random()::text
+                            )
+                            FROM 1 FOR 9
+                        )
+                    ) AS referral_code
+                FROM users
+                WHERE referral_code IS NULL OR referral_code = ''
+            )
+            UPDATE users AS u
+            SET referral_code = g.referral_code
+            FROM generated_codes AS g
+            WHERE u.user_id = g.user_id
+            """
+        )
+    )
+
+    connection.execute(
+        text(
+            """
+            CREATE UNIQUE INDEX IF NOT EXISTS uq_users_referral_code
+            ON users (referral_code)
+            WHERE referral_code IS NOT NULL
+            """
+        )
+    )
+
+
+def _migration_0003_normalize_referral_codes(connection: Connection) -> None:
+    inspector = inspect(connection)
+    columns: Set[str] = {col["name"] for col in inspector.get_columns("users")}
+    if "referral_code" not in columns:
+        return
+
+    connection.execute(
+        text(
+            """
+            UPDATE users
+            SET referral_code = UPPER(referral_code)
+            WHERE referral_code IS NOT NULL
+              AND referral_code <> UPPER(referral_code)
+            """
+        )
+    )
+
 MIGRATIONS: List[Migration] = [
     Migration(
         id="0001_add_channel_subscription_fields",
         description="Add columns to track required channel subscription verification",
         upgrade=_migration_0001_add_channel_subscription_fields,
     ),
+    Migration(
+        id="0002_add_referral_code",
+        description="Store short referral codes for users and backfill existing rows",
+        upgrade=_migration_0002_add_referral_code,
+    ),
+    Migration(
+        id="0003_normalize_referral_codes",
+        description="Normalize referral codes to uppercase for consistent lookups",
+        upgrade=_migration_0003_normalize_referral_codes,
+    ),
 ]
 
 

db/models.py

@@ -21,6 +21,7 @@ class User(Base):
                                server_default=func.now())
     is_banned = Column(Boolean, default=False)
     panel_user_uuid = Column(String, nullable=True, unique=True, index=True)
+    referral_code = Column(String(16), nullable=True, unique=True, index=True)
     referred_by_id = Column(BigInteger,
                             ForeignKey("users.user_id"),
                             nullable=True)