Private
Public Access
forked from FOSS/AutoShop-Djimbo-Simple
feat(referral): add referral system with withdrawals
This commit is contained in:
@@ -6,7 +6,29 @@ from .db_purchases import PurchaseModel, Purchasesx
|
||||
from .db_refill import RefillModel, Refillx
|
||||
from .db_settings import SettingsModel, Settingsx
|
||||
from .db_users import UserModel, UsersRepository, Userx
|
||||
from .entities import Category, Item, Payments, Position, Purchase, Refill, Settings, User
|
||||
|
||||
from .db_referral import (
|
||||
ReferralModel,
|
||||
ReferralTransactionModel,
|
||||
ReferralWithdrawalModel,
|
||||
Referralx,
|
||||
ReferralTransactionx,
|
||||
ReferralWithdrawalx,
|
||||
)
|
||||
|
||||
from .entities import (
|
||||
Category,
|
||||
Item,
|
||||
Payments,
|
||||
Position,
|
||||
Purchase,
|
||||
Refill,
|
||||
Settings,
|
||||
User,
|
||||
Referral,
|
||||
ReferralTransaction,
|
||||
ReferralWithdrawal,
|
||||
)
|
||||
|
||||
ModelCategory = Category
|
||||
ModelItem = Item
|
||||
@@ -18,3 +40,7 @@ ModelSettings = Settings
|
||||
ModelUser = User
|
||||
ModelUsers = User
|
||||
SettingsRepository = Settingsx
|
||||
|
||||
ModelReferral = Referral
|
||||
ModelReferralTransaction = ReferralTransaction
|
||||
ModelReferralWithdrawal = ReferralWithdrawal
|
||||
|
||||
@@ -0,0 +1,728 @@
|
||||
# - *- coding: utf- 8 - *-
|
||||
from math import isfinite
|
||||
import re
|
||||
from urllib.parse import urlparse
|
||||
from uuid import uuid4
|
||||
|
||||
from sqlalchemy import BigInteger, Float, Integer, String, UniqueConstraint, text
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
from sqlalchemy.ext.asyncio import AsyncConnection
|
||||
|
||||
from typing import Any, Dict, Optional, Tuple, Union
|
||||
|
||||
from tgbot.database.core import Base, session_factory
|
||||
from tgbot.database.entities import Referral, ReferralTransaction, ReferralWithdrawal
|
||||
from tgbot.database.repository import BaseRepository
|
||||
from tgbot.utils.const_functions import get_unix
|
||||
|
||||
|
||||
def is_valid_referral_withdrawal_recipient(
|
||||
withdrawal_method: str, recipient: str
|
||||
) -> bool:
|
||||
if withdrawal_method == "Cryptobot":
|
||||
return bool(re.fullmatch(r"@[A-Za-z][A-Za-z0-9_]{4,31}", recipient))
|
||||
|
||||
if withdrawal_method == "Lolzteam":
|
||||
parsed_url = urlparse(recipient)
|
||||
return parsed_url.scheme in ("http", "https") and parsed_url.hostname in (
|
||||
"zelenka.guru",
|
||||
"lolz.team",
|
||||
"lolz.live",
|
||||
"lolz.guru",
|
||||
)
|
||||
|
||||
return False
|
||||
|
||||
|
||||
async def ensure_referrals_schema(conn: AsyncConnection) -> None:
|
||||
user_columns_result = await conn.exec_driver_sql("PRAGMA table_info(storage_users)")
|
||||
user_columns = {column[1] for column in user_columns_result.fetchall()}
|
||||
|
||||
# Обновление полей пользователя
|
||||
if user_columns:
|
||||
if "user_referrer_id" not in user_columns:
|
||||
await conn.exec_driver_sql(
|
||||
"ALTER TABLE storage_users ADD COLUMN user_referrer_id BIGINT"
|
||||
)
|
||||
if "user_referral_balance" not in user_columns:
|
||||
await conn.exec_driver_sql(
|
||||
"ALTER TABLE storage_users "
|
||||
"ADD COLUMN user_referral_balance FLOAT NOT NULL DEFAULT 0"
|
||||
)
|
||||
if "user_referral_hold" not in user_columns:
|
||||
await conn.exec_driver_sql(
|
||||
"ALTER TABLE storage_users "
|
||||
"ADD COLUMN user_referral_hold FLOAT NOT NULL DEFAULT 0"
|
||||
)
|
||||
|
||||
await conn.exec_driver_sql(
|
||||
"CREATE INDEX IF NOT EXISTS ix_storage_users_user_referrer_id "
|
||||
"ON storage_users (user_referrer_id)"
|
||||
)
|
||||
|
||||
settings_columns_result = await conn.exec_driver_sql(
|
||||
"PRAGMA table_info(storage_settings)"
|
||||
)
|
||||
settings_columns = {column[1] for column in settings_columns_result.fetchall()}
|
||||
|
||||
# Обновление полей настроек
|
||||
if settings_columns:
|
||||
if "status_referral" not in settings_columns:
|
||||
await conn.exec_driver_sql(
|
||||
"ALTER TABLE storage_settings "
|
||||
"ADD COLUMN status_referral VARCHAR(16) NOT NULL DEFAULT 'False'"
|
||||
)
|
||||
if "referral_bonus_rub" not in settings_columns:
|
||||
await conn.exec_driver_sql(
|
||||
"ALTER TABLE storage_settings "
|
||||
"ADD COLUMN referral_bonus_rub FLOAT NOT NULL DEFAULT 0"
|
||||
)
|
||||
if "referral_refill_percent" not in settings_columns:
|
||||
await conn.exec_driver_sql(
|
||||
"ALTER TABLE storage_settings "
|
||||
"ADD COLUMN referral_refill_percent FLOAT NOT NULL DEFAULT 0"
|
||||
)
|
||||
|
||||
|
||||
class ReferralModel(Base):
|
||||
__tablename__ = "storage_referrals"
|
||||
|
||||
increment: Mapped[int] = mapped_column(
|
||||
Integer, primary_key=True, autoincrement=True
|
||||
)
|
||||
referrer_id: Mapped[int] = mapped_column(BigInteger, nullable=False, index=True)
|
||||
referral_id: Mapped[int] = mapped_column(
|
||||
BigInteger, nullable=False, unique=True, index=True
|
||||
)
|
||||
referral_unix: Mapped[int] = mapped_column(
|
||||
Integer, nullable=False, default=get_unix
|
||||
)
|
||||
|
||||
|
||||
class ReferralTransactionModel(Base):
|
||||
__tablename__ = "storage_referral_transactions"
|
||||
__table_args__ = (
|
||||
UniqueConstraint(
|
||||
"user_id",
|
||||
"transaction_type",
|
||||
"source_id",
|
||||
name="uq_referral_transaction_source",
|
||||
),
|
||||
)
|
||||
|
||||
increment: Mapped[int] = mapped_column(
|
||||
Integer, primary_key=True, autoincrement=True
|
||||
)
|
||||
user_id: Mapped[int] = mapped_column(BigInteger, nullable=False, index=True)
|
||||
related_user_id: Mapped[Optional[int]] = mapped_column(BigInteger, nullable=True)
|
||||
|
||||
transaction_type: Mapped[str] = mapped_column(String(64), nullable=False)
|
||||
source_id: Mapped[str] = mapped_column(String(128), nullable=False)
|
||||
amount: Mapped[float] = mapped_column(Float, nullable=False, default=0)
|
||||
transaction_unix: Mapped[int] = mapped_column(
|
||||
Integer, nullable=False, default=get_unix
|
||||
)
|
||||
|
||||
|
||||
class ReferralWithdrawalModel(Base):
|
||||
__tablename__ = "storage_referral_withdrawals"
|
||||
|
||||
increment: Mapped[int] = mapped_column(
|
||||
Integer, primary_key=True, autoincrement=True
|
||||
)
|
||||
user_id: Mapped[int] = mapped_column(BigInteger, nullable=False, index=True)
|
||||
withdrawal_amount: Mapped[float] = mapped_column(Float, nullable=False, default=0)
|
||||
withdrawal_method: Mapped[str] = mapped_column(String(32), nullable=False)
|
||||
withdrawal_recipient: Mapped[str] = mapped_column(String(255), nullable=False)
|
||||
withdrawal_status: Mapped[str] = mapped_column(
|
||||
String(32), nullable=False, default="pending"
|
||||
)
|
||||
withdrawal_unix: Mapped[int] = mapped_column(
|
||||
Integer, nullable=False, default=get_unix
|
||||
)
|
||||
|
||||
processed_unix: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
|
||||
processed_admin_id: Mapped[Optional[int]] = mapped_column(BigInteger, nullable=True)
|
||||
|
||||
|
||||
class Referralx(BaseRepository[ReferralModel, Referral]):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.table_model = ReferralModel
|
||||
self.entity_model = Referral
|
||||
self.storage_name = ReferralModel.__tablename__
|
||||
|
||||
async def add(self, referrer_id: int, referral_id: int) -> Referral:
|
||||
return await self._insert(
|
||||
referrer_id=referrer_id,
|
||||
referral_id=referral_id,
|
||||
referral_unix=get_unix(),
|
||||
)
|
||||
|
||||
async def register_referral(self, referral_id: int, referrer_id: int) -> str:
|
||||
if referral_id == referrer_id:
|
||||
return "SELF_REFERRAL"
|
||||
|
||||
async with session_factory() as session:
|
||||
try:
|
||||
await session.execute(text("BEGIN IMMEDIATE"))
|
||||
|
||||
settings_result = await session.execute(
|
||||
text(
|
||||
"SELECT status_referral, referral_bonus_rub "
|
||||
"FROM storage_settings WHERE id = 1"
|
||||
)
|
||||
)
|
||||
settings = settings_result.mappings().first()
|
||||
|
||||
if settings is None or settings["status_referral"] != "True":
|
||||
await session.rollback()
|
||||
return "DISABLED"
|
||||
|
||||
referrer_result = await session.execute(
|
||||
text(
|
||||
"SELECT user_referral_balance FROM storage_users "
|
||||
"WHERE user_id = :user_id"
|
||||
),
|
||||
{"user_id": referrer_id},
|
||||
)
|
||||
referrer = referrer_result.mappings().first()
|
||||
|
||||
if referrer is None:
|
||||
await session.rollback()
|
||||
return "REFERRER_NOT_FOUND"
|
||||
|
||||
referral_result = await session.execute(
|
||||
text(
|
||||
"SELECT user_balance, user_referrer_id "
|
||||
"FROM storage_users WHERE user_id = :user_id"
|
||||
),
|
||||
{"user_id": referral_id},
|
||||
)
|
||||
referral = referral_result.mappings().first()
|
||||
|
||||
if referral is None:
|
||||
await session.rollback()
|
||||
return "REFERRAL_NOT_FOUND"
|
||||
|
||||
existing_referral_result = await session.execute(
|
||||
text(
|
||||
"SELECT increment FROM storage_referrals "
|
||||
"WHERE referral_id = :user_id LIMIT 1"
|
||||
),
|
||||
{"user_id": referral_id},
|
||||
)
|
||||
|
||||
if (
|
||||
referral["user_referrer_id"] is not None
|
||||
or existing_referral_result.mappings().first() is not None
|
||||
):
|
||||
await session.rollback()
|
||||
return "ALREADY"
|
||||
|
||||
bonus = round(float(settings["referral_bonus_rub"]), 2)
|
||||
source_id = f"registration:{referral_id}"
|
||||
now_unix = get_unix()
|
||||
|
||||
# Приглашённому — бонус на основной баланс.
|
||||
await session.execute(
|
||||
text(
|
||||
"INSERT INTO storage_referrals "
|
||||
"(referrer_id, referral_id, referral_unix) "
|
||||
"VALUES (:referrer_id, :referral_id, :referral_unix)"
|
||||
),
|
||||
{
|
||||
"referrer_id": referrer_id,
|
||||
"referral_id": referral_id,
|
||||
"referral_unix": now_unix,
|
||||
},
|
||||
)
|
||||
await session.execute(
|
||||
text(
|
||||
"UPDATE storage_users "
|
||||
"SET user_referrer_id = :referrer_id, "
|
||||
"user_balance = :balance "
|
||||
"WHERE user_id = :user_id"
|
||||
),
|
||||
{
|
||||
"referrer_id": referrer_id,
|
||||
"balance": round(float(referral["user_balance"]) + bonus, 2),
|
||||
"user_id": referral_id,
|
||||
},
|
||||
)
|
||||
|
||||
# Рефереру — бонус на реферальный баланс.
|
||||
await session.execute(
|
||||
text(
|
||||
"UPDATE storage_users "
|
||||
"SET user_referral_balance = :balance "
|
||||
"WHERE user_id = :user_id"
|
||||
),
|
||||
{
|
||||
"balance": round(
|
||||
float(referrer["user_referral_balance"]) + bonus, 2
|
||||
),
|
||||
"user_id": referrer_id,
|
||||
},
|
||||
)
|
||||
await session.execute(
|
||||
text(
|
||||
"INSERT INTO storage_referral_transactions "
|
||||
"(user_id, related_user_id, transaction_type, source_id, amount, transaction_unix) "
|
||||
"VALUES (:user_id, :related_user_id, :transaction_type, :source_id, :amount, :transaction_unix)"
|
||||
),
|
||||
{
|
||||
"user_id": referral_id,
|
||||
"related_user_id": referrer_id,
|
||||
"transaction_type": "registration_bonus_main",
|
||||
"source_id": source_id,
|
||||
"amount": bonus,
|
||||
"transaction_unix": now_unix,
|
||||
},
|
||||
)
|
||||
await session.execute(
|
||||
text(
|
||||
"INSERT INTO storage_referral_transactions "
|
||||
"(user_id, related_user_id, transaction_type, source_id, amount, transaction_unix) "
|
||||
"VALUES (:user_id, :related_user_id, :transaction_type, :source_id, :amount, :transaction_unix)"
|
||||
),
|
||||
{
|
||||
"user_id": referrer_id,
|
||||
"related_user_id": referral_id,
|
||||
"transaction_type": "registration_bonus_referral",
|
||||
"source_id": source_id,
|
||||
"amount": bonus,
|
||||
"transaction_unix": now_unix,
|
||||
},
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
raise
|
||||
|
||||
return "ok"
|
||||
|
||||
async def transfer_to_main_balance(self, user_id: int, amount: float) -> str:
|
||||
try:
|
||||
transfer_amount = round(float(amount), 2)
|
||||
except (TypeError, ValueError):
|
||||
return "INVALID_AMOUNT"
|
||||
|
||||
if not isfinite(transfer_amount) or transfer_amount <= 0:
|
||||
return "INVALID_AMOUNT"
|
||||
|
||||
async with session_factory() as session:
|
||||
try:
|
||||
await session.execute(text("BEGIN IMMEDIATE"))
|
||||
|
||||
user_result = await session.execute(
|
||||
text(
|
||||
"SELECT user_balance, user_referral_balance, user_referral_hold "
|
||||
"FROM storage_users WHERE user_id = :user_id"
|
||||
),
|
||||
{"user_id": user_id},
|
||||
)
|
||||
get_user = user_result.mappings().first()
|
||||
|
||||
if get_user is None:
|
||||
await session.rollback()
|
||||
return "USER_NOT_FOUND"
|
||||
|
||||
available_balance = round(
|
||||
float(get_user["user_referral_balance"])
|
||||
- float(get_user["user_referral_hold"]),
|
||||
2,
|
||||
)
|
||||
|
||||
if transfer_amount > available_balance:
|
||||
await session.rollback()
|
||||
return "INSUFFICIENT_FUNDS"
|
||||
|
||||
await session.execute(
|
||||
text(
|
||||
"UPDATE storage_users "
|
||||
"SET user_balance = :user_balance, "
|
||||
"user_referral_balance = :referral_balance "
|
||||
"WHERE user_id = :user_id"
|
||||
),
|
||||
{
|
||||
"user_balance": round(
|
||||
float(get_user["user_balance"]) + transfer_amount,
|
||||
2,
|
||||
),
|
||||
"referral_balance": round(
|
||||
float(get_user["user_referral_balance"]) - transfer_amount,
|
||||
2,
|
||||
),
|
||||
"user_id": user_id,
|
||||
},
|
||||
)
|
||||
await session.execute(
|
||||
text(
|
||||
"INSERT INTO storage_referral_transactions "
|
||||
"(user_id, related_user_id, transaction_type, source_id, amount, transaction_unix) "
|
||||
"VALUES (:user_id, NULL, :transaction_type, :source_id, :amount, :transaction_unix)"
|
||||
),
|
||||
{
|
||||
"user_id": user_id,
|
||||
"transaction_type": "transfer_to_main",
|
||||
"source_id": f"transfer:{user_id}:{uuid4().hex}",
|
||||
"amount": transfer_amount,
|
||||
"transaction_unix": get_unix(),
|
||||
},
|
||||
)
|
||||
await session.commit()
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
raise
|
||||
|
||||
return "ok"
|
||||
|
||||
|
||||
class ReferralTransactionx(
|
||||
BaseRepository[ReferralTransactionModel, ReferralTransaction]
|
||||
):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.table_model = ReferralTransactionModel
|
||||
self.entity_model = ReferralTransaction
|
||||
self.storage_name = ReferralTransactionModel.__tablename__
|
||||
|
||||
async def get_refill_reward(self, source_id: str) -> Optional[ReferralTransaction]:
|
||||
return await self.get(
|
||||
transaction_type="refill_percent",
|
||||
source_id=source_id,
|
||||
)
|
||||
|
||||
async def add(
|
||||
self,
|
||||
user_id: int,
|
||||
transaction_type: str,
|
||||
source_id: str,
|
||||
amount: float,
|
||||
related_user_id: Optional[int] = None,
|
||||
) -> ReferralTransaction:
|
||||
return await self._insert(
|
||||
user_id=user_id,
|
||||
related_user_id=related_user_id,
|
||||
transaction_type=transaction_type,
|
||||
source_id=source_id,
|
||||
amount=round(amount, 2),
|
||||
transaction_unix=get_unix(),
|
||||
)
|
||||
|
||||
|
||||
class ReferralWithdrawalx(BaseRepository[ReferralWithdrawalModel, ReferralWithdrawal]):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.table_model = ReferralWithdrawalModel
|
||||
self.entity_model = ReferralWithdrawal
|
||||
self.storage_name = ReferralWithdrawalModel.__tablename__
|
||||
|
||||
async def create_pending(
|
||||
self,
|
||||
user_id: int,
|
||||
withdrawal_amount: float,
|
||||
withdrawal_method: str,
|
||||
withdrawal_recipient: str,
|
||||
) -> Tuple[str, Optional[ReferralWithdrawal]]:
|
||||
try:
|
||||
amount = round(float(withdrawal_amount), 2)
|
||||
except (TypeError, ValueError):
|
||||
return "INVALID_AMOUNT", None
|
||||
|
||||
recipient = str(withdrawal_recipient).strip()
|
||||
|
||||
if not isfinite(amount) or amount <= 0:
|
||||
return "INVALID_AMOUNT", None
|
||||
|
||||
if withdrawal_method not in ("Cryptobot", "Lolzteam"):
|
||||
return "INVALID_METHOD", None
|
||||
|
||||
if not is_valid_referral_withdrawal_recipient(withdrawal_method, recipient):
|
||||
return "INVALID_RECIPIENT", None
|
||||
|
||||
withdrawal_id: Optional[int] = None
|
||||
|
||||
async with session_factory() as session:
|
||||
try:
|
||||
await session.execute(text("BEGIN IMMEDIATE"))
|
||||
|
||||
user_result = await session.execute(
|
||||
text(
|
||||
"SELECT user_referral_balance, user_referral_hold "
|
||||
"FROM storage_users WHERE user_id = :user_id"
|
||||
),
|
||||
{"user_id": user_id},
|
||||
)
|
||||
get_user = user_result.mappings().first()
|
||||
|
||||
if get_user is None:
|
||||
await session.rollback()
|
||||
return "USER_NOT_FOUND", None
|
||||
|
||||
available_balance = round(
|
||||
float(get_user["user_referral_balance"])
|
||||
- float(get_user["user_referral_hold"]),
|
||||
2,
|
||||
)
|
||||
|
||||
if amount > available_balance:
|
||||
await session.rollback()
|
||||
return "INSUFFICIENT_FUNDS", None
|
||||
|
||||
now_unix = get_unix()
|
||||
withdrawal_result = await session.execute(
|
||||
text(
|
||||
"INSERT INTO storage_referral_withdrawals "
|
||||
"(user_id, withdrawal_amount, withdrawal_method, withdrawal_recipient, withdrawal_status, withdrawal_unix) "
|
||||
"VALUES (:user_id, :amount, :method, :recipient, 'pending', :unix)"
|
||||
),
|
||||
{
|
||||
"user_id": user_id,
|
||||
"amount": amount,
|
||||
"method": withdrawal_method,
|
||||
"recipient": recipient,
|
||||
"unix": now_unix,
|
||||
},
|
||||
)
|
||||
withdrawal_id = int(withdrawal_result.lastrowid)
|
||||
|
||||
await session.execute(
|
||||
text(
|
||||
"UPDATE storage_users "
|
||||
"SET user_referral_hold = :hold "
|
||||
"WHERE user_id = :user_id"
|
||||
),
|
||||
{
|
||||
"hold": round(
|
||||
float(get_user["user_referral_hold"]) + amount,
|
||||
2,
|
||||
),
|
||||
"user_id": user_id,
|
||||
},
|
||||
)
|
||||
await session.execute(
|
||||
text(
|
||||
"INSERT INTO storage_referral_transactions "
|
||||
"(user_id, related_user_id, transaction_type, source_id, amount, transaction_unix) "
|
||||
"VALUES (:user_id, NULL, :transaction_type, :source_id, :amount, :unix)"
|
||||
),
|
||||
{
|
||||
"user_id": user_id,
|
||||
"transaction_type": "withdrawal_hold",
|
||||
"source_id": f"withdrawal:{withdrawal_id}",
|
||||
"amount": amount,
|
||||
"unix": now_unix,
|
||||
},
|
||||
)
|
||||
await session.commit()
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
raise
|
||||
|
||||
return "ok", await self.get(increment=withdrawal_id)
|
||||
|
||||
async def complete(
|
||||
self, withdrawal_id: int, admin_id: int
|
||||
) -> Tuple[str, Optional[ReferralWithdrawal]]:
|
||||
async with session_factory() as session:
|
||||
try:
|
||||
await session.execute(text("BEGIN IMMEDIATE"))
|
||||
|
||||
withdrawal_result = await session.execute(
|
||||
text(
|
||||
"SELECT user_id, withdrawal_amount, withdrawal_status "
|
||||
"FROM storage_referral_withdrawals WHERE increment = :withdrawal_id"
|
||||
),
|
||||
{"withdrawal_id": withdrawal_id},
|
||||
)
|
||||
withdrawal = withdrawal_result.mappings().first()
|
||||
|
||||
if withdrawal is None:
|
||||
await session.rollback()
|
||||
return "NOT_FOUND", None
|
||||
if withdrawal["withdrawal_status"] != "pending":
|
||||
await session.rollback()
|
||||
return "ALREADY_PROCESSED", None
|
||||
|
||||
user_result = await session.execute(
|
||||
text(
|
||||
"SELECT user_referral_balance, user_referral_hold "
|
||||
"FROM storage_users WHERE user_id = :user_id"
|
||||
),
|
||||
{"user_id": withdrawal["user_id"]},
|
||||
)
|
||||
get_user = user_result.mappings().first()
|
||||
amount = float(withdrawal["withdrawal_amount"])
|
||||
|
||||
if (
|
||||
get_user is None
|
||||
or float(get_user["user_referral_balance"]) < amount
|
||||
or float(get_user["user_referral_hold"]) < amount
|
||||
):
|
||||
await session.rollback()
|
||||
return "BALANCE_ERROR", None
|
||||
|
||||
now_unix = get_unix()
|
||||
await session.execute(
|
||||
text(
|
||||
"UPDATE storage_users "
|
||||
"SET user_referral_balance = :balance, user_referral_hold = :hold "
|
||||
"WHERE user_id = :user_id"
|
||||
),
|
||||
{
|
||||
"balance": round(
|
||||
float(get_user["user_referral_balance"]) - amount,
|
||||
2,
|
||||
),
|
||||
"hold": round(
|
||||
float(get_user["user_referral_hold"]) - amount,
|
||||
2,
|
||||
),
|
||||
"user_id": withdrawal["user_id"],
|
||||
},
|
||||
)
|
||||
await session.execute(
|
||||
text(
|
||||
"UPDATE storage_referral_withdrawals "
|
||||
"SET withdrawal_status = 'completed', processed_unix = :unix, processed_admin_id = :admin_id "
|
||||
"WHERE increment = :withdrawal_id"
|
||||
),
|
||||
{
|
||||
"unix": now_unix,
|
||||
"admin_id": admin_id,
|
||||
"withdrawal_id": withdrawal_id,
|
||||
},
|
||||
)
|
||||
await session.execute(
|
||||
text(
|
||||
"INSERT INTO storage_referral_transactions "
|
||||
"(user_id, related_user_id, transaction_type, source_id, amount, transaction_unix) "
|
||||
"VALUES (:user_id, NULL, :transaction_type, :source_id, :amount, :unix)"
|
||||
),
|
||||
{
|
||||
"user_id": withdrawal["user_id"],
|
||||
"transaction_type": "withdrawal_completed",
|
||||
"source_id": f"withdrawal:{withdrawal_id}",
|
||||
"amount": amount,
|
||||
"unix": now_unix,
|
||||
},
|
||||
)
|
||||
await session.commit()
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
raise
|
||||
|
||||
return "ok", await self.get(increment=withdrawal_id)
|
||||
|
||||
async def reject(
|
||||
self, withdrawal_id: int, admin_id: int
|
||||
) -> Tuple[str, Optional[ReferralWithdrawal]]:
|
||||
async with session_factory() as session:
|
||||
try:
|
||||
await session.execute(text("BEGIN IMMEDIATE"))
|
||||
|
||||
withdrawal_result = await session.execute(
|
||||
text(
|
||||
"SELECT user_id, withdrawal_amount, withdrawal_status "
|
||||
"FROM storage_referral_withdrawals WHERE increment = :withdrawal_id"
|
||||
),
|
||||
{"withdrawal_id": withdrawal_id},
|
||||
)
|
||||
withdrawal = withdrawal_result.mappings().first()
|
||||
|
||||
if withdrawal is None:
|
||||
await session.rollback()
|
||||
return "NOT_FOUND", None
|
||||
if withdrawal["withdrawal_status"] != "pending":
|
||||
await session.rollback()
|
||||
return "ALREADY_PROCESSED", None
|
||||
|
||||
user_result = await session.execute(
|
||||
text(
|
||||
"SELECT user_referral_hold FROM storage_users "
|
||||
"WHERE user_id = :user_id"
|
||||
),
|
||||
{"user_id": withdrawal["user_id"]},
|
||||
)
|
||||
get_user = user_result.mappings().first()
|
||||
amount = float(withdrawal["withdrawal_amount"])
|
||||
|
||||
if get_user is None or float(get_user["user_referral_hold"]) < amount:
|
||||
await session.rollback()
|
||||
return "BALANCE_ERROR", None
|
||||
|
||||
now_unix = get_unix()
|
||||
await session.execute(
|
||||
text(
|
||||
"UPDATE storage_users "
|
||||
"SET user_referral_hold = :hold WHERE user_id = :user_id"
|
||||
),
|
||||
{
|
||||
"hold": round(
|
||||
float(get_user["user_referral_hold"]) - amount,
|
||||
2,
|
||||
),
|
||||
"user_id": withdrawal["user_id"],
|
||||
},
|
||||
)
|
||||
await session.execute(
|
||||
text(
|
||||
"UPDATE storage_referral_withdrawals "
|
||||
"SET withdrawal_status = 'rejected', processed_unix = :unix, processed_admin_id = :admin_id "
|
||||
"WHERE increment = :withdrawal_id"
|
||||
),
|
||||
{
|
||||
"unix": now_unix,
|
||||
"admin_id": admin_id,
|
||||
"withdrawal_id": withdrawal_id,
|
||||
},
|
||||
)
|
||||
await session.execute(
|
||||
text(
|
||||
"INSERT INTO storage_referral_transactions "
|
||||
"(user_id, related_user_id, transaction_type, source_id, amount, transaction_unix) "
|
||||
"VALUES (:user_id, NULL, :transaction_type, :source_id, :amount, :unix)"
|
||||
),
|
||||
{
|
||||
"user_id": withdrawal["user_id"],
|
||||
"transaction_type": "withdrawal_rejected",
|
||||
"source_id": f"withdrawal:{withdrawal_id}",
|
||||
"amount": amount,
|
||||
"unix": now_unix,
|
||||
},
|
||||
)
|
||||
await session.commit()
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
raise
|
||||
|
||||
return "ok", await self.get(increment=withdrawal_id)
|
||||
|
||||
async def add(
|
||||
self,
|
||||
user_id: int,
|
||||
withdrawal_amount: float,
|
||||
withdrawal_method: str,
|
||||
withdrawal_recipient: str,
|
||||
) -> ReferralWithdrawal:
|
||||
return await self._insert(
|
||||
user_id=user_id,
|
||||
withdrawal_amount=round(withdrawal_amount, 2),
|
||||
withdrawal_method=withdrawal_method,
|
||||
withdrawal_recipient=withdrawal_recipient,
|
||||
withdrawal_status="pending",
|
||||
withdrawal_unix=get_unix(),
|
||||
)
|
||||
|
||||
async def update(
|
||||
self,
|
||||
where: Optional[Union[Dict[str, Any], int]] = None,
|
||||
**kwargs,
|
||||
) -> int:
|
||||
if isinstance(where, int):
|
||||
where = {"increment": where}
|
||||
|
||||
return await self._update(where=where, **kwargs)
|
||||
@@ -90,7 +90,7 @@ class Refillx(BaseRepository[RefillModel, Refill]):
|
||||
user_result = await session.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT user_balance, user_refill
|
||||
SELECT user_balance, user_refill, user_referrer_id
|
||||
FROM storage_users
|
||||
WHERE user_id = :user_id
|
||||
"""
|
||||
@@ -103,6 +103,14 @@ class Refillx(BaseRepository[RefillModel, Refill]):
|
||||
await session.rollback()
|
||||
return "USER_NOT_FOUND"
|
||||
|
||||
settings_result = await session.execute(
|
||||
text(
|
||||
"SELECT status_referral, referral_refill_percent "
|
||||
"FROM storage_settings WHERE id = 1"
|
||||
)
|
||||
)
|
||||
get_settings = settings_result.mappings().first()
|
||||
|
||||
new_balance = round(float(get_user["user_balance"]) + float(pay_amount), 2)
|
||||
new_refill = round(float(get_user["user_refill"]) + float(pay_amount), 2)
|
||||
|
||||
@@ -138,6 +146,66 @@ class Refillx(BaseRepository[RefillModel, Refill]):
|
||||
),
|
||||
{"balance": new_balance, "refill": new_refill, "user_id": user_id},
|
||||
)
|
||||
|
||||
referrer_id = get_user["user_referrer_id"]
|
||||
referral_reward = 0.0
|
||||
|
||||
if (
|
||||
get_settings is not None
|
||||
and get_settings["status_referral"] == "True"
|
||||
and referrer_id is not None
|
||||
and referrer_id != user_id
|
||||
):
|
||||
referral_reward = round(
|
||||
float(pay_amount)
|
||||
* float(get_settings["referral_refill_percent"])
|
||||
/ 100,
|
||||
2,
|
||||
)
|
||||
|
||||
if referral_reward > 0:
|
||||
referrer_result = await session.execute(
|
||||
text(
|
||||
"SELECT user_referral_balance FROM storage_users "
|
||||
"WHERE user_id = :user_id"
|
||||
),
|
||||
{"user_id": referrer_id},
|
||||
)
|
||||
get_referrer = referrer_result.mappings().first()
|
||||
|
||||
if get_referrer is not None:
|
||||
source_id = f"refill:{pay_comment or pay_receipt}"
|
||||
|
||||
await session.execute(
|
||||
text(
|
||||
"UPDATE storage_users "
|
||||
"SET user_referral_balance = :balance "
|
||||
"WHERE user_id = :user_id"
|
||||
),
|
||||
{
|
||||
"balance": round(
|
||||
float(get_referrer["user_referral_balance"])
|
||||
+ referral_reward,
|
||||
2,
|
||||
),
|
||||
"user_id": referrer_id,
|
||||
},
|
||||
)
|
||||
await session.execute(
|
||||
text(
|
||||
"INSERT INTO storage_referral_transactions "
|
||||
"(user_id, related_user_id, transaction_type, source_id, amount, transaction_unix) "
|
||||
"VALUES (:user_id, :related_user_id, :transaction_type, :source_id, :amount, :transaction_unix)"
|
||||
),
|
||||
{
|
||||
"user_id": referrer_id,
|
||||
"related_user_id": user_id,
|
||||
"transaction_type": "refill_percent",
|
||||
"source_id": source_id,
|
||||
"amount": referral_reward,
|
||||
"transaction_unix": get_unix(),
|
||||
},
|
||||
)
|
||||
await session.commit()
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# - *- coding: utf- 8 - *-
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from sqlalchemy import Integer, String
|
||||
from sqlalchemy import Integer, String, Float
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from tgbot.database.core import Base
|
||||
@@ -46,6 +46,15 @@ class SettingsModel(Base):
|
||||
misc_profit_week: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
||||
misc_profit_month: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
||||
|
||||
# Реферальная система
|
||||
referral_refill_percent: Mapped[float] = mapped_column(
|
||||
Float, nullable=False, default=0
|
||||
)
|
||||
status_referral: Mapped[str] = mapped_column(
|
||||
String(16), nullable=False, default="False"
|
||||
)
|
||||
referral_bonus_rub: Mapped[float] = mapped_column(Float, nullable=False, default=0)
|
||||
|
||||
|
||||
ModelBase = Settings
|
||||
BaseModel = Settings
|
||||
|
||||
+33
-17
@@ -15,8 +15,12 @@ from tgbot.utils.const_functions import get_unix
|
||||
class UserModel(Base):
|
||||
__tablename__ = "storage_users"
|
||||
|
||||
increment: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
||||
user_id: Mapped[int] = mapped_column(BigInteger, nullable=False, unique=True, index=True)
|
||||
increment: Mapped[int] = mapped_column(
|
||||
Integer, primary_key=True, autoincrement=True
|
||||
)
|
||||
user_id: Mapped[int] = mapped_column(
|
||||
BigInteger, nullable=False, unique=True, index=True
|
||||
)
|
||||
user_login: Mapped[str] = mapped_column(String(255), nullable=False, default="")
|
||||
user_name: Mapped[str] = mapped_column(String(255), nullable=False, default="")
|
||||
user_surname: Mapped[str] = mapped_column(String(255), nullable=False, default="")
|
||||
@@ -26,6 +30,13 @@ class UserModel(Base):
|
||||
user_give: Mapped[float] = mapped_column(Float, nullable=False, default=0)
|
||||
user_unix: Mapped[int] = mapped_column(Integer, nullable=False, default=get_unix)
|
||||
|
||||
# Поля для реферальной системы
|
||||
user_referrer_id: Mapped[int] = mapped_column(BigInteger, nullable=True, index=True)
|
||||
user_referral_balance: Mapped[float] = mapped_column(
|
||||
Float, nullable=False, default=0
|
||||
)
|
||||
user_referral_hold: Mapped[float] = mapped_column(Float, nullable=False, default=0)
|
||||
|
||||
|
||||
ModelBase = User
|
||||
BaseModel = User
|
||||
@@ -41,31 +52,34 @@ class UsersRepository(BaseRepository[UserModel, User]):
|
||||
|
||||
# Добавление или обновление пользователя
|
||||
async def add(
|
||||
self,
|
||||
user_id: int,
|
||||
user_login: str,
|
||||
user_name: str,
|
||||
user_surname: str = "",
|
||||
user_fullname: str = "",
|
||||
self,
|
||||
user_id: int,
|
||||
user_login: str,
|
||||
user_name: str,
|
||||
user_surname: str = "",
|
||||
user_fullname: str = "",
|
||||
) -> User:
|
||||
return await self.upsert(
|
||||
user_id=user_id,
|
||||
user_login=user_login,
|
||||
user_name=user_name,
|
||||
user_surname=user_surname,
|
||||
user_fullname=user_fullname or " ".join(filter(None, [user_name, user_surname])),
|
||||
user_fullname=user_fullname
|
||||
or " ".join(filter(None, [user_name, user_surname])),
|
||||
)
|
||||
|
||||
# Выполнение upsert пользователя по телеграм ID
|
||||
async def upsert(
|
||||
self,
|
||||
user_id: int,
|
||||
user_login: str,
|
||||
user_name: str,
|
||||
user_surname: str = "",
|
||||
user_fullname: str = "",
|
||||
self,
|
||||
user_id: int,
|
||||
user_login: str,
|
||||
user_name: str,
|
||||
user_surname: str = "",
|
||||
user_fullname: str = "",
|
||||
) -> User:
|
||||
user_fullname = user_fullname or " ".join(filter(None, [user_name, user_surname]))
|
||||
user_fullname = user_fullname or " ".join(
|
||||
filter(None, [user_name, user_surname])
|
||||
)
|
||||
user_table = UserModel.__table__
|
||||
statement = insert(UserModel).values(
|
||||
user_id=user_id,
|
||||
@@ -105,7 +119,9 @@ class UsersRepository(BaseRepository[UserModel, User]):
|
||||
return user
|
||||
|
||||
# Обновление пользователя по ID или фильтру
|
||||
async def update(self, where: Optional[Union[Dict[str, Any], int]] = None, **kwargs) -> int:
|
||||
async def update(
|
||||
self, where: Optional[Union[Dict[str, Any], int]] = None, **kwargs
|
||||
) -> int:
|
||||
if isinstance(where, int):
|
||||
where = {"user_id": where}
|
||||
|
||||
|
||||
@@ -97,6 +97,11 @@ class Settings:
|
||||
misc_profit_week: int
|
||||
misc_profit_month: int
|
||||
|
||||
# Реферальная система
|
||||
referral_refill_percent: float
|
||||
status_referral: str
|
||||
referral_bonus_rub: float
|
||||
|
||||
|
||||
@dataclass
|
||||
class User:
|
||||
@@ -110,3 +115,40 @@ class User:
|
||||
user_refill: float
|
||||
user_give: float
|
||||
user_unix: int
|
||||
|
||||
# Поля для реферальной системы
|
||||
user_referrer_id: Optional[int]
|
||||
user_referral_balance: float
|
||||
user_referral_hold: float
|
||||
|
||||
|
||||
@dataclass
|
||||
class Referral:
|
||||
increment: int
|
||||
referrer_id: int
|
||||
referral_id: int
|
||||
referral_unix: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReferralTransaction:
|
||||
increment: int
|
||||
user_id: int
|
||||
related_user_id: Optional[int]
|
||||
transaction_type: str
|
||||
source_id: str
|
||||
amount: float
|
||||
transaction_unix: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReferralWithdrawal:
|
||||
increment: int
|
||||
user_id: int
|
||||
withdrawal_amount: float
|
||||
withdrawal_method: str
|
||||
withdrawal_recipient: str
|
||||
withdrawal_status: str
|
||||
withdrawal_unix: int
|
||||
processed_unix: Optional[int]
|
||||
processed_admin_id: Optional[int]
|
||||
|
||||
@@ -160,10 +160,13 @@ async def prepare_database() -> None:
|
||||
import tgbot.database.db_refill # noqa: F401
|
||||
import tgbot.database.db_settings # noqa: F401
|
||||
import tgbot.database.db_users # noqa: F401
|
||||
import tgbot.database.db_referral # noqa: F401
|
||||
|
||||
async with engine.begin() as conn:
|
||||
from tgbot.database.db_payments import ensure_payments_schema
|
||||
from tgbot.database.db_referral import ensure_referrals_schema
|
||||
|
||||
await ensure_referrals_schema(conn)
|
||||
await ensure_payments_schema(conn)
|
||||
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
|
||||
Reference in New Issue
Block a user