feat: add tariff config and database schema

This commit is contained in:
3252a8
2026-04-28 14:33:38 +03:00
parent aabc0e312d
commit 87b2443a06
10 changed files with 686 additions and 1 deletions
+97
View File
@@ -0,0 +1,97 @@
from typing import Any, Dict, List, Optional
from sqlalchemy import and_, delete, select
from sqlalchemy.ext.asyncio import AsyncSession
from db.models import TariffChange, TrafficTopup, TrafficWarning
async def create_traffic_topup(
session: AsyncSession,
*,
subscription_id: int,
payment_id: Optional[int],
purchased_bytes: int,
kind: str,
) -> TrafficTopup:
record = TrafficTopup(
subscription_id=subscription_id,
payment_id=payment_id,
purchased_bytes=purchased_bytes,
kind=kind,
)
session.add(record)
await session.flush()
await session.refresh(record)
return record
async def create_tariff_change(
session: AsyncSession,
change_data: Dict[str, Any],
) -> TariffChange:
record = TariffChange(**change_data)
session.add(record)
await session.flush()
await session.refresh(record)
return record
async def get_warning(
session: AsyncSession,
*,
subscription_id: int,
period_start_at,
level: int,
traffic_limit_bytes: Optional[int] = None,
) -> Optional[TrafficWarning]:
conditions = [
TrafficWarning.subscription_id == subscription_id,
TrafficWarning.level == level,
]
if period_start_at is None:
conditions.append(TrafficWarning.period_start_at.is_(None))
if traffic_limit_bytes is not None:
conditions.append(TrafficWarning.traffic_limit_bytes == traffic_limit_bytes)
else:
conditions.append(TrafficWarning.period_start_at == period_start_at)
result = await session.execute(select(TrafficWarning).where(and_(*conditions)).limit(1))
return result.scalar_one_or_none()
async def create_warning(
session: AsyncSession,
*,
subscription_id: int,
period_start_at,
level: int,
traffic_limit_bytes: Optional[int],
) -> TrafficWarning:
record = TrafficWarning(
subscription_id=subscription_id,
period_start_at=period_start_at,
level=level,
traffic_limit_bytes=traffic_limit_bytes,
)
session.add(record)
await session.flush()
await session.refresh(record)
return record
async def clear_period_warnings(session: AsyncSession, subscription_id: int) -> int:
result = await session.execute(
delete(TrafficWarning).where(TrafficWarning.subscription_id == subscription_id)
)
return result.rowcount or 0
async def get_tariff_changes_for_subscription(
session: AsyncSession, subscription_id: int
) -> List[TariffChange]:
result = await session.execute(
select(TariffChange)
.where(TariffChange.subscription_id == subscription_id)
.order_by(TariffChange.created_at.desc())
)
return list(result.scalars().all())
+54
View File
@@ -72,6 +72,7 @@ async def init_db(settings: Settings, session_factory: sessionmaker):
async with session_factory() as session:
from .dal.panel_sync_dal import get_panel_sync_status, update_panel_sync_status
from sqlalchemy import text
try:
current_status = await get_panel_sync_status(session)
if current_status is None:
@@ -87,3 +88,56 @@ async def init_db(settings: Settings, session_factory: sessionmaker):
logging.error(
f"Failed to initialize PanelSyncStatus: {e_sync_init}",
exc_info=True)
if settings.tariffs_config:
try:
default_tariff = settings.tariffs_config.default
default_price = default_tariff.period_price(1, "rub") or default_tariff.min_period_price_rub()
await session.execute(
text(
"""
UPDATE subscriptions AS s
SET
tariff_key = COALESCE(s.tariff_key, :tariff_key),
tier_baseline_bytes = COALESCE(s.tier_baseline_bytes, s.traffic_limit_bytes, :baseline),
topup_balance_bytes = COALESCE(s.topup_balance_bytes, 0),
period_start_at = COALESCE(
s.period_start_at,
(
SELECT p.created_at
FROM payments p
WHERE p.user_id = s.user_id
AND p.status = 'succeeded'
ORDER BY p.created_at DESC
LIMIT 1
),
s.start_date,
NOW()
),
effective_monthly_price_rub = COALESCE(
s.effective_monthly_price_rub,
(
SELECT p.amount / GREATEST(COALESCE(p.subscription_duration_months, 1), 1)
FROM payments p
WHERE p.user_id = s.user_id
AND p.status = 'succeeded'
AND COALESCE(p.subscription_duration_months, 0) > 0
ORDER BY p.created_at DESC
LIMIT 1
),
:default_price
)
WHERE s.is_active = TRUE
AND s.tariff_key IS NULL
"""
),
{
"tariff_key": default_tariff.key,
"baseline": default_tariff.monthly_bytes,
"default_price": default_price,
},
)
await session.commit()
except Exception:
await session.rollback()
logging.exception("Failed to backfill existing subscriptions for tariffs config.")
+99
View File
@@ -303,6 +303,100 @@ def _migration_0010_add_email_magic_token_hash(connection: Connection) -> None:
)
def _migration_0011_add_tariffs_schema(connection: Connection) -> None:
inspector = inspect(connection)
sub_columns: Set[str] = {col["name"] for col in inspector.get_columns("subscriptions")}
sub_statements: List[str] = []
if "tariff_key" not in sub_columns:
sub_statements.append("ALTER TABLE subscriptions ADD COLUMN tariff_key VARCHAR")
if "tier_baseline_bytes" not in sub_columns:
sub_statements.append("ALTER TABLE subscriptions ADD COLUMN tier_baseline_bytes BIGINT")
if "topup_balance_bytes" not in sub_columns:
sub_statements.append("ALTER TABLE subscriptions ADD COLUMN topup_balance_bytes BIGINT NOT NULL DEFAULT 0")
if "period_start_at" not in sub_columns:
sub_statements.append("ALTER TABLE subscriptions ADD COLUMN period_start_at TIMESTAMPTZ")
if "is_throttled" not in sub_columns:
sub_statements.append("ALTER TABLE subscriptions ADD COLUMN is_throttled BOOLEAN NOT NULL DEFAULT FALSE")
if "effective_monthly_price_rub" not in sub_columns:
sub_statements.append("ALTER TABLE subscriptions ADD COLUMN effective_monthly_price_rub NUMERIC")
for stmt in sub_statements:
connection.execute(text(stmt))
payment_columns: Set[str] = {col["name"] for col in inspector.get_columns("payments")}
payment_statements: List[str] = []
if "sale_mode" not in payment_columns:
payment_statements.append("ALTER TABLE payments ADD COLUMN sale_mode VARCHAR")
if "tariff_key" not in payment_columns:
payment_statements.append("ALTER TABLE payments ADD COLUMN tariff_key VARCHAR")
if "purchased_gb" not in payment_columns:
payment_statements.append("ALTER TABLE payments ADD COLUMN purchased_gb DOUBLE PRECISION")
for stmt in payment_statements:
connection.execute(text(stmt))
connection.execute(
text(
"""
CREATE TABLE IF NOT EXISTS traffic_topups (
topup_id SERIAL PRIMARY KEY,
subscription_id INTEGER NOT NULL REFERENCES subscriptions(subscription_id),
payment_id INTEGER NULL REFERENCES payments(payment_id),
purchased_bytes BIGINT NOT NULL,
kind VARCHAR NOT NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
)
"""
)
)
connection.execute(
text(
"""
CREATE TABLE IF NOT EXISTS traffic_warnings (
warning_id SERIAL PRIMARY KEY,
subscription_id INTEGER NOT NULL REFERENCES subscriptions(subscription_id),
period_start_at TIMESTAMPTZ NULL,
level INTEGER NOT NULL,
traffic_limit_bytes BIGINT NULL,
sent_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
CONSTRAINT uq_traffic_warning_period_level UNIQUE (subscription_id, period_start_at, level)
)
"""
)
)
connection.execute(
text(
"""
CREATE TABLE IF NOT EXISTS tariff_changes (
change_id SERIAL PRIMARY KEY,
subscription_id INTEGER NOT NULL REFERENCES subscriptions(subscription_id),
from_tariff_key VARCHAR NULL,
to_tariff_key VARCHAR NOT NULL,
mode VARCHAR NOT NULL,
payment_id INTEGER NULL REFERENCES payments(payment_id),
days_before INTEGER NULL,
days_after INTEGER NULL,
converted_bytes BIGINT NULL,
eff_price_before NUMERIC NULL,
eff_price_after NUMERIC NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
)
"""
)
)
for stmt in [
"CREATE INDEX IF NOT EXISTS ix_subscriptions_tariff_key ON subscriptions (tariff_key)",
"CREATE INDEX IF NOT EXISTS ix_subscriptions_is_throttled ON subscriptions (is_throttled)",
"CREATE INDEX IF NOT EXISTS ix_payments_sale_mode ON payments (sale_mode)",
"CREATE INDEX IF NOT EXISTS ix_payments_tariff_key ON payments (tariff_key)",
"CREATE INDEX IF NOT EXISTS ix_traffic_topups_subscription_id ON traffic_topups (subscription_id)",
"CREATE INDEX IF NOT EXISTS ix_traffic_topups_payment_id ON traffic_topups (payment_id)",
"CREATE INDEX IF NOT EXISTS ix_traffic_topups_kind ON traffic_topups (kind)",
"CREATE INDEX IF NOT EXISTS ix_traffic_warnings_subscription_id ON traffic_warnings (subscription_id)",
"CREATE INDEX IF NOT EXISTS ix_tariff_changes_subscription_id ON tariff_changes (subscription_id)",
]:
connection.execute(text(stmt))
def _migration_0009_add_composite_indexes(connection: Connection) -> None:
connection.execute(
text(
@@ -381,6 +475,11 @@ MIGRATIONS: List[Migration] = [
description="Store hashed magic-link tokens for email login deeplinks",
upgrade=_migration_0010_add_email_magic_token_hash,
),
Migration(
id="0011_add_tariffs_schema",
description="Add tariff catalog columns and traffic accounting tables",
upgrade=_migration_0011_add_tariffs_schema,
),
]
+60 -1
View File
@@ -1,4 +1,4 @@
from sqlalchemy import create_engine, Column, Integer, String, Boolean, DateTime, Float, ForeignKey, UniqueConstraint, Text, BigInteger, Index
from sqlalchemy import create_engine, Column, Integer, String, Boolean, DateTime, Float, ForeignKey, UniqueConstraint, Text, BigInteger, Index, Numeric
from sqlalchemy.orm import relationship, DeclarativeBase
from sqlalchemy.ext.asyncio import AsyncAttrs
from sqlalchemy.sql import func
@@ -87,6 +87,12 @@ class Subscription(Base):
provider = Column(String, nullable=True)
skip_notifications = Column(Boolean, default=False)
auto_renew_enabled = Column(Boolean, default=True, index=True)
tariff_key = Column(String, nullable=True, index=True)
tier_baseline_bytes = Column(BigInteger, nullable=True)
topup_balance_bytes = Column(BigInteger, nullable=False, default=0)
period_start_at = Column(DateTime(timezone=True), nullable=True)
is_throttled = Column(Boolean, nullable=False, default=False, index=True)
effective_monthly_price_rub = Column(Numeric, nullable=True)
user = relationship("User", back_populates="subscriptions")
@@ -158,6 +164,9 @@ class Payment(Base):
status = Column(String, nullable=False, index=True)
description = Column(String, nullable=True)
subscription_duration_months = Column(Integer, nullable=True)
sale_mode = Column(String, nullable=True, index=True)
tariff_key = Column(String, nullable=True, index=True)
purchased_gb = Column(Float, nullable=True)
promo_code_id = Column(Integer,
ForeignKey("promo_codes.promo_code_id"),
nullable=True)
@@ -171,6 +180,56 @@ class Payment(Base):
back_populates="payments_where_used")
class TrafficTopup(Base):
__tablename__ = "traffic_topups"
topup_id = Column(Integer, primary_key=True, autoincrement=True)
subscription_id = Column(Integer, ForeignKey("subscriptions.subscription_id"), nullable=False, index=True)
payment_id = Column(Integer, ForeignKey("payments.payment_id"), nullable=True, index=True)
purchased_bytes = Column(BigInteger, nullable=False)
kind = Column(String, nullable=False, index=True)
created_at = Column(DateTime(timezone=True), server_default=func.now())
subscription = relationship("Subscription")
payment = relationship("Payment")
class TrafficWarning(Base):
__tablename__ = "traffic_warnings"
__table_args__ = (
UniqueConstraint("subscription_id", "period_start_at", "level", name="uq_traffic_warning_period_level"),
)
warning_id = Column(Integer, primary_key=True, autoincrement=True)
subscription_id = Column(Integer, ForeignKey("subscriptions.subscription_id"), nullable=False, index=True)
period_start_at = Column(DateTime(timezone=True), nullable=True)
level = Column(Integer, nullable=False)
traffic_limit_bytes = Column(BigInteger, nullable=True)
sent_at = Column(DateTime(timezone=True), server_default=func.now())
subscription = relationship("Subscription")
class TariffChange(Base):
__tablename__ = "tariff_changes"
change_id = Column(Integer, primary_key=True, autoincrement=True)
subscription_id = Column(Integer, ForeignKey("subscriptions.subscription_id"), nullable=False, index=True)
from_tariff_key = Column(String, nullable=True)
to_tariff_key = Column(String, nullable=False)
mode = Column(String, nullable=False, index=True)
payment_id = Column(Integer, ForeignKey("payments.payment_id"), nullable=True, index=True)
days_before = Column(Integer, nullable=True)
days_after = Column(Integer, nullable=True)
converted_bytes = Column(BigInteger, nullable=True)
eff_price_before = Column(Numeric, nullable=True)
eff_price_after = Column(Numeric, nullable=True)
created_at = Column(DateTime(timezone=True), server_default=func.now())
subscription = relationship("Subscription")
payment = relationship("Payment")
class UserBilling(Base):
__tablename__ = "user_billing"