octopus/myproject/tenants/management/commands/clear_tenant_data.py

# -*- coding: utf-8 -*-
"""
Management команда для полной очистки данных тенанта без удаления схемы и самого тенанта.

ВАЖНО: Это необратимая операция! Все данные тенанта будут удалены, но схема и запись в Client останутся.

Использование:
    # Очистка всех данных тенанта
    python manage.py clear_tenant_data --schema=anatol --noinput
    
    # С подтверждением
    python manage.py clear_tenant_data --schema=anatol
"""
from django.core.management.base import BaseCommand
from django.db import connection
from django.apps import apps
from django_tenants.utils import schema_context
from tenants.models import Client


class Command(BaseCommand):
    help = 'Полная очистка всех данных в схеме тенанта без удаления схемы и Client'

    def add_arguments(self, parser):
        parser.add_argument(
            '--schema',
            type=str,
            required=True,
            help='Имя схемы БД тенанта для очистки (пример: anatol)'
        )
        parser.add_argument(
            '--noinput',
            action='store_true',
            help='Не запрашивать подтверждение (для автоматизации)'
        )

    def handle(self, *args, **options):
        self.stdout.write(self.style.SUCCESS('\n=== Очистка данных тенанта ===\n'))

        # Получаем параметры
        schema_name = options.get('schema')
        noinput = options.get('noinput', False)

        # Проверяем что тенант существует
        try:
            tenant = Client.objects.get(schema_name=schema_name)
        except Client.DoesNotExist:
            self.stdout.write(self.style.ERROR(f'\nОШИБКА: Тенант со схемой "{schema_name}" не найден\n'))
            return

        # Получаем список всех таблиц в схеме (кроме служебных django)
        tables = self._get_tenant_tables(schema_name)
        
        # Показываем информацию о том что будет удалено
        self.stdout.write('='*70)
        self.stdout.write(self.style.WARNING('ВНИМАНИЕ! Следующие данные будут удалены:'))
        self.stdout.write('='*70)
        self.stdout.write(f'\n[Тенант]')
        self.stdout.write(f'   Название: {tenant.name}')
        self.stdout.write(f'   Схема: {tenant.schema_name}')
        self.stdout.write(f'   Владелец: {tenant.owner_email}')
        
        self.stdout.write(f'\n[База данных]')
        self.stdout.write(f'   Схема: "{schema_name}" (НЕ будет удалена)')
        self.stdout.write(f'   Будет очищено таблиц: {len(tables)}')
        
        if tables:
            self.stdout.write(f'\n[Таблицы для очистки]')
            for table in tables[:20]:  # Показываем первые 20
                self.stdout.write(f'   - {table}')
            if len(tables) > 20:
                self.stdout.write(f'   ... и ещё {len(tables) - 20} таблиц')

        self.stdout.write(f'\n[Что НЕ будет удалено]')
        self.stdout.write(f'   ✓ Запись Client (тенант)')
        self.stdout.write(f'   ✓ Схема БД "{schema_name}"')
        self.stdout.write(f'   ✓ Домены тенанта')

        self.stdout.write('\n' + '='*70)

        # Запрашиваем подтверждение
        if not noinput:
            confirm = input(self.style.WARNING('\nВы уверены? Введите "yes" для подтверждения: '))
            if confirm.lower() not in ['yes', 'да', 'y']:
                self.stdout.write(self.style.ERROR('\nОтменено\n'))
                return

        # Начинаем очистку
        self.stdout.write('\n' + '='*70)
        self.stdout.write(self.style.SUCCESS('Начинаю очистку данных...'))
        self.stdout.write('='*70 + '\n')

        try:
            deleted_count = self._clear_tenant_data(schema_name, tables)
            
            # Итоговое сообщение
            self.stdout.write('\n' + '='*70)
            self.stdout.write(self.style.SUCCESS('УСПЕХ: Данные тенанта очищены!'))
            self.stdout.write('='*70)
            self.stdout.write(self.style.WARNING('\nИтоги очистки:'))
            self.stdout.write(f'  - Очищено таблиц: {deleted_count}')
            self.stdout.write(f'  - Схема "{schema_name}" сохранена')
            self.stdout.write(f'  - Тенант "{tenant.name}" сохранён')
            
            # Предлагаем инициализировать системные данные
            self._suggest_init_data(schema_name)

        except Exception as e:
            self.stdout.write(self.style.ERROR(f'\nОШИБКА при очистке данных: {e}'))
            self.stdout.write(self.style.ERROR('   Данные могут быть частично удалены. Проверьте состояние БД.\n'))
            raise

    def _get_tenant_tables(self, schema_name):
        """
        Получает список всех таблиц в схеме тенанта (кроме служебных).
        """
        with connection.cursor() as cursor:
            cursor.execute("""
                SELECT table_name 
                FROM information_schema.tables 
                WHERE table_schema = %s 
                AND table_type = 'BASE TABLE'
                AND table_name NOT IN (
                    'django_migrations',
                    'django_content_type',
                    'django_admin_log',
                    'django_session'
                )
                ORDER BY table_name
            """, [schema_name])
            
            return [row[0] for row in cursor.fetchall()]

    def _clear_tenant_data(self, schema_name, tables):
        """
        Очищает все данные в таблицах тенанта.
        Использует TRUNCATE CASCADE для удаления связанных данных.
        """
        deleted_count = 0
        
        with connection.cursor() as cursor:
            # Устанавливаем search_path на схему тенанта
            cursor.execute(f'SET search_path TO {schema_name}')
            
            # Отключаем триггеры и ограничения для быстрой очистки
            self.stdout.write('[1] Отключение триггеров и ограничений...')
            cursor.execute('SET session_replication_role = replica;')
            
            # Очищаем каждую таблицу
            self.stdout.write(f'[2] Очистка {len(tables)} таблиц...')
            for i, table in enumerate(tables, 1):
                try:
                    # Используем TRUNCATE CASCADE для удаления связанных данных
                    cursor.execute(f'TRUNCATE TABLE {table} CASCADE;')
                    deleted_count += 1
                    
                    if i % 10 == 0:
                        self.stdout.write(f'   Очищено таблиц: {i}/{len(tables)}')
                        
                except Exception as e:
                    self.stdout.write(
                        self.style.WARNING(f'   Предупреждение: не удалось очистить таблицу {table}: {e}')
                    )
            
            # Включаем триггеры и ограничения обратно
            self.stdout.write('[3] Включение триггеров и ограничений...')
            cursor.execute('SET session_replication_role = DEFAULT;')
            
            # Сбрасываем последовательности (sequences) для ID
            self.stdout.write('[4] Сброс последовательностей (ID sequences)...')
            cursor.execute("""
                SELECT 'SELECT SETVAL(' ||
                       quote_literal(quote_ident(sequence_schema) || '.' || quote_ident(sequence_name)) ||
                       ', 1, false);'
                FROM information_schema.sequences
                WHERE sequence_schema = %s;
            """, [schema_name])
            
            for row in cursor.fetchall():
                try:
                    cursor.execute(row[0])
                except Exception as e:
                    # Игнорируем ошибки при сбросе последовательностей
                    pass
            
            self.stdout.write(self.style.SUCCESS(f'   OK: Все данные очищены'))
        
        return deleted_count
    
    def _suggest_init_data(self, schema_name):
        """
        Предлагает пользователю инициализировать системные данные.
        """
        self.stdout.write('\n' + '-'*70)
        self.stdout.write(self.style.WARNING('Рекомендация:'))
        self.stdout.write('-'*70)
        self.stdout.write('Для работы системы необходимы системные данные.')
        self.stdout.write('Выполните команду для их создания:\n')
        self.stdout.write(self.style.SUCCESS(f'  python manage.py init_tenant_data --schema={schema_name}'))
        self.stdout.write('')