db: add composite indexes and merge-user optimizations
This commit is contained in:
+10
-11
@@ -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),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
|
||||
+9
-1
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user