@@ -294,10 +294,14 @@ CELERY_BEAT_SCHEDULE = { 'task': 'payments.tasks.send_low_balance_message', 'schedule': crontab(0, 8), }, - # 'execute_recurring_payments': { - # 'task': 'payments.tasks.execute_recurring_payments', - # 'schedule': crontab(*env.list('RECURRING_PAYMENT_CRONTAB_SCHEDULE', [])), - # }, + 'execute_recurring_payments': { + 'task': 'payments.tasks.execute_recurring_payments', + 'schedule': crontab(*env.list('RECURRING_PAYMENT_CRONTAB_SCHEDULE', [])), + }, + 'revoke_recurring_payments': { + 'task': 'payments.tasks.revoke_recurring_payments', + 'schedule': crontab(*env.list('RECURRING_PAYMENT_CRONTAB_SCHEDULE', [])), + }, } CACHES = { @@ -1,6 +1,7 @@ from types import MappingProxyType -from typing import List, Mapping +from typing import List, Mapping, Never +import orjson from UnleashClient import UnleashClient from django.conf import settings @@ -16,7 +17,7 @@ class UnleashFeatureFlagService(FeatureFlagService): app_name=settings.FEATURE_FLAG_APP_NAME, instance_id=settings.FEATURE_FLAG_INSTANCE_ID, cache=UnleashRedisCache(), - environment=settings.FEATURE_FLAG_APP_NAME + environment=settings.FEATURE_FLAG_APP_NAME, ) def get_flag_state_by_emails(self, name: str, emails: List[Email]) -> Mapping[Email, State]: @@ -24,3 +25,18 @@ class UnleashFeatureFlagService(FeatureFlagService): def get_flag_state(self, name: str, email: Email) -> State: return self.client.is_enabled(feature_name=name, context={'userId': email}) + + def is_feature_enabled(self, name: str) -> bool: + for feature in orjson.loads(self.client.engine.get_state())['features']: + if feature['name'] == name: + return bool(feature.get('enabled')) + raise ValueError(f'Feature flag "{name}" not found in Unleash state') + + def get_user_emails(self, name: str) -> List[Email] | List[Never]: + for feature in orjson.loads(self.client.engine.get_state())['features']: + if feature['name'] == name: + userIds = feature['strategies'][0]['parameters'].get('userIds') + if userIds: + return userIds.replace(' ', '').split(',') + return [] + raise ValueError(f'Feature flag "{name}" not found in Unleash state') @@ -1,3 +1,4 @@ +import hashlib import logging from datetime import timedelta @@ -102,7 +103,8 @@ class PaymentService: return await sync_to_async(YookassaPayment.find_one)(payment_id) def handle_captured_payment(self, payment_id: UUID) -> None: - YookassaPayment.capture(str(payment_id)) + idempotency_key = hashlib.sha256(f'capture:{payment_id}'.encode('utf-8')).hexdigest() + YookassaPayment.capture(str(payment_id), idempotency_key=idempotency_key) logger.info('Payment captured: payment_id=%s email=%s', payment_id, self.user.email) def calculate_buying_tokens(self, plan: PaymentPlan): @@ -1,8 +1,13 @@ +import hashlib +from datetime import timezone as dt_timezone from decimal import Decimal -from uuid import UUID, uuid4 +from itertools import islice +from uuid import UUID from celery import shared_task from celery.utils.log import get_task_logger +from django.core.cache import cache +from django.db import transaction from django.db.models import F from django.utils import timezone @@ -10,14 +15,15 @@ from authentication.models.business_host import BusinessUserHost from authentication.models.user import CustomUserModel from authentication.services.email_service import EmailService from lib.unleash.client import celery_client -from payments.models import PaymentPlanUserInfo, PaymentMethod -from payments.selectors.payment_plan_selector import PaymentPlanSelector +from payments.models import PaymentMethod, PaymentPlan, PaymentPlanUserInfo from payments.services.payment_plan_service import PaymentPlanService - from yookassa import Payment as YookassaPayment logger = get_task_logger(__name__) +CHUNK_SIZE = 1000 +LOCK_TIMEOUT_SECONDS = 10 * 60 + @shared_task def send_low_balance_message(): @@ -38,39 +44,42 @@ def withdraw(user_id: UUID, amount: Decimal): @shared_task def execute_recurring_payments() -> None: - overdue_payments = PaymentPlanUserInfo.objects.select_related('user', 'plan', 'method').filter( - next_payment_at__isnull=False, - next_payment_at__lte=timezone.now(), - plan__price__gt=0, - plan__individual=False, + if not celery_client.is_feature_enabled('recurring_payments'): + logger.info('Recurring payments feature disabled, skipping execute') + return + + emails = celery_client.get_user_emails('recurring_payments') + overdue_payments = ( + PaymentPlanUserInfo.objects.select_related('user', 'plan', 'method') + .filter( + next_payment_at__isnull=False, + next_payment_at__lte=timezone.now(), + plan__price__gt=0, + plan__individual=False, + plan__is_corporate=False, + method__isnull=False, + ) + .only( + 'uid', + 'next_payment_at', + 'user_id', + 'plan_id', + 'method_id', + 'user__uid', + 'user__email', + 'plan__uid', + 'plan__price', + 'plan__tokens_per_plan', + 'method__uid', + 'method__payment_method_id', + 'method__attempts', + ) ) - canceled_recurring_payments = [] - logger.info('Recurring payments task started: overdue_count=%s', overdue_payments.count()) - for overdue_payment in overdue_payments: + if emails: + overdue_payments = overdue_payments.filter(user__email__in=emails) + for overdue_payment in overdue_payments.iterator(chunk_size=CHUNK_SIZE): customer = overdue_payment.user plan = overdue_payment.plan - if not celery_client.get_flag_state('recurring_payments', overdue_payment.user.email): - overdue_payment.next_payment_at = None - canceled_recurring_payments.append(overdue_payment) - logger.info( - 'Recurring payment canceled by feature flag: email=%s plan_uid=%s', - customer.email, - plan.uid, - ) - continue - if not overdue_payment.is_recurring: - free_plan = PaymentPlanSelector(customer).get_free_plan(plan.is_corporate) - overdue_payment.next_payment_at = None - overdue_payment.plan = free_plan - overdue_payment.current_token_balance = 0 - canceled_recurring_payments.append(overdue_payment) - logger.info( - 'Recurring payment canceled due to missing method: email=%s plan_uid=%s switched_to_free_plan_uid=%s', - customer.email, - plan.uid, - free_plan.uid, - ) - continue receipt_data = { 'customer': {'email': customer.email}, 'items': [ @@ -93,7 +102,14 @@ def execute_recurring_payments() -> None: 'plan_uid': str(plan.uid), }, } - YookassaPayment.create(payment_data, uuid4()) + dt = overdue_payment.next_payment_at + if timezone.is_naive(dt): + dt = timezone.make_aware(dt) + period = dt.astimezone(dt_timezone.utc).replace(microsecond=0).isoformat() + idempotency_key = hashlib.sha256( + f'recurring:{customer.uid}:{plan.uid}:{period}:{overdue_payment.method.attempts}'.encode('utf-8') + ).hexdigest() + YookassaPayment.create(payment_data, idempotency_key=idempotency_key) logger.info( 'Recurring payment initiated: email=%s plan_uid=%s amount=%s method_uid=%s', customer.email, @@ -101,13 +117,63 @@ def execute_recurring_payments() -> None: plan.price, overdue_payment.method.uid, ) - methods_for_delete = [crp.method.uid for crp in canceled_recurring_payments if crp.method] - PaymentPlanUserInfo.objects.bulk_update( - canceled_recurring_payments, fields=['next_payment_at', 'plan', 'current_token_balance'] - ) - deleted_methods_count, deleted_details = PaymentMethod.objects.filter(uid__in=methods_for_delete).delete() - logger.info( - 'Recurring payments task finished: canceled_count=%s deleted_methods=%s', - len(canceled_recurring_payments), - deleted_methods_count, - ) \ No newline at end of file + + +@shared_task +def revoke_recurring_payments() -> None: + lock_key = 'recurring_payments:revoke_lock' + if not cache.add(lock_key, 'locked', timeout=LOCK_TIMEOUT_SECONDS): + logger.info('Revoke recurring already running, skipping') + return + try: + feature_name = 'recurring_payments' + base_qs = PaymentPlanUserInfo.objects.filter( + next_payment_at__isnull=False, + next_payment_at__lte=timezone.now(), + plan__price__gt=0, + plan__individual=False, + plan__is_corporate=False, + ) + revoked_count = 0 + canceled_count = 0 + deleted_methods_count = 0 + + if not celery_client.is_feature_enabled(feature_name): + flag_off_qs = base_qs + flag_on_qs = base_qs.none() + else: + allowed_emails = celery_client.get_user_emails(feature_name) + if allowed_emails: + flag_off_qs = base_qs.exclude(user__email__in=allowed_emails) + flag_on_qs = base_qs.filter(user__email__in=allowed_emails, method__isnull=True) + else: + flag_off_qs = base_qs.none() + flag_on_qs = base_qs.filter(method__isnull=True) + + uid_iter = flag_off_qs.values_list('uid', flat=True).iterator(chunk_size=CHUNK_SIZE) + while uids := list(islice(uid_iter, CHUNK_SIZE)): + with transaction.atomic(): + deleted_methods_count += PaymentMethod.objects.filter( + user_plan_info__uid__in=uids + ).delete()[0] + revoked_count += flag_off_qs.filter(uid__in=uids).update(next_payment_at=None) + + if flag_on_qs.exists(): + free_regular_plan = PaymentPlan.objects.get_or_create(price=0, is_corporate=False)[0] + flag_on_iter = flag_on_qs.values_list('uid', flat=True).iterator(chunk_size=CHUNK_SIZE) + while uids := list(islice(flag_on_iter, CHUNK_SIZE)): + with transaction.atomic(): + canceled_count += flag_on_qs.filter(uid__in=uids).update( + next_payment_at=None, + plan_id=free_regular_plan.pk, + current_token_balance=0, + ) + + logger.info( + 'Revoke recurring finished: revoked=%s free_regular=%s deleted_methods=%s', + revoked_count, + canceled_count, + deleted_methods_count, + ) + finally: + cache.delete(lock_key)