Migrating to Postgres from sqlite3
This commit is contained in:
machka
@@ -0,0 +1,76 @@
|
||||
import logging
|
||||
from typing import Optional, List
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
from sqlalchemy import func, or_
|
||||
|
||||
from ..models import MessageLog, User
|
||||
|
||||
|
||||
async def create_message_log(session: AsyncSession,
|
||||
log_data: dict) -> Optional[MessageLog]:
|
||||
|
||||
try:
|
||||
log_entry = await create_message_log_no_commit(session, log_data)
|
||||
await session.commit()
|
||||
await session.refresh(log_entry)
|
||||
return log_entry
|
||||
except Exception as e:
|
||||
await session.rollback()
|
||||
logging.error(f"Failed to create and commit message log: {e}",
|
||||
exc_info=True)
|
||||
return None
|
||||
|
||||
|
||||
async def get_all_message_logs(session: AsyncSession, limit: int,
|
||||
offset: int) -> List[MessageLog]:
|
||||
stmt = select(MessageLog).order_by(
|
||||
MessageLog.timestamp.desc()).limit(limit).offset(offset)
|
||||
result = await session.execute(stmt)
|
||||
return result.scalars().all()
|
||||
|
||||
|
||||
async def count_all_message_logs(session: AsyncSession) -> int:
|
||||
stmt = select(func.count()).select_from(MessageLog)
|
||||
result = await session.execute(stmt)
|
||||
return result.scalar_one()
|
||||
|
||||
|
||||
async def get_user_message_logs(session: AsyncSession, user_id_to_search: int,
|
||||
limit: int, offset: int) -> List[MessageLog]:
|
||||
stmt = (select(MessageLog).where(
|
||||
or_(MessageLog.user_id == user_id_to_search,
|
||||
MessageLog.target_user_id == user_id_to_search)).order_by(
|
||||
MessageLog.timestamp.desc()).limit(limit).offset(offset))
|
||||
result = await session.execute(stmt)
|
||||
return result.scalars().all()
|
||||
|
||||
|
||||
async def count_user_message_logs(session: AsyncSession,
|
||||
user_id_to_search: int) -> int:
|
||||
stmt = (select(func.count()).select_from(MessageLog).where(
|
||||
or_(MessageLog.user_id == user_id_to_search,
|
||||
MessageLog.target_user_id == user_id_to_search)))
|
||||
result = await session.execute(stmt)
|
||||
return result.scalar_one()
|
||||
|
||||
|
||||
async def create_message_log_no_commit(session: AsyncSession,
|
||||
log_data: dict) -> MessageLog:
|
||||
|
||||
if log_data.get("target_user_id"):
|
||||
from .user_dal import get_user_by_id
|
||||
target_user = await get_user_by_id(session, log_data["target_user_id"])
|
||||
if not target_user:
|
||||
logging.warning(
|
||||
f"Target user {log_data['target_user_id']} not found for message log. Setting to NULL."
|
||||
)
|
||||
log_data["target_user_id"] = None
|
||||
|
||||
new_log = MessageLog(**log_data)
|
||||
session.add(new_log)
|
||||
|
||||
logging.debug(
|
||||
f"Message log added to session: user {log_data.get('user_id')}, event {log_data.get('event_type')}"
|
||||
)
|
||||
return new_log
|
||||
@@ -0,0 +1,50 @@
|
||||
import logging
|
||||
from typing import Optional
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
from sqlalchemy import update
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from db.models import PanelSyncStatus
|
||||
|
||||
SINGLETON_ID = 1
|
||||
|
||||
|
||||
async def get_panel_sync_status(
|
||||
session: AsyncSession) -> Optional[PanelSyncStatus]:
|
||||
return await session.get(PanelSyncStatus, SINGLETON_ID)
|
||||
|
||||
|
||||
async def update_panel_sync_status(
|
||||
session: AsyncSession,
|
||||
status: str,
|
||||
details: str,
|
||||
users_processed: int = 0,
|
||||
subs_synced: int = 0,
|
||||
last_sync_time: Optional[datetime] = None) -> PanelSyncStatus:
|
||||
if last_sync_time is None:
|
||||
last_sync_time = datetime.now(timezone.utc)
|
||||
|
||||
sync_record = await get_panel_sync_status(session)
|
||||
if sync_record:
|
||||
sync_record.last_sync_time = last_sync_time
|
||||
sync_record.status = status
|
||||
sync_record.details = details
|
||||
sync_record.users_processed_from_panel = users_processed
|
||||
sync_record.subscriptions_synced = subs_synced
|
||||
else:
|
||||
sync_record = PanelSyncStatus(
|
||||
id=SINGLETON_ID,
|
||||
last_sync_time=last_sync_time,
|
||||
status=status,
|
||||
details=details,
|
||||
users_processed_from_panel=users_processed,
|
||||
subscriptions_synced=subs_synced)
|
||||
session.add(sync_record)
|
||||
|
||||
await session.flush()
|
||||
await session.refresh(sync_record)
|
||||
logging.info(
|
||||
f"Panel sync status updated: {status}, Users: {users_processed}, Subs: {subs_synced}"
|
||||
)
|
||||
return sync_record
|
||||
@@ -0,0 +1,115 @@
|
||||
import logging
|
||||
from typing import Optional, List, Dict, Any
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
from sqlalchemy import update, func
|
||||
from sqlalchemy.orm import selectinload
|
||||
|
||||
from db.models import Payment, User
|
||||
|
||||
|
||||
async def create_payment_record(session: AsyncSession,
|
||||
payment_data: Dict[str, Any]) -> Payment:
|
||||
|
||||
from .user_dal import get_user_by_id
|
||||
user = await get_user_by_id(session, payment_data["user_id"])
|
||||
if not user:
|
||||
|
||||
raise ValueError(
|
||||
f"User with id {payment_data['user_id']} not found for creating payment."
|
||||
)
|
||||
|
||||
if payment_data.get("promo_code_id"):
|
||||
from .promo_code_dal import get_promo_code_by_id
|
||||
promo = await get_promo_code_by_id(session,
|
||||
payment_data["promo_code_id"])
|
||||
if not promo:
|
||||
raise ValueError(
|
||||
f"Promo code with id {payment_data['promo_code_id']} not found."
|
||||
)
|
||||
|
||||
new_payment = Payment(**payment_data)
|
||||
session.add(new_payment)
|
||||
await session.flush()
|
||||
await session.refresh(new_payment)
|
||||
logging.info(
|
||||
f"Payment record {new_payment.payment_id} created for user {new_payment.user_id}"
|
||||
)
|
||||
return new_payment
|
||||
|
||||
|
||||
async def get_payment_by_yookassa_id(
|
||||
session: AsyncSession, yookassa_payment_id: str) -> Optional[Payment]:
|
||||
stmt = select(Payment).where(
|
||||
Payment.yookassa_payment_id == yookassa_payment_id)
|
||||
result = await session.execute(stmt)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
|
||||
async def get_payment_by_db_id(session: AsyncSession,
|
||||
payment_db_id: int) -> Optional[Payment]:
|
||||
|
||||
stmt = select(Payment).where(Payment.payment_id == payment_db_id).options(
|
||||
selectinload(Payment.user), selectinload(Payment.promo_code_used))
|
||||
result = await session.execute(stmt)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
|
||||
async def get_payment_by_db_id_with_promo(
|
||||
session: AsyncSession, payment_db_id: int) -> Optional[Payment]:
|
||||
|
||||
stmt = select(Payment).where(Payment.payment_id == payment_db_id).options(
|
||||
selectinload(Payment.promo_code_used))
|
||||
result = await session.execute(stmt)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
|
||||
async def update_payment_status_by_db_id(
|
||||
session: AsyncSession,
|
||||
payment_db_id: int,
|
||||
new_status: str,
|
||||
yk_payment_id: Optional[str] = None) -> Optional[Payment]:
|
||||
payment = await get_payment_by_db_id(session, payment_db_id)
|
||||
if payment:
|
||||
payment.status = new_status
|
||||
payment.updated_at = func.now()
|
||||
if yk_payment_id and payment.yookassa_payment_id is None:
|
||||
payment.yookassa_payment_id = yk_payment_id
|
||||
await session.flush()
|
||||
await session.refresh(payment)
|
||||
logging.info(
|
||||
f"Payment record {payment.payment_id} status updated to {new_status}."
|
||||
)
|
||||
else:
|
||||
logging.warning(
|
||||
f"Payment record with DB ID {payment_db_id} not found for status update."
|
||||
)
|
||||
return payment
|
||||
|
||||
|
||||
async def update_payment_status_by_yk_id(session: AsyncSession,
|
||||
yookassa_payment_id: str,
|
||||
new_status: str) -> Optional[Payment]:
|
||||
payment = await get_payment_by_yookassa_id(session, yookassa_payment_id)
|
||||
if payment:
|
||||
payment.status = new_status
|
||||
payment.updated_at = func.now()
|
||||
await session.flush()
|
||||
await session.refresh(payment)
|
||||
logging.info(
|
||||
f"Payment record with YK ID {yookassa_payment_id} status updated to {new_status}."
|
||||
)
|
||||
else:
|
||||
logging.warning(
|
||||
f"Payment record with YK ID {yookassa_payment_id} not found for status update."
|
||||
)
|
||||
return payment
|
||||
|
||||
|
||||
async def get_recent_payment_logs_with_user(session: AsyncSession,
|
||||
limit: int = 20,
|
||||
offset: int = 0) -> List[Payment]:
|
||||
stmt = (select(Payment).options(selectinload(Payment.user)).order_by(
|
||||
Payment.created_at.desc()).limit(limit).offset(offset))
|
||||
result = await session.execute(stmt)
|
||||
return result.scalars().all()
|
||||
@@ -0,0 +1,126 @@
|
||||
import logging
|
||||
from typing import Optional, List, Dict, Any
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
from sqlalchemy import update, func, and_, or_
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from db.models import PromoCode, PromoCodeActivation, User, Payment
|
||||
|
||||
|
||||
async def create_promo_code(session: AsyncSession,
|
||||
promo_data: Dict[str, Any]) -> PromoCode:
|
||||
|
||||
new_promo = PromoCode(**promo_data)
|
||||
session.add(new_promo)
|
||||
await session.flush()
|
||||
await session.refresh(new_promo)
|
||||
logging.info(
|
||||
f"Promo code '{new_promo.code}' created with ID {new_promo.promo_code_id}"
|
||||
)
|
||||
return new_promo
|
||||
|
||||
|
||||
async def get_promo_code_by_id(session: AsyncSession,
|
||||
promo_code_id: int) -> Optional[PromoCode]:
|
||||
return await session.get(PromoCode, promo_code_id)
|
||||
|
||||
|
||||
async def get_active_promo_code_by_code_str(
|
||||
session: AsyncSession, code_str: str) -> Optional[PromoCode]:
|
||||
stmt = select(PromoCode).where(
|
||||
PromoCode.code == code_str.upper(), PromoCode.is_active == True,
|
||||
PromoCode.current_activations < PromoCode.max_activations,
|
||||
or_(PromoCode.valid_until == None, PromoCode.valid_until
|
||||
> datetime.now(timezone.utc)))
|
||||
result = await session.execute(stmt)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
|
||||
async def get_all_active_promo_codes(session: AsyncSession,
|
||||
limit: int = 20,
|
||||
offset: int = 0) -> List[PromoCode]:
|
||||
stmt = (select(PromoCode).where(
|
||||
PromoCode.is_active == True,
|
||||
or_(PromoCode.valid_until == None, PromoCode.valid_until
|
||||
> datetime.now(timezone.utc))).order_by(
|
||||
PromoCode.created_at.desc()).limit(limit).offset(offset))
|
||||
result = await session.execute(stmt)
|
||||
return result.scalars().all()
|
||||
|
||||
|
||||
async def increment_promo_code_usage(
|
||||
session: AsyncSession, promo_code_id: int) -> Optional[PromoCode]:
|
||||
promo = await get_promo_code_by_id(session, promo_code_id)
|
||||
if promo:
|
||||
if promo.current_activations < promo.max_activations:
|
||||
promo.current_activations += 1
|
||||
await session.flush()
|
||||
await session.refresh(promo)
|
||||
return promo
|
||||
else:
|
||||
logging.warning(
|
||||
f"Promo code {promo.code} (ID: {promo_code_id}) already reached max activations."
|
||||
)
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
async def get_user_activation_for_promo(
|
||||
session: AsyncSession, promo_code_id: int,
|
||||
user_id: int) -> Optional[PromoCodeActivation]:
|
||||
|
||||
stmt = select(PromoCodeActivation).where(
|
||||
PromoCodeActivation.promo_code_id == promo_code_id,
|
||||
PromoCodeActivation.user_id == user_id).limit(1)
|
||||
result = await session.execute(stmt)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
|
||||
async def record_promo_activation(
|
||||
session: AsyncSession,
|
||||
promo_code_id: int,
|
||||
user_id: int,
|
||||
payment_id: Optional[int] = None) -> Optional[PromoCodeActivation]:
|
||||
|
||||
existing_activation = await get_user_activation_for_promo(
|
||||
session, promo_code_id, user_id)
|
||||
if existing_activation:
|
||||
logging.info(
|
||||
f"User {user_id} has already activated promo code {promo_code_id}. Activation ID: {existing_activation.activation_id}"
|
||||
)
|
||||
return existing_activation
|
||||
|
||||
from .user_dal import get_user_by_id
|
||||
user = await get_user_by_id(session, user_id)
|
||||
promo = await get_promo_code_by_id(session, promo_code_id)
|
||||
if not user or not promo:
|
||||
logging.error(
|
||||
f"Cannot record promo activation: User {user_id} or Promo {promo_code_id} not found."
|
||||
)
|
||||
return None
|
||||
|
||||
if payment_id:
|
||||
from .payment_dal import get_payment_by_db_id
|
||||
payment = await get_payment_by_db_id(session, payment_id)
|
||||
if not payment:
|
||||
logging.error(
|
||||
f"Cannot record promo activation: Payment {payment_id} not found."
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
activation_data = {
|
||||
"promo_code_id": promo_code_id,
|
||||
"user_id": user_id,
|
||||
"payment_id": payment_id,
|
||||
"activated_at": datetime.now(timezone.utc)
|
||||
}
|
||||
new_activation = PromoCodeActivation(**activation_data)
|
||||
session.add(new_activation)
|
||||
await session.flush()
|
||||
await session.refresh(new_activation)
|
||||
logging.info(
|
||||
f"Promo code {promo_code_id} activated by user {user_id}. Activation ID: {new_activation.activation_id}"
|
||||
)
|
||||
return new_activation
|
||||
@@ -0,0 +1,205 @@
|
||||
import logging
|
||||
from typing import Optional, List, Dict, Any
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
from sqlalchemy import update, delete, func, and_, or_
|
||||
from sqlalchemy.orm import selectinload
|
||||
from datetime import datetime, timezone, timedelta
|
||||
|
||||
from db.models import Subscription, User
|
||||
|
||||
|
||||
async def get_active_subscription_by_user_id(
|
||||
session: AsyncSession,
|
||||
user_id: int,
|
||||
panel_user_uuid: Optional[str] = None) -> Optional[Subscription]:
|
||||
stmt = select(Subscription).where(
|
||||
Subscription.user_id == user_id, Subscription.is_active == True,
|
||||
Subscription.end_date > datetime.now(timezone.utc))
|
||||
if panel_user_uuid:
|
||||
stmt = stmt.where(Subscription.panel_user_uuid == panel_user_uuid)
|
||||
stmt = stmt.order_by(Subscription.end_date.desc())
|
||||
result = await session.execute(stmt)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
|
||||
async def get_subscription_by_panel_subscription_uuid(
|
||||
session: AsyncSession, panel_sub_uuid: str) -> Optional[Subscription]:
|
||||
stmt = select(Subscription).where(
|
||||
Subscription.panel_subscription_uuid == panel_sub_uuid)
|
||||
result = await session.execute(stmt)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
|
||||
async def create_subscription(session: AsyncSession,
|
||||
sub_data: Dict[str, Any]) -> Subscription:
|
||||
from .user_dal import get_user_by_id
|
||||
|
||||
if "user_id" not in sub_data or sub_data["user_id"] is None:
|
||||
raise ValueError(
|
||||
"user_id is required to create a subscription directly.")
|
||||
user = await get_user_by_id(session, sub_data["user_id"])
|
||||
if not user:
|
||||
raise ValueError(
|
||||
f"User with id {sub_data['user_id']} not found for creating subscription."
|
||||
)
|
||||
|
||||
new_sub = Subscription(**sub_data)
|
||||
session.add(new_sub)
|
||||
await session.flush()
|
||||
await session.refresh(new_sub)
|
||||
logging.info(
|
||||
f"Subscription {new_sub.subscription_id} created for user {new_sub.user_id}"
|
||||
)
|
||||
return new_sub
|
||||
|
||||
|
||||
async def update_subscription(
|
||||
session: AsyncSession, subscription_id: int,
|
||||
update_data: Dict[str, Any]) -> Optional[Subscription]:
|
||||
sub = await session.get(Subscription, subscription_id)
|
||||
if sub:
|
||||
for key, value in update_data.items():
|
||||
setattr(sub, key, value)
|
||||
await session.flush()
|
||||
await session.refresh(sub)
|
||||
return sub
|
||||
|
||||
|
||||
async def upsert_subscription(session: AsyncSession,
|
||||
sub_payload: Dict[str, Any]) -> Subscription:
|
||||
panel_sub_uuid = sub_payload.get("panel_subscription_uuid")
|
||||
if not panel_sub_uuid:
|
||||
raise ValueError("panel_subscription_uuid is required for upsert.")
|
||||
|
||||
existing_sub = await get_subscription_by_panel_subscription_uuid(
|
||||
session, panel_sub_uuid)
|
||||
|
||||
if existing_sub:
|
||||
logging.info(
|
||||
f"Updating existing subscription {existing_sub.subscription_id} by panel_sub_uuid {panel_sub_uuid}"
|
||||
)
|
||||
for key, value in sub_payload.items():
|
||||
if hasattr(existing_sub, key):
|
||||
setattr(existing_sub, key, value)
|
||||
await session.flush()
|
||||
await session.refresh(existing_sub)
|
||||
return existing_sub
|
||||
else:
|
||||
logging.info(
|
||||
f"Creating new subscription with panel_sub_uuid {panel_sub_uuid}")
|
||||
|
||||
if sub_payload.get(
|
||||
"user_id") is None and "panel_user_uuid" not in sub_payload:
|
||||
raise ValueError(
|
||||
"For a new subscription without user_id, panel_user_uuid is required."
|
||||
)
|
||||
if "end_date" not in sub_payload:
|
||||
raise ValueError("Missing 'end_date' for new subscription.")
|
||||
if sub_payload.get("user_id") is not None:
|
||||
from .user_dal import get_user_by_id
|
||||
user = await get_user_by_id(session, sub_payload["user_id"])
|
||||
if not user:
|
||||
raise ValueError(
|
||||
f"User {sub_payload['user_id']} not found for new subscription with panel_uuid {panel_sub_uuid}."
|
||||
)
|
||||
|
||||
new_sub = Subscription(**sub_payload)
|
||||
session.add(new_sub)
|
||||
await session.flush()
|
||||
await session.refresh(new_sub)
|
||||
return new_sub
|
||||
|
||||
|
||||
async def deactivate_other_active_subscriptions(
|
||||
session: AsyncSession, panel_user_uuid: str,
|
||||
current_panel_subscription_uuid: Optional[str]):
|
||||
stmt = (update(Subscription).where(
|
||||
Subscription.panel_user_uuid == panel_user_uuid,
|
||||
Subscription.is_active == True,
|
||||
).values(is_active=False, status_from_panel="INACTIVE_BY_BOT_SYNC"))
|
||||
if current_panel_subscription_uuid:
|
||||
stmt = stmt.where(Subscription.panel_subscription_uuid !=
|
||||
current_panel_subscription_uuid)
|
||||
|
||||
result = await session.execute(stmt)
|
||||
if result.rowcount > 0:
|
||||
logging.info(
|
||||
f"Deactivated {result.rowcount} other active subscriptions for panel_user_uuid {panel_user_uuid}."
|
||||
)
|
||||
|
||||
|
||||
async def update_subscription_end_date(
|
||||
session: AsyncSession, subscription_id: int,
|
||||
new_end_date: datetime) -> Optional[Subscription]:
|
||||
|
||||
return await update_subscription(
|
||||
session, subscription_id, {
|
||||
"end_date": new_end_date,
|
||||
"last_notification_sent": None,
|
||||
"is_active": True,
|
||||
"status_from_panel": "ACTIVE_EXTENDED_BY_BOT"
|
||||
})
|
||||
|
||||
|
||||
async def has_any_subscription_for_user(session: AsyncSession,
|
||||
user_id: int) -> bool:
|
||||
stmt = select(Subscription.subscription_id).where(
|
||||
Subscription.user_id == user_id).limit(1)
|
||||
result = await session.execute(stmt)
|
||||
return result.scalar_one_or_none() is not None
|
||||
|
||||
|
||||
async def get_subscriptions_near_expiration(
|
||||
session: AsyncSession, days_threshold: int) -> List[Subscription]:
|
||||
now_utc = datetime.now(timezone.utc)
|
||||
threshold_date = now_utc + timedelta(days=days_threshold)
|
||||
|
||||
stmt = (select(Subscription).join(Subscription.user).where(
|
||||
Subscription.is_active == True, Subscription.end_date > now_utc,
|
||||
Subscription.end_date <= threshold_date,
|
||||
or_(
|
||||
Subscription.last_notification_sent == None,
|
||||
func.date(Subscription.last_notification_sent)
|
||||
< func.date(now_utc))).order_by(
|
||||
Subscription.end_date.asc()).options(
|
||||
selectinload(Subscription.user)))
|
||||
result = await session.execute(stmt)
|
||||
return result.scalars().all()
|
||||
|
||||
|
||||
async def update_subscription_notification_time(
|
||||
session: AsyncSession, subscription_id: int,
|
||||
notification_time: datetime) -> Optional[Subscription]:
|
||||
return await update_subscription(
|
||||
session, subscription_id,
|
||||
{"last_notification_sent": notification_time})
|
||||
|
||||
|
||||
async def get_user_active_subscription_end_date_str(
|
||||
session: AsyncSession, user_id: int) -> Optional[str]:
|
||||
stmt = (select(Subscription.end_date).where(
|
||||
Subscription.user_id == user_id, Subscription.is_active == True,
|
||||
Subscription.end_date > datetime.now(timezone.utc)).order_by(
|
||||
Subscription.end_date.desc()).limit(1))
|
||||
result = await session.execute(stmt)
|
||||
end_date_obj = result.scalar_one_or_none()
|
||||
return end_date_obj.strftime('%Y-%m-%d') if end_date_obj else None
|
||||
|
||||
|
||||
async def find_subscription_for_notification_update(
|
||||
session: AsyncSession, user_id: int,
|
||||
subscription_end_date_to_match: datetime) -> Optional[Subscription]:
|
||||
|
||||
if subscription_end_date_to_match.tzinfo is None:
|
||||
subscription_end_date_to_match = subscription_end_date_to_match.replace(
|
||||
tzinfo=timezone.utc)
|
||||
|
||||
stmt = select(Subscription).where(
|
||||
Subscription.user_id == user_id, Subscription.is_active == True,
|
||||
Subscription.end_date
|
||||
>= subscription_end_date_to_match - timedelta(seconds=1),
|
||||
Subscription.end_date
|
||||
<= subscription_end_date_to_match + timedelta(seconds=1)).limit(1)
|
||||
result = await session.execute(stmt)
|
||||
return result.scalar_one_or_none()
|
||||
@@ -0,0 +1,124 @@
|
||||
import logging
|
||||
from typing import Optional, List, Dict, Any, Tuple
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
from sqlalchemy.orm import selectinload
|
||||
from sqlalchemy import update, delete, func, and_
|
||||
from datetime import datetime
|
||||
|
||||
from ..models import User, Subscription
|
||||
|
||||
|
||||
async def get_user_by_id(session: AsyncSession,
|
||||
user_id: int) -> Optional[User]:
|
||||
stmt = select(User).where(User.user_id == user_id)
|
||||
result = await session.execute(stmt)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
|
||||
async def get_user_by_username(session: AsyncSession,
|
||||
username: str) -> Optional[User]:
|
||||
clean_username = username.lstrip('@').lower()
|
||||
stmt = select(User).where(func.lower(User.username) == clean_username)
|
||||
result = await session.execute(stmt)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
|
||||
async def get_user_by_panel_uuid(session: AsyncSession,
|
||||
panel_uuid: str) -> Optional[User]:
|
||||
stmt = select(User).where(User.panel_user_uuid == panel_uuid)
|
||||
result = await session.execute(stmt)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
|
||||
async def create_user(session: AsyncSession, user_data: Dict[str,
|
||||
Any]) -> User:
|
||||
|
||||
if 'registration_date' not in user_data:
|
||||
user_data['registration_date'] = datetime.now()
|
||||
|
||||
new_user = User(**user_data)
|
||||
session.add(new_user)
|
||||
await session.flush()
|
||||
await session.refresh(new_user)
|
||||
logging.info(
|
||||
f"New user {new_user.user_id} created in DAL. Referred by: {new_user.referred_by_id or 'N/A'}."
|
||||
)
|
||||
return new_user
|
||||
|
||||
|
||||
async def update_user(session: AsyncSession, user_id: int,
|
||||
update_data: Dict[str, Any]) -> Optional[User]:
|
||||
user = await get_user_by_id(session, user_id)
|
||||
if user:
|
||||
for key, value in update_data.items():
|
||||
setattr(user, key, value)
|
||||
await session.flush()
|
||||
await session.refresh(user)
|
||||
return user
|
||||
|
||||
|
||||
async def update_user_language(session: AsyncSession, user_id: int,
|
||||
lang_code: str) -> bool:
|
||||
stmt = update(User).where(User.user_id == user_id).values(
|
||||
language_code=lang_code)
|
||||
result = await session.execute(stmt)
|
||||
return result.rowcount > 0
|
||||
|
||||
|
||||
async def set_user_ban_status(session: AsyncSession, user_id: int,
|
||||
is_banned: bool) -> bool:
|
||||
user = await get_user_by_id(session, user_id)
|
||||
if user:
|
||||
user.is_banned = is_banned
|
||||
await session.flush()
|
||||
await session.refresh(user)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
async def get_banned_users_paginated(session: AsyncSession, limit: int,
|
||||
offset: int) -> Tuple[List[User], int]:
|
||||
stmt_users = select(User).where(User.is_banned == True).order_by(
|
||||
User.registration_date.desc()).limit(limit).offset(offset)
|
||||
result_users = await session.execute(stmt_users)
|
||||
users_list = result_users.scalars().all()
|
||||
|
||||
stmt_count = select(
|
||||
func.count()).select_from(User).where(User.is_banned == True)
|
||||
result_count = await session.execute(stmt_count)
|
||||
total_banned = result_count.scalar_one()
|
||||
|
||||
return users_list, total_banned
|
||||
|
||||
|
||||
async def get_all_active_user_ids_for_broadcast(
|
||||
session: AsyncSession) -> List[int]:
|
||||
stmt = select(User.user_id).where(User.is_banned == False)
|
||||
result = await session.execute(stmt)
|
||||
return result.scalars().all()
|
||||
|
||||
|
||||
async def get_user_count_stats_dal(session: AsyncSession) -> Dict[str, int]:
|
||||
total_users_stmt = select(func.count(User.user_id)).select_from(User)
|
||||
banned_users_stmt = select(func.count(
|
||||
User.user_id)).select_from(User).where(User.is_banned == True)
|
||||
|
||||
active_subs_stmt = (select(func.count(
|
||||
func.distinct(Subscription.user_id))).join(
|
||||
User, Subscription.user_id == User.user_id).where(
|
||||
Subscription.is_active == True).where(
|
||||
Subscription.end_date > datetime.now()))
|
||||
|
||||
total_users = (await
|
||||
session.execute(total_users_stmt)).scalar_one_or_none() or 0
|
||||
banned_users = (
|
||||
await session.execute(banned_users_stmt)).scalar_one_or_none() or 0
|
||||
active_subs_users = (
|
||||
await session.execute(active_subs_stmt)).scalar_one_or_none() or 0
|
||||
|
||||
return {
|
||||
"total_users": total_users,
|
||||
"banned_users": banned_users,
|
||||
"users_with_active_subscriptions": active_subs_users,
|
||||
}
|
||||
-1081
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,85 @@
|
||||
import logging
|
||||
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from config.settings import Settings
|
||||
from .models import Base
|
||||
|
||||
async_engine = None
|
||||
|
||||
|
||||
def init_db_connection(settings: Settings) -> sessionmaker:
|
||||
global async_engine
|
||||
|
||||
if async_engine is None:
|
||||
logging.info(
|
||||
f"Attempting to create SQLAlchemy engine with URL: {settings.DATABASE_URL}"
|
||||
)
|
||||
async_engine = create_async_engine(
|
||||
settings.DATABASE_URL,
|
||||
echo=False,
|
||||
pool_pre_ping=True,
|
||||
)
|
||||
|
||||
local_async_session_factory = async_sessionmaker(
|
||||
bind=async_engine,
|
||||
class_=AsyncSession,
|
||||
expire_on_commit=False,
|
||||
autocommit=False,
|
||||
autoflush=False,
|
||||
)
|
||||
logging.info(
|
||||
f"SQLAlchemy Async Engine and SessionFactory configured for PostgreSQL."
|
||||
)
|
||||
return local_async_session_factory
|
||||
|
||||
|
||||
async def get_async_session(session_factory: sessionmaker) -> AsyncSession:
|
||||
|
||||
if session_factory is None:
|
||||
raise RuntimeError(
|
||||
"AsyncSessionFactory is not provided or initialized.")
|
||||
|
||||
async_session = session_factory()
|
||||
try:
|
||||
yield async_session
|
||||
finally:
|
||||
await async_session.close()
|
||||
|
||||
|
||||
async def init_db(settings: Settings, session_factory: sessionmaker):
|
||||
|
||||
global async_engine
|
||||
if async_engine is None:
|
||||
|
||||
logging.warning(
|
||||
"init_db: async_engine was None, re-initializing via init_db_connection."
|
||||
)
|
||||
|
||||
raise RuntimeError(
|
||||
"async_engine is not initialized. Call init_db_connection and get session_factory first."
|
||||
)
|
||||
|
||||
async with async_engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
logging.info(
|
||||
"PostgreSQL database initialized/checked successfully using SQLAlchemy."
|
||||
)
|
||||
|
||||
async with session_factory() as session:
|
||||
from .dal.panel_sync_dal import get_panel_sync_status, update_panel_sync_status
|
||||
try:
|
||||
current_status = await get_panel_sync_status(session)
|
||||
if current_status is None:
|
||||
logging.info("Initializing panel_sync_status record.")
|
||||
await update_panel_sync_status(session,
|
||||
status="never_run",
|
||||
details="System initialized",
|
||||
users_processed=0,
|
||||
subs_synced=0)
|
||||
await session.commit()
|
||||
except Exception as e_sync_init:
|
||||
await session.rollback()
|
||||
logging.error(
|
||||
f"Failed to initialize PanelSyncStatus: {e_sync_init}",
|
||||
exc_info=True)
|
||||
+193
@@ -0,0 +1,193 @@
|
||||
from sqlalchemy import create_engine, Column, Integer, String, Boolean, DateTime, Float, ForeignKey, UniqueConstraint, Text, BigInteger
|
||||
from sqlalchemy.orm import relationship, DeclarativeBase
|
||||
from sqlalchemy.ext.asyncio import AsyncAttrs
|
||||
from sqlalchemy.sql import func
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class Base(AsyncAttrs, DeclarativeBase):
|
||||
pass
|
||||
|
||||
|
||||
class User(Base):
|
||||
__tablename__ = "users"
|
||||
|
||||
user_id = Column(BigInteger, primary_key=True, index=True)
|
||||
username = Column(String, nullable=True, index=True)
|
||||
first_name = Column(String, nullable=True)
|
||||
last_name = Column(String, nullable=True)
|
||||
language_code = Column(String, default="ru")
|
||||
registration_date = Column(DateTime(timezone=True),
|
||||
server_default=func.now())
|
||||
is_banned = Column(Boolean, default=False)
|
||||
panel_user_uuid = Column(String, nullable=True, unique=True, index=True)
|
||||
referred_by_id = Column(BigInteger,
|
||||
ForeignKey("users.user_id"),
|
||||
nullable=True)
|
||||
|
||||
referrer = relationship("User", remote_side=[user_id], backref="referrals")
|
||||
subscriptions = relationship("Subscription",
|
||||
back_populates="user",
|
||||
cascade="all, delete-orphan")
|
||||
payments = relationship("Payment",
|
||||
back_populates="user",
|
||||
cascade="all, delete-orphan")
|
||||
promo_code_activations = relationship("PromoCodeActivation",
|
||||
back_populates="user",
|
||||
cascade="all, delete-orphan")
|
||||
message_logs_authored = relationship("MessageLog",
|
||||
foreign_keys="MessageLog.user_id",
|
||||
back_populates="author_user",
|
||||
cascade="all, delete-orphan")
|
||||
message_logs_targeted = relationship(
|
||||
"MessageLog",
|
||||
foreign_keys="MessageLog.target_user_id",
|
||||
back_populates="target_user",
|
||||
cascade="all, delete-orphan")
|
||||
|
||||
def __repr__(self):
|
||||
return f"<User(user_id={self.user_id}, username='{self.username}')>"
|
||||
|
||||
|
||||
class Subscription(Base):
|
||||
__tablename__ = "subscriptions"
|
||||
|
||||
subscription_id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
user_id = Column(BigInteger,
|
||||
ForeignKey("users.user_id"),
|
||||
nullable=False,
|
||||
index=True)
|
||||
panel_user_uuid = Column(String, nullable=False, index=True)
|
||||
panel_subscription_uuid = Column(String,
|
||||
unique=True,
|
||||
index=True,
|
||||
nullable=True)
|
||||
start_date = Column(DateTime(timezone=True), nullable=True)
|
||||
end_date = Column(DateTime(timezone=True), nullable=False, index=True)
|
||||
duration_months = Column(Integer, nullable=True)
|
||||
is_active = Column(Boolean, default=True, index=True)
|
||||
status_from_panel = Column(String, nullable=True)
|
||||
traffic_limit_bytes = Column(BigInteger, nullable=True)
|
||||
traffic_used_bytes = Column(BigInteger, nullable=True)
|
||||
last_notification_sent = Column(DateTime(timezone=True), nullable=True)
|
||||
|
||||
user = relationship("User", back_populates="subscriptions")
|
||||
|
||||
def __repr__(self):
|
||||
return f"<Subscription(id={self.subscription_id}, user_id={self.user_id}, panel_uuid='{self.panel_user_uuid}', ends='{self.end_date}')>"
|
||||
|
||||
|
||||
class Payment(Base):
|
||||
__tablename__ = "payments"
|
||||
|
||||
payment_id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
user_id = Column(BigInteger,
|
||||
ForeignKey("users.user_id"),
|
||||
nullable=False,
|
||||
index=True)
|
||||
yookassa_payment_id = Column(String,
|
||||
unique=True,
|
||||
index=True,
|
||||
nullable=True)
|
||||
idempotence_key = Column(String, unique=True, nullable=True)
|
||||
amount = Column(Float, nullable=False)
|
||||
currency = Column(String, nullable=False)
|
||||
status = Column(String, nullable=False, index=True)
|
||||
description = Column(String, nullable=True)
|
||||
subscription_duration_months = Column(Integer, nullable=True)
|
||||
promo_code_id = Column(Integer,
|
||||
ForeignKey("promo_codes.promo_code_id"),
|
||||
nullable=True)
|
||||
created_at = Column(DateTime(timezone=True), server_default=func.now())
|
||||
updated_at = Column(DateTime(timezone=True),
|
||||
onupdate=func.now(),
|
||||
nullable=True)
|
||||
|
||||
user = relationship("User", back_populates="payments")
|
||||
promo_code_used = relationship("PromoCode",
|
||||
back_populates="payments_where_used")
|
||||
|
||||
|
||||
class PromoCode(Base):
|
||||
__tablename__ = "promo_codes"
|
||||
|
||||
promo_code_id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
code = Column(String, unique=True, nullable=False, index=True)
|
||||
bonus_days = Column(Integer, nullable=False)
|
||||
max_activations = Column(Integer, nullable=False)
|
||||
current_activations = Column(Integer, default=0)
|
||||
is_active = Column(Boolean, default=True)
|
||||
created_by_admin_id = Column(BigInteger, nullable=False)
|
||||
created_at = Column(DateTime(timezone=True), server_default=func.now())
|
||||
valid_until = Column(DateTime(timezone=True), nullable=True)
|
||||
|
||||
activations = relationship("PromoCodeActivation",
|
||||
back_populates="promo_code",
|
||||
cascade="all, delete-orphan")
|
||||
payments_where_used = relationship("Payment",
|
||||
back_populates="promo_code_used")
|
||||
|
||||
|
||||
class PromoCodeActivation(Base):
|
||||
__tablename__ = "promo_code_activations"
|
||||
|
||||
activation_id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
promo_code_id = Column(Integer,
|
||||
ForeignKey("promo_codes.promo_code_id"),
|
||||
nullable=False)
|
||||
user_id = Column(BigInteger, ForeignKey("users.user_id"), nullable=False)
|
||||
activated_at = Column(DateTime(timezone=True), server_default=func.now())
|
||||
payment_id = Column(Integer,
|
||||
ForeignKey("payments.payment_id"),
|
||||
nullable=True)
|
||||
|
||||
promo_code = relationship("PromoCode", back_populates="activations")
|
||||
user = relationship("User", back_populates="promo_code_activations")
|
||||
payment = relationship("Payment")
|
||||
|
||||
__table_args__ = (UniqueConstraint('promo_code_id',
|
||||
'user_id',
|
||||
name='uq_promo_user_activation'), )
|
||||
|
||||
|
||||
class MessageLog(Base):
|
||||
__tablename__ = "message_logs"
|
||||
|
||||
log_id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
user_id = Column(BigInteger,
|
||||
ForeignKey("users.user_id"),
|
||||
nullable=True,
|
||||
index=True)
|
||||
telegram_username = Column(String, nullable=True)
|
||||
telegram_first_name = Column(String, nullable=True)
|
||||
event_type = Column(String, nullable=False, index=True)
|
||||
content = Column(Text, nullable=True)
|
||||
raw_update_preview = Column(Text, nullable=True)
|
||||
timestamp = Column(DateTime(timezone=True),
|
||||
server_default=func.now(),
|
||||
index=True)
|
||||
is_admin_event = Column(Boolean, default=False)
|
||||
target_user_id = Column(BigInteger,
|
||||
ForeignKey("users.user_id"),
|
||||
nullable=True,
|
||||
index=True)
|
||||
|
||||
author_user = relationship("User",
|
||||
foreign_keys=[user_id],
|
||||
back_populates="message_logs_authored")
|
||||
target_user = relationship("User",
|
||||
foreign_keys=[target_user_id],
|
||||
back_populates="message_logs_targeted")
|
||||
|
||||
|
||||
class PanelSyncStatus(Base):
|
||||
__tablename__ = "panel_sync_status"
|
||||
|
||||
id = Column(Integer, primary_key=True, default=1, autoincrement=False)
|
||||
last_sync_time = Column(DateTime(timezone=True), nullable=True)
|
||||
status = Column(String, nullable=True)
|
||||
details = Column(Text, nullable=True)
|
||||
users_processed_from_panel = Column(Integer, default=0)
|
||||
subscriptions_synced = Column(Integer, default=0)
|
||||
|
||||
__table_args__ = (UniqueConstraint('id'), )
|
||||
Reference in New Issue
Block a user