From 12376e507067f5b0159c2e1e6e434d54dcee8b1e Mon Sep 17 00:00:00 2001 From: 3252a8 <3252a8@proton.me> Date: Sun, 26 Apr 2026 19:47:14 +0300 Subject: [PATCH] db: add composite indexes and merge-user optimizations --- db/dal/user_dal.py | 21 ++++--- db/database_setup.py | 2 + db/migrator.py | 68 +++++++++++++++++++++ db/models.py | 10 +++- tests/test_user_dal.py | 133 +++++++++++++++++++++++++++++++++++++++++ 5 files changed, 222 insertions(+), 12 deletions(-) create mode 100644 tests/test_user_dal.py diff --git a/db/dal/user_dal.py b/db/dal/user_dal.py index 31d346a..108a7cd 100644 --- a/db/dal/user_dal.py +++ b/db/dal/user_dal.py @@ -629,23 +629,22 @@ async def get_user_ids_without_active_subscription(session: AsyncSession) -> Lis from datetime import datetime, timezone now = datetime.now(timezone.utc) - # Subquery for users with active subscription - active_subs_subq = ( - select(Subscription.user_id) - .where( - and_( - Subscription.is_active == True, - Subscription.end_date > now, - ) - ) - ).scalar_subquery() + active_subs = aliased(Subscription) stmt = ( select(User.user_id) + .outerjoin( + active_subs, + and_( + active_subs.user_id == User.user_id, + active_subs.is_active == True, + active_subs.end_date > now, + ), + ) .where( and_( User.is_banned == False, - ~User.user_id.in_(active_subs_subq), + active_subs.user_id.is_(None), ) ) ) diff --git a/db/database_setup.py b/db/database_setup.py index 711c11b..30fc731 100644 --- a/db/database_setup.py +++ b/db/database_setup.py @@ -20,6 +20,8 @@ def init_db_connection(settings: Settings) -> sessionmaker: settings.DATABASE_URL, echo=False, pool_pre_ping=True, + pool_size=20, + max_overflow=10, ) local_async_session_factory = async_sessionmaker( diff --git a/db/migrator.py b/db/migrator.py index e2cf15d..83da49e 100644 --- a/db/migrator.py +++ b/db/migrator.py @@ -251,6 +251,64 @@ def _migration_0007_add_telegram_photo_url(connection: Connection) -> None: ) +def _migration_0008_add_email_verification_code_status(connection: Connection) -> None: + inspector = inspect(connection) + columns: Set[str] = {col["name"] for col in inspector.get_columns("email_verification_codes")} + + if "status" not in columns: + connection.execute( + text( + "ALTER TABLE email_verification_codes ADD COLUMN status VARCHAR NOT NULL DEFAULT 'active'" + ) + ) + else: + connection.execute( + text( + """ + UPDATE email_verification_codes + SET status = 'active' + WHERE status IS NULL OR status = '' + """ + ) + ) + + connection.execute( + text( + """ + CREATE INDEX IF NOT EXISTS ix_email_verification_codes_status + ON email_verification_codes (status) + """ + ) + ) + + +def _migration_0009_add_composite_indexes(connection: Connection) -> None: + connection.execute( + text( + """ + CREATE INDEX IF NOT EXISTS ix_subscriptions_is_active_end_date + ON subscriptions (is_active, end_date) + """ + ) + ) + connection.execute( + text( + """ + CREATE INDEX IF NOT EXISTS ix_subscriptions_user_id_is_active + ON subscriptions (user_id, is_active) + """ + ) + ) + connection.execute( + text( + """ + CREATE INDEX IF NOT EXISTS ix_payments_user_id_status + ON payments (user_id, status) + """ + ) + ) + + MIGRATIONS: List[Migration] = [ Migration( id="0001_add_channel_subscription_fields", @@ -287,6 +345,16 @@ MIGRATIONS: List[Migration] = [ description="Store Telegram profile photo URLs for linked users", upgrade=_migration_0007_add_telegram_photo_url, ), + Migration( + id="0008_add_email_verification_code_status", + description="Track superseded email verification codes explicitly", + upgrade=_migration_0008_add_email_verification_code_status, + ), + Migration( + id="0009_add_composite_indexes", + description="Add composite indexes for subscription and payment lookups", + upgrade=_migration_0009_add_composite_indexes, + ), ] diff --git a/db/models.py b/db/models.py index 2f182ae..a3dc4be 100644 --- a/db/models.py +++ b/db/models.py @@ -1,4 +1,4 @@ -from sqlalchemy import create_engine, Column, Integer, String, Boolean, DateTime, Float, ForeignKey, UniqueConstraint, Text, BigInteger +from sqlalchemy import create_engine, Column, Integer, String, Boolean, DateTime, Float, ForeignKey, UniqueConstraint, Text, BigInteger, Index from sqlalchemy.orm import relationship, DeclarativeBase from sqlalchemy.ext.asyncio import AsyncAttrs from sqlalchemy.sql import func @@ -61,6 +61,10 @@ class User(Base): class Subscription(Base): __tablename__ = "subscriptions" + __table_args__ = ( + Index("ix_subscriptions_is_active_end_date", "is_active", "end_date"), + Index("ix_subscriptions_user_id_is_active", "user_id", "is_active"), + ) subscription_id = Column(Integer, primary_key=True, autoincrement=True) user_id = Column(BigInteger, @@ -105,6 +109,7 @@ class EmailVerificationCode(Base): ) expires_at = Column(DateTime(timezone=True), nullable=False, index=True) consumed_at = Column(DateTime(timezone=True), nullable=True) + status = Column(String, nullable=False, default="active", index=True) attempts = Column(Integer, nullable=False, default=0) created_at = Column(DateTime(timezone=True), server_default=func.now()) @@ -131,6 +136,9 @@ class SecurityThrottle(Base): class Payment(Base): __tablename__ = "payments" + __table_args__ = ( + Index("ix_payments_user_id_status", "user_id", "status"), + ) payment_id = Column(Integer, primary_key=True, autoincrement=True) user_id = Column(BigInteger, diff --git a/tests/test_user_dal.py b/tests/test_user_dal.py new file mode 100644 index 0000000..7f62f4d --- /dev/null +++ b/tests/test_user_dal.py @@ -0,0 +1,133 @@ +import unittest +from datetime import datetime, timezone +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +from sqlalchemy.dialects import postgresql +from sqlalchemy.sql.dml import Delete, Update + +from db.dal import user_dal + + +class FakeResult: + def __init__(self, scalar_value=None, rowcount=1): + self._scalar_value = scalar_value + self.rowcount = rowcount + + def scalar_one_or_none(self): + return self._scalar_value + + def scalars(self): + return self + + def all(self): + if self._scalar_value is None: + return [] + if isinstance(self._scalar_value, list): + return self._scalar_value + return [self._scalar_value] + + +class UserDalMergeTests(unittest.IsolatedAsyncioTestCase): + async def test_get_user_ids_without_active_subscription_uses_left_join_null_check(self): + session = SimpleNamespace( + execute=AsyncMock(return_value=FakeResult([2, 3])), + ) + + result = await user_dal.get_user_ids_without_active_subscription(session) + + self.assertEqual(result, [2, 3]) + stmt = session.execute.await_args.args[0] + sql = str( + stmt.compile( + dialect=postgresql.dialect(), + compile_kwargs={"literal_binds": True}, + ) + ).upper() + self.assertIn("LEFT OUTER JOIN", sql) + self.assertIn("IS NULL", sql) + + async def test_merge_users_uses_bulk_updates_for_related_tables(self): + source = SimpleNamespace( + user_id=1, + email="source@example.com", + telegram_id=111, + panel_user_uuid="panel-source", + email_verified_at=datetime.now(timezone.utc), + username="source-user", + first_name="Source", + last_name="User", + language_code="ru", + telegram_photo_url="https://example.com/source.jpg", + channel_subscription_verified=True, + channel_subscription_checked_at=datetime.now(timezone.utc), + channel_subscription_verified_for=1, + lifetime_used_traffic_bytes=512, + referred_by_id=999, + referral_code="SRC123", + ) + target = SimpleNamespace( + user_id=2, + email=None, + telegram_id=None, + panel_user_uuid=None, + email_verified_at=None, + username=None, + first_name=None, + last_name=None, + language_code=None, + telegram_photo_url=None, + channel_subscription_verified=False, + channel_subscription_checked_at=None, + channel_subscription_verified_for=None, + lifetime_used_traffic_bytes=128, + referred_by_id=None, + referral_code=None, + ) + session = SimpleNamespace( + execute=AsyncMock(side_effect=lambda stmt: FakeResult()), + delete=AsyncMock(), + flush=AsyncMock(), + refresh=AsyncMock(), + ) + + async def fake_get_user_by_id(_session, user_id): + if user_id == source.user_id: + return source + if user_id == target.user_id: + return target + return None + + with ( + patch("db.dal.user_dal.get_user_by_id", side_effect=fake_get_user_by_id), + patch("db.dal.user_dal._get_active_subscription_for_user", return_value=None), + patch("db.dal.user_dal._get_latest_subscription_for_user", return_value=None), + ): + merged = await user_dal.merge_users( + session, + source_user_id=source.user_id, + target_user_id=target.user_id, + ) + + self.assertIs(merged, target) + + update_tables = [] + delete_tables = [] + for call in session.execute.await_args_list: + stmt = call.args[0] + if isinstance(stmt, Update): + update_tables.append(stmt.table.name) + elif isinstance(stmt, Delete): + delete_tables.append(stmt.table.name) + + self.assertIn("user_billing", update_tables) + self.assertIn("ad_attributions", update_tables) + self.assertIn("subscriptions", update_tables) + self.assertIn("payments", update_tables) + self.assertIn("promo_code_activations", update_tables) + self.assertIn("user_payment_methods", update_tables) + self.assertIn("message_logs", update_tables) + self.assertIn("users", update_tables) + self.assertIn("user_payment_methods", delete_tables) + self.assertIn("promo_code_activations", delete_tables) + session.delete.assert_awaited_once_with(source)