@@ -0,0 +1,230 @@ +import re +from datetime import datetime +from decimal import Decimal, InvalidOperation +from pathlib import Path +from typing import Any, NamedTuple + +from django.core.management.base import BaseCommand, CommandError +from django.db.migrations.loader import MigrationLoader + +from ml_model.models import NeuronModel +from payments.models import PaymentPlanFeature + +MEASUREMENT_UNITS = {choice.value for choice in PaymentPlanFeature.MeasurementUnitChoices} +PAYMENTS_APP_LABEL = 'payments' + + +def _sanitize_slug(slug: str) -> str: + return re.sub(r'[^a-zA-Z0-9_]', '_', slug) + + +class PaymentsMigrationSlot(NamedTuple): + """dependency — лист графа для dependencies; next_number — префикс нового файла.""" + + dependency: str + next_number: int + + +def _resolve_payments_migration_slot() -> PaymentsMigrationSlot: + """Один проход MigrationLoader: хвост ветки и номер следующей миграции.""" + loader = MigrationLoader(None, ignore_no_migrations=True) + + leaves = sorted(name for app, name in loader.graph.leaf_nodes() if app == PAYMENTS_APP_LABEL) + if not leaves: + raise CommandError(f'В приложении {PAYMENTS_APP_LABEL} нет миграций.') + if len(leaves) > 1: + raise CommandError( + f'У {PAYMENTS_APP_LABEL} несколько концов веток: {", ".join(leaves)}. ' + f'Сначала смержите: python manage.py makemigrations {PAYMENTS_APP_LABEL} --merge' + ) + + dependency = leaves[0] + prefix = dependency.split('_', 1)[0] + if not prefix.isdigit(): + raise CommandError(f'Не удалось вычислить номер следующей миграции из листа "{dependency}".') + return PaymentsMigrationSlot(dependency=dependency, next_number=int(prefix) + 1) + + +def _migration_basename(model_slug: str, next_number: int) -> str: + safe_slug = _sanitize_slug(model_slug) + return f'{next_number:04d}_add_{safe_slug}_payment_features' + + +def _build_migration_source( + *, + func_name: str, + model_slug: str, + price: Decimal, + measurement_unit: str, + price_threshold: int, + latest_migration: str, +) -> str: + return f"""# Generated by makemigration_payment_features on {datetime.now():%Y-%m-%d %H:%M} + +import math +from decimal import Decimal + +from django.db import migrations +from django.db.models import Max + + +def {func_name}(apps, schema_editor): + PaymentPlan = apps.get_model('payments', 'PaymentPlan') + PaymentPlanFeature = apps.get_model('payments', 'PaymentPlanFeature') + NeuronModel = apps.get_model('ml_model', 'NeuronModel') + + price = Decimal('{price}') + measurement_unit = '{measurement_unit}' + price_threshold = {price_threshold} + model = NeuronModel.objects.get(slug='{model_slug}') + category = model.category + + max_order_by_plan_id = {{ + row['plan_id']: row['max_order'] + for row in PaymentPlanFeature.objects.filter(model__category=category) + .values('plan_id') + .annotate(max_order=Max('order')) + }} + + features = [] + for plan in PaymentPlan.objects.filter(price__gt=price_threshold): + quantity = math.floor(plan.tokens_per_plan / price) + next_order = (max_order_by_plan_id.get(plan.pk) or -1) + 1 + max_order_by_plan_id[plan.pk] = next_order + features.append( + PaymentPlanFeature( + plan=plan, + model=model, + quantity=quantity, + measurement_unit=measurement_unit, + order=next_order, + ) + ) + PaymentPlanFeature.objects.bulk_create( + features, + update_conflicts=True, + update_fields=['quantity', 'measurement_unit'], + unique_fields=['plan', 'model'], + ) + + +class Migration(migrations.Migration): + + dependencies = [ + ('payments', '{latest_migration}'), + ] + + operations = [ + migrations.RunPython({func_name}, migrations.RunPython.noop), + ] +""" + + +class Command(BaseCommand): + help = ( + 'Создаёт data-миграцию payments с PaymentPlanFeature для указанной модели. ' + 'Применение: python manage.py migrate payments' + ) + + def add_arguments(self, parser): + parser.add_argument('modelname', type=str, help='Slug модели (NeuronModel)') + parser.add_argument('--price', type=str, help='Цена за единицу измерения (токены)') + parser.add_argument( + '--measurement-unit', + type=str, + choices=sorted(MEASUREMENT_UNITS), + help='Единица измерения: text_page, file, time', + ) + parser.add_argument( + '--price-threshold', + type=int, + help='Создавать фичи только для планов с price строго больше этого значения', + ) + + def handle(self, *args: Any, **options: Any) -> None: + model_slug = options['modelname'].strip() + self._validate_model_name(model_slug) + price = self._read_price(options.get('price')) + measurement_unit = self._read_measurement_unit(options.get('measurement_unit')) + price_threshold = self._read_price_threshold(options.get('price_threshold')) + + self.stdout.write( + f'Миграция для model={model_slug}, price={price}, ' + f'measurement_unit={measurement_unit}, price_threshold={price_threshold}' + ) + self.stdout.write('Расчёт quantity выполнится на стейдже/проде при migrate по планам в БД.') + + slot = _resolve_payments_migration_slot() + migration_name = _migration_basename(model_slug, slot.next_number) + func_name = f'add_{_sanitize_slug(model_slug)}_payment_features' + source = _build_migration_source( + func_name=func_name, + model_slug=model_slug, + price=price, + measurement_unit=measurement_unit, + price_threshold=price_threshold, + latest_migration=slot.dependency, + ) + + migration_path = ( + Path(__file__).resolve().parent.parent.parent / 'migrations' / f'{migration_name}.py' + ) + if migration_path.exists(): + raise CommandError(f'Файл миграции уже существует: {migration_path}') + + migration_path.write_text(source, encoding='utf-8') + self.stdout.write(self.style.SUCCESS(f'Создана миграция: {migration_path}')) + self.stdout.write('Примените: python manage.py migrate payments') + + def _read_price(self, value: str | None) -> Decimal: + if value is not None: + try: + price = Decimal(value) + except InvalidOperation as exc: + raise CommandError(f'Некорректная цена: {value}') from exc + if price <= 0: + raise CommandError('Цена должна быть больше 0.') + return price + + while True: + raw = input('Введите прайс за единицу измерения (токены): ').strip() + try: + price = Decimal(raw) + except InvalidOperation: + self.stderr.write('Введите число.') + continue + if price <= 0: + self.stderr.write('Цена должна быть больше 0.') + continue + return price + + def _read_measurement_unit(self, value: str | None) -> str: + units = ', '.join(sorted(MEASUREMENT_UNITS)) + + if value is not None: + if value not in MEASUREMENT_UNITS: + raise CommandError(f'Недопустимая ед. измерения: {value}. Допустимые: {units}') + return value + + while True: + raw = input(f'Введите ед. измерения ({units}): ').strip() + if raw in MEASUREMENT_UNITS: + return raw + self.stderr.write(f'Допустимые значения: {units}') + + def _read_price_threshold(self, value: int | None) -> int: + if value is not None: + return value + + while True: + raw = input('Введите планы выше какого price будут учитываться: ').strip() + try: + return int(raw) + except ValueError: + self.stderr.write('Введите целое число.') + + def _validate_model_name(self, model_slug: str) -> None: + if NeuronModel.objects.filter(slug=model_slug).exists(): + return + + raise CommandError(f'NeuronModel со slug "{model_slug}" не найдена.')