austnv/remnawave-minishop
1
0
Fork 0
Code Issues 1 Pull Requests Activity

Merge pull request #129 from 3252a8/feature/ref-link

Browse Source 
Ref link improvement
This commit is contained in:
machka pasla
2025-11-13 23:00:43 +03:00
committed by GitHub
parent 643bf5fabb caf5b88eef
commit ead990e61c
9 changed files with 236 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.env.example
+1
View File
@@ -93,6 +93,7 @@ SUBSCRIPTION_NOTIFY_DAYS_BEFORE=3 # Da
REFERRAL_ONE_BONUS_PER_REFEREE=False # Give a bonus only once per referee
LEGACY_REFS=true # Allow ref_<tg_id> links. Leave unset/true unless you want to disable old links
# Referral Bonus Days
# Bonus for the inviting user
REFERRAL_BONUS_DAYS_1_MONTH=3
bot/handlers/inline_mode.py
+23 -5
View File
@@ -38,7 +38,13 @@ async def inline_query_handler(inline_query: InlineQuery,
# For all users: referral functionality
if not query or "реф" in query or "ref" in query or "друг" in query or "friend" in query:
referral_result = await create_referral_result(
inline_query, bot, referral_service, i18n, current_lang, settings
inline_query,
bot,
referral_service,
i18n,
current_lang,
settings,
session,
)
if referral_result:
results.append(referral_result)
@@ -67,9 +73,15 @@ async def inline_query_handler(inline_query: InlineQuery,
await inline_query.answer(results=[], cache_time=10)
async def create_referral_result(inline_query: InlineQuery, bot: Bot,
referral_service: ReferralService,
i18n_instance, lang: str, settings: Settings) -> Optional[InlineQueryResultArticle]:
async def create_referral_result(
inline_query: InlineQuery,
bot: Bot,
referral_service: ReferralService,
i18n_instance,
lang: str,
settings: Settings,
session: AsyncSession,
) -> Optional[InlineQueryResultArticle]:
"""Create referral link result for inline query"""
_ = lambda key, **kwargs: i18n_instance.gettext(lang, key, **kwargs)
@@ -80,7 +92,13 @@ async def create_referral_result(inline_query: InlineQuery, bot: Bot,
return None
user_id = inline_query.from_user.id
referral_link = referral_service.generate_referral_link(bot_username, user_id)
referral_link = await referral_service.generate_referral_link(
session, bot_username, user_id
)
if not referral_link:
logging.warning("Could not produce referral link for inline user %s", user_id)
return None
# Create message content (use same text as friend message)
message_text = _(
bot/handlers/user/referral.py
+22 -3
View File
@@ -60,8 +60,18 @@ async def referral_command_handler(event: Union[types.Message,
return
inviter_user_id = event.from_user.id
referral_link = referral_service.generate_referral_link(
bot_username, inviter_user_id)
referral_link = await referral_service.generate_referral_link(
session, bot_username, inviter_user_id)
if not referral_link:
logging.error(
"Failed to generate referral link for user %s (probably missing DB record).",
inviter_user_id,
)
await target_message_obj.answer(_("error_generating_referral_link"))
if isinstance(event, types.CallbackQuery):
await event.answer()
return
bonus_info_parts = []
if settings.subscription_options:
@@ -132,7 +142,16 @@ async def referral_action_handler(callback: types.CallbackQuery, settings: Setti
return
inviter_user_id = callback.from_user.id
referral_link = referral_service.generate_referral_link(bot_username, inviter_user_id)
referral_link = await referral_service.generate_referral_link(
session, bot_username, inviter_user_id)
if not referral_link:
logging.error(
"Failed to generate referral link for user %s via inline button.",
inviter_user_id,
)
await callback.answer(_("error_generating_referral_link"), show_alert=True)
return
friend_message = _("referral_friend_message", referral_link=referral_link)
bot/handlers/user/start.py
+18 -4
View File
@@ -302,7 +302,7 @@ async def ensure_required_channel_subscription(
@router.message(CommandStart())
@router.message(CommandStart(magic=F.args.regexp(r"^ref_(\d+)$").as_("ref_match")))
@router.message(CommandStart(magic=F.args.regexp(r"^ref_((?:[uU][A-Za-z0-9]{9})|(?:[A-Za-z0-9]{9})|\d+)$").as_("ref_match")))
@router.message(CommandStart(magic=F.args.regexp(r"^promo_(\w+)$").as_("promo_match")))
@router.message(CommandStart(magic=F.args.regexp(r"^(?!ref_|promo_)([A-Za-z0-9_\-]{2,64})$").as_("ad_param_match")))
async def start_command_handler(message: types.Message,
@@ -328,9 +328,23 @@ async def start_command_handler(message: types.Message,
ad_start_param: Optional[str] = None
if ref_match:
potential_referrer_id = int(ref_match.group(1))
if await user_dal.get_user_by_id(session, potential_referrer_id):
referred_by_user_id = potential_referrer_id
raw_ref_value = ref_match.group(1)
if raw_ref_value.isdigit():
if settings.LEGACY_REFS:
potential_referrer_id = int(raw_ref_value)
if potential_referrer_id != user_id and await user_dal.get_user_by_id(
session, potential_referrer_id):
referred_by_user_id = potential_referrer_id
else:
normalized_code = raw_ref_value.strip()
if normalized_code and normalized_code[0].lower() == "u":
normalized_code = normalized_code[1:]
ref_user = None
if normalized_code:
ref_user = await user_dal.get_user_by_referral_code(
session, normalized_code)
if ref_user and ref_user.user_id != user_id:
referred_by_user_id = ref_user.user_id
elif promo_match:
promo_code_to_apply = promo_match.group(1)
logging.info(f"User {user_id} started with promo code: {promo_code_to_apply}")
bot/services/referral_service.py
+29 -3
View File
@@ -257,9 +257,35 @@ class ReferralService:
raise
def generate_referral_link(self, bot_username: str,
inviter_user_id: int) -> str:
return f"https://t.me/{bot_username}?start=ref_{inviter_user_id}"
async def generate_referral_link(self, session: AsyncSession,
bot_username: str,
inviter_user_id: int) -> Optional[str]:
try:
user = await user_dal.get_user_by_id(session, inviter_user_id)
if not user:
logging.warning(
"Unable to generate referral link: user %s not found.",
inviter_user_id,
)
return None
referral_code = await user_dal.ensure_referral_code(session, user)
if not referral_code:
logging.warning(
"User %s has no referral code even after regeneration attempt.",
inviter_user_id,
)
return None
return f"https://t.me/{bot_username}?start=ref_u{referral_code}"
except Exception as exc:
logging.error(
"Failed to generate referral link for user %s: %s",
inviter_user_id,
exc,
exc_info=True,
)
return None
async def get_referral_stats(self, session: AsyncSession, user_id: int) -> dict:
"""Get referral statistics for a user"""
config/settings.py
+5 -1
View File
@@ -118,7 +118,11 @@ class Settings(BaseSettings):
# Referral program configuration
REFERRAL_ONE_BONUS_PER_REFEREE: bool = Field(
default=True,
description="When true, referral bonuses (for inviter and referee) are applied only once per invited user – on their first successful payment."
description="When true, referral bonuses (for inviter and referee) are applied only once per invited user - on their first successful payment."
)
LEGACY_REFS: bool = Field(
default=True,
description="Allow legacy referral links like ref_<telegram_id> to continue working. Defaults to True when unset."
)
PANEL_API_URL: Optional[str] = None
db/dal/user_dal.py
+63
View File
@@ -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
+74
View File
@@ -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
+1
View File
@@ -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)