Optimizing database queries/ Combining cash registers into a single point/ Caching key moments/ API versioning/ ruff formatting

2026-02-20 00:27:10 +03:00
parent 2580835f55
commit 05a19f01b1
102 changed files with 3815 additions and 858 deletions
@@ -2,6 +2,7 @@ from .bans import *
 from .coupons import *
 from .db import async_session_maker
 from .gifts import *
+from . import identities
 from .hot_leads import *
 from .init_db import *
 from .keys import *
@@ -9,6 +10,7 @@ from .notifications import *
 from .payments import *
 from .referrals import *
 from .servers import *
+from .settings_cache import settings_cache
 from .statistics import *
 from .tariffs import *
 from .temporary_data import *
@@ -8,7 +8,7 @@ from config import DATABASE_URL, DB_MAX_OVERFLOW, DB_POOL_SIZE


 CONCURRENT_UPDATES_LIMIT = DB_POOL_SIZE + DB_MAX_OVERFLOW
-MAX_UPDATE_AGE_SEC = 28
+MAX_UPDATE_AGE_SEC = 15

 engine = create_async_engine(
    DATABASE_URL,
@@ -0,0 +1,197 @@
+import hashlib
+import secrets
+from datetime import datetime, timedelta
+
+import bcrypt
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from config import API_TOKEN_TTL_DAYS
+from database.models import Admin, Identity, User
+
+
+_BCRYPT_MAX_PASSWORD_BYTES = 72
+_BCRYPT_ROUNDS = 12
+
+
+def _password_bytes(password: str) -> bytes:
+    """Пароль в байтах, не длиннее 72 байт (ограничение bcrypt)."""
+    raw = password.encode("utf-8")
+    if len(raw) > _BCRYPT_MAX_PASSWORD_BYTES:
+        return raw[:_BCRYPT_MAX_PASSWORD_BYTES]
+    return raw
+
+
+def hash_token(token: str) -> str:
+    return hashlib.sha256(token.encode()).hexdigest()
+
+
+def hash_password(password: str) -> str:
+    """Хеш пароля через bcrypt (соль уникальна на каждый пароль)."""
+    salt = bcrypt.gensalt(rounds=_BCRYPT_ROUNDS)
+    return bcrypt.hashpw(_password_bytes(password), salt).decode("ascii")
+
+
+def check_password(password: str, password_hash: str | None) -> bool:
+    if not password_hash:
+        return False
+    try:
+        return bcrypt.checkpw(_password_bytes(password), password_hash.encode("ascii"))
+    except Exception:
+        return False
+
+
+def generate_token() -> str:
+    return secrets.token_urlsafe(32)
+
+
+async def create_identity(
+    session: AsyncSession,
+    email: str | None = None,
+    tg_id: int | None = None,
+) -> Identity:
+    """Создаёт идентичность; можно задать email и/или tg_id."""
+    identity = Identity(email=email.strip().lower() if email else None, tg_id=tg_id)
+    session.add(identity)
+    await session.flush()
+    if tg_id:
+        await session.execute(User.__table__.update().where(User.tg_id == tg_id).values(identity_id=identity.id))
+    await session.commit()
+    await session.refresh(identity)
+    return identity
+
+
+async def get_identity_by_id(session: AsyncSession, identity_id: str) -> Identity | None:
+    """Возвращает идентичность по id."""
+    result = await session.execute(select(Identity).where(Identity.id == identity_id))
+    return result.scalar_one_or_none()
+
+
+async def get_identity_by_email(session: AsyncSession, email: str) -> Identity | None:
+    """Возвращает идентичность по email."""
+    if not email or not email.strip():
+        return None
+    result = await session.execute(select(Identity).where(Identity.email == email.strip().lower()))
+    return result.scalar_one_or_none()
+
+
+async def get_identity_by_tg_id(session: AsyncSession, tg_id: int) -> Identity | None:
+    """Возвращает идентичность по tg_id."""
+    result = await session.execute(select(Identity).where(Identity.tg_id == tg_id))
+    return result.scalar_one_or_none()
+
+
+async def get_identity_by_token_hash(session: AsyncSession, token_hash: str) -> Identity | None:
+    """Возвращает идентичность по хешу API-токена."""
+    result = await session.execute(select(Identity).where(Identity.api_token_hash == token_hash))
+    return result.scalar_one_or_none()
+
+
+async def issue_token_for_identity(session: AsyncSession, identity: Identity) -> str:
+    """Генерирует токен, сохраняет хеш и token_issued_at в identity, возвращает токен (показать один раз)."""
+    token = generate_token()
+    identity.api_token_hash = hash_token(token)
+    identity.token_issued_at = datetime.utcnow()
+    await session.commit()
+    await session.refresh(identity)
+    return token
+
+
+def _is_token_expired(identity: Identity) -> bool:
+    """Проверяет, истёк ли срок действия токена (если задан API_TOKEN_TTL_DAYS)."""
+    if API_TOKEN_TTL_DAYS is None or identity.token_issued_at is None:
+        return False
+    expiry = identity.token_issued_at + timedelta(days=API_TOKEN_TTL_DAYS)
+    return datetime.utcnow() >= expiry
+
+
+async def create_identity_with_token(
+    session: AsyncSession,
+    email: str | None = None,
+    password: str | None = None,
+    tg_id: int | None = None,
+) -> tuple[Identity, str]:
+    """Создаёт идентичность и выдаёт API-токен. При регистрации по почте передать email и password."""
+    identity = await create_identity(session, email=email, tg_id=tg_id)
+    if password:
+        identity.password_hash = hash_password(password)
+        await session.commit()
+        await session.refresh(identity)
+    token = await issue_token_for_identity(session, identity)
+    return identity, token
+
+
+async def verify_identity_token(session: AsyncSession, identity_id: str, token: str) -> Identity | None:
+    """Проверяет пару identity_id + token и срок действия токена; возвращает Identity или None."""
+    identity = await get_identity_by_id(session, identity_id)
+    if not identity or not identity.api_token_hash:
+        return None
+    if hash_token(token) != identity.api_token_hash:
+        return None
+    if _is_token_expired(identity):
+        return None
+    return identity
+
+
+async def login_by_email(session: AsyncSession, email: str, password: str) -> tuple[Identity, str] | None:
+    """Вход по email и паролю: проверяет пароль, выдаёт новый токен; возвращает (identity, token) или None."""
+    identity = await get_identity_by_email(session, email)
+    if not identity or not check_password(password, identity.password_hash):
+        return None
+    token = await issue_token_for_identity(session, identity)
+    return identity, token
+
+
+async def resolve_tg_id(session: AsyncSession, identity_id: str) -> int | None:
+    """По identity_id возвращает tg_id, если привязан."""
+    identity = await get_identity_by_id(session, identity_id)
+    return identity.tg_id if identity else None
+
+
+async def attach_email(session: AsyncSession, identity_id: str, email: str) -> Identity | None:
+    """Привязывает email к идентичности."""
+    identity = await get_identity_by_id(session, identity_id)
+    if not identity:
+        return None
+    email_clean = email.strip().lower() if email else None
+    if not email_clean:
+        return identity
+    existing = await get_identity_by_email(session, email_clean)
+    if existing and existing.id != identity_id:
+        return None
+    identity.email = email_clean
+    await session.commit()
+    await session.refresh(identity)
+    return identity
+
+
+async def attach_telegram(session: AsyncSession, identity_id: str, tg_id: int) -> Identity | None:
+    """Привязывает Telegram (tg_id) к идентичности и связывает User с identity. Если tg_id в admins — выставляет is_admin."""
+    identity = await get_identity_by_id(session, identity_id)
+    if not identity:
+        return None
+    existing = await get_identity_by_tg_id(session, tg_id)
+    if existing and existing.id != identity_id:
+        return None
+    identity.tg_id = tg_id
+    admin_row = await session.execute(select(Admin).where(Admin.tg_id == tg_id))
+    if admin_row.scalar_one_or_none():
+        identity.is_admin = True
+    await session.execute(User.__table__.update().where(User.tg_id == tg_id).values(identity_id=identity_id))
+    await session.commit()
+    await session.refresh(identity)
+    return identity
+
+
+async def get_or_create_identity_for_tg(session: AsyncSession, tg_id: int) -> Identity:
+    """Для tg_id возвращает существующую идентичность или создаёт новую и привязывает User."""
+    identity = await get_identity_by_tg_id(session, tg_id)
+    if identity:
+        return identity
+    identity = Identity(tg_id=tg_id)
+    session.add(identity)
+    await session.flush()
+    await session.execute(User.__table__.update().where(User.tg_id == tg_id).values(identity_id=identity.id))
+    await session.commit()
+    await session.refresh(identity)
+    return identity
@@ -5,6 +5,7 @@ from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.ext.asyncio import AsyncSession

 from database.models import Key, User
+from database.users import invalidate_user_snapshot
 from logger import logger


@@ -81,6 +82,7 @@ async def store_key(
            logger.info(f"[Store Key] Ключ создан: tg_id={tg_id}, client_id={client_id}, server_id={server_id}")

        await session.commit()
+        invalidate_user_snapshot(tg_id)

    except SQLAlchemyError as e:
        logger.error(f"❌ Ошибка при сохранении ключа: {e}")
@@ -157,10 +159,18 @@ async def get_key_count(session: AsyncSession, tg_id: int) -> int:


 async def delete_key(session: AsyncSession, identifier: int | str, commit: bool = True):
-    stmt = delete(Key).where(Key.tg_id == identifier if str(identifier).isdigit() else Key.client_id == identifier)
+    tg_id_for_cache = None
+    if isinstance(identifier, str):
+        res = await session.execute(select(Key.tg_id).where(Key.client_id == identifier).limit(1))
+        tg_id_for_cache = res.scalar_one_or_none()
+    else:
+        tg_id_for_cache = identifier
+    stmt = delete(Key).where(Key.tg_id == identifier if isinstance(identifier, int) else Key.client_id == identifier)
    await session.execute(stmt)
    if commit:
        await session.commit()
+    if tg_id_for_cache is not None:
+        invalidate_user_snapshot(tg_id_for_cache)
    logger.info(f"Ключ с идентификатором {identifier} удалён")


@@ -36,10 +36,32 @@ class DictLikeMixin:
        return {column.name: getattr(self, column.name) for column in self.__table__.columns}


+class Identity(DictLikeMixin, Base):
+    """Слой идентификации: к одному identity можно привязать email и/или Telegram (tg_id)."""
+
+    __tablename__ = "identities"
+
+    id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
+    email = Column(String(255), unique=True, nullable=True, index=True)
+    tg_id = Column(BigInteger, unique=True, nullable=True, index=True)
+    api_token_hash = Column(String(64), nullable=True, index=True)
+    token_issued_at = Column(DateTime, nullable=True)
+    password_hash = Column(String(64), nullable=True)
+    is_admin = Column(Boolean, nullable=False, server_default=text("false"))
+    created_at = Column(DateTime, default=datetime.utcnow)
+    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
+
+
 class User(DictLikeMixin, Base):
    __tablename__ = "users"

    tg_id = Column(BigInteger, primary_key=True)
+    identity_id = Column(
+        String(36),
+        ForeignKey("identities.id", ondelete="SET NULL", onupdate="CASCADE"),
+        nullable=True,
+        index=True,
+    )
    username = Column(String)
    first_name = Column(String)
    last_name = Column(String)
@@ -162,7 +184,7 @@ class ServerSpecialgroup(DictLikeMixin, Base):
 class Payment(DictLikeMixin, Base):
    __tablename__ = "payments"

-    id = Column(Integer, primary_key=True)
+    id = Column(Integer, primary_key=True, autoincrement=True)
    tg_id = Column(BigInteger, ForeignKey("users.tg_id"))
    amount = Column(Float)
    payment_system = Column(String)
@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import threading
+from datetime import datetime
+from typing import Any
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from .models import Setting
+
+
+class SettingsCache:
+    _cache: dict[str, dict[str, Any]]
+    _lock: threading.Lock
+
+    def __init__(self) -> None:
+        self._cache = {}
+        self._lock = threading.Lock()
+
+    @staticmethod
+    def _row_to_item(s: Setting) -> dict[str, Any]:
+        return {
+            "key": s.key,
+            "value": s.value,
+            "description": s.description,
+            "created_at": getattr(s, "created_at", None),
+            "updated_at": getattr(s, "updated_at", None),
+        }
+
+    async def load(self, session: AsyncSession) -> None:
+        result = await session.execute(select(Setting))
+        rows = result.scalars().all()
+        with self._lock:
+            self._cache.clear()
+            for s in rows:
+                self._cache[s.key] = self._row_to_item(s)
+
+    def get_all(self) -> list[dict[str, Any]]:
+        with self._lock:
+            return [dict(x) for x in self._cache.values()]
+
+    def get(self, key: str) -> dict[str, Any] | None:
+        with self._lock:
+            return dict(self._cache[key]) if key in self._cache else None
+
+    def update(
+        self,
+        key: str,
+        value: Any,
+        description: str | None = None,
+        created_at: datetime | None = None,
+        updated_at: datetime | None = None,
+    ) -> None:
+        now = datetime.utcnow()
+        with self._lock:
+            if key in self._cache:
+                self._cache[key]["value"] = value
+                self._cache[key]["updated_at"] = updated_at if updated_at is not None else now
+                if description is not None:
+                    self._cache[key]["description"] = description
+            else:
+                self._cache[key] = {
+                    "key": key,
+                    "value": value,
+                    "description": description,
+                    "created_at": created_at if created_at is not None else now,
+                    "updated_at": updated_at if updated_at is not None else now,
+                }
+
+    def delete(self, key: str) -> None:
+        with self._lock:
+            self._cache.pop(key, None)
+
+
+settings_cache = SettingsCache()
+
+
+async def load(session: AsyncSession) -> None:
+    await settings_cache.load(session)
@@ -1,5 +1,6 @@
 from datetime import datetime

+from cachetools import TTLCache
 from sqlalchemy import delete, exists, func, or_, select, update
 from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.exc import SQLAlchemyError
@@ -20,6 +21,13 @@ from database.models import (
 )
 from logger import logger

+_SNAPSHOT_CACHE: TTLCache[int, tuple[int, int]] = TTLCache(maxsize=150_000, ttl=30)
+_EXISTS_CACHE: TTLCache[int, bool] = TTLCache(maxsize=150_000, ttl=60)
+
+
+def invalidate_user_snapshot(tg_id: int) -> None:
+    _SNAPSHOT_CACHE.pop(tg_id, None)
+

 async def add_user(
    session: AsyncSession,
@@ -53,6 +61,7 @@ async def add_user(
            return False
        if commit:
            await session.commit()
+        _EXISTS_CACHE[tg_id] = True
        logger.info(f"[DB] Новый пользователь добавлен: {tg_id} (source: {source_code})")
        return True
    except SQLAlchemyError as e:
@@ -82,9 +91,15 @@ async def update_balance(session: AsyncSession, tg_id: int, amount: float) -> No


 async def check_user_exists(session: AsyncSession, tg_id: int) -> bool:
+    try:
+        return _EXISTS_CACHE[tg_id]
+    except KeyError:
+        pass
    stmt = select(exists().where(User.tg_id == tg_id))
    result = await session.execute(stmt)
-    return result.scalar()
+    value = result.scalar()
+    _EXISTS_CACHE[tg_id] = value
+    return value


 async def get_balance(session: AsyncSession, tg_id: int) -> float:
@@ -107,6 +122,7 @@ async def update_trial(session: AsyncSession, tg_id: int, status: int):
    try:
        await session.execute(update(User).where(User.tg_id == tg_id).values(trial=status))
        await session.commit()
+        invalidate_user_snapshot(tg_id)
        logger.info(f"[DB] Триал статус обновлён для пользователя {tg_id}: {status}")
    except SQLAlchemyError as e:
        logger.error(f"[DB] Ошибка при обновлении триала пользователя {tg_id}: {e}")
@@ -158,6 +174,7 @@ async def upsert_user(
            if row is None:
                return None
            await session.commit()
+            _EXISTS_CACHE[tg_id] = True
            return dict(row)

        res = await session.execute(
@@ -187,6 +204,7 @@ async def upsert_user(
        )
        row = res.mappings().one()
        await session.commit()
+        _EXISTS_CACHE[tg_id] = True
        return dict(row)
    except SQLAlchemyError as e:
        logger.error(f"[DB] Ошибка при UPSERT пользователя {tg_id}: {e}")
@@ -222,16 +240,22 @@ async def delete_user_data(session: AsyncSession, tg_id: int):
 async def mark_trial_extended(tg_id: int, session: AsyncSession):
    await session.execute(update(User).where(User.tg_id == tg_id).values(trial=-1))
    await session.commit()
+    invalidate_user_snapshot(tg_id)


 async def get_user_snapshot(session: AsyncSession, tg_id: int) -> tuple[int, int] | None:
+    try:
+        return _SNAPSHOT_CACHE[tg_id]
+    except KeyError:
+        pass
    keys_count_sq = select(func.count(Key.client_id)).where(Key.tg_id == tg_id).scalar_subquery()
-
    res = await session.execute(select(func.coalesce(User.trial, 0), keys_count_sq).where(User.tg_id == tg_id))
    row = res.first()
    if row is None:
        return None
-    return int(row[0]), int(row[1])
+    value = (int(row[0]), int(row[1]))
+    _SNAPSHOT_CACHE[tg_id] = value
+    return value


 async def upsert_source_if_empty(