from django.core.management.base import BaseCommand from django.db import connections, transaction from tqdm import tqdm from account.models import OldRole, Role, User, UserRole from establishment.management.commands.add_position import namedtuplefetchall from main.models import SiteSettings class Command(BaseCommand): help = '''Add site affilations from old db to new db. Run after migrate account models!!!''' def map_role_sql(self): with connections['legacy'].cursor() as cursor: cursor.execute(''' select distinct case when role = 'news_editor' then 'CONTENT_PAGE_MANAGER' when role in ('reviewer', 'reviwer', 'reviewer_manager') then 'REVIEWER_MANGER' when role = 'admin' then 'SUPERUSER' when role ='community_manager' then 'COUNTRY_ADMIN' when role = 'site_admin' then 'COUNTRY_ADMIN' when role = 'wine_reviewer' then 'WINERY_REVIEWER' when role in ('salesman', 'sales_man') then 'SALES_MAN' when role = 'seller' then 'SELLER' else role end as new_role, case when role = 'GUEST' then null else role end as role from ( SELECT DISTINCT COALESCE(role, 'GUEST') as role FROM site_affiliations AS sa ) t ''') return namedtuplefetchall(cursor) def add_old_roles(self): objects = [] OldRole.objects.all().delete() for s in tqdm(self.map_role_sql(), desc='Add permissions old'): objects.append( OldRole(new_role=s.new_role, old_role=s.role) ) OldRole.objects.bulk_create(objects) self.stdout.write(self.style.WARNING(f'Migrated old roles.')) def site_role_sql(self): with connections['legacy'].cursor() as cursor: cursor.execute(''' select site_id, role from ( SELECT DISTINCT site_id, COALESCE(role, 'GUEST') as role FROM site_affiliations AS sa ) t where t.role not in ('admin', 'GUEST') ''') return namedtuplefetchall(cursor) def add_site_role(self): objects = [] for s in tqdm(self.site_role_sql(), desc='Add site role'): old_role = OldRole.objects.get(old_role=s.role) role_choice = getattr(Role, old_role.new_role) sites = SiteSettings.objects.filter(old_id=s.site_id) for site in sites: data = {'site': site, 'role': role_choice} role = Role.objects.filter(**data) if not role.exists(): objects.append( Role(**data) ) Role.objects.bulk_create(objects) self.stdout.write(self.style.WARNING(f'Added site roles.')) def update_site_role(self): roles = Role.objects.filter(country__isnull=True).select_related('site') \ .filter(site__id__isnull=False).select_for_update() with transaction.atomic(): for role in tqdm(roles, desc='Update role country'): role.country = role.site.country role.save() self.stdout.write(self.style.WARNING(f'Updated site roles.')) def user_role_sql(self): with connections['legacy'].cursor() as cursor: cursor.execute(''' select t.* from ( SELECT site_id, account_id, COALESCE(role, 'GUEST') as role FROM site_affiliations AS sa ) t join accounts a on a.id = t.account_id where t.role not in ('admin', 'GUEST') ''') return namedtuplefetchall(cursor) def add_role_user(self): for s in tqdm(self.user_role_sql(), desc='Add role to user'): sites = SiteSettings.objects.filter(old_id=s.site_id) old_role = OldRole.objects.get(old_role=s.role) role_choice = getattr(Role, old_role.new_role) roles = Role.objects.filter(site__in=[site for site in sites], role=role_choice) users = User.objects.filter(old_id=s.account_id) for user in users: for role in roles: UserRole.objects.get_or_create(user=user, role=role, state=UserRole.VALIDATED) self.stdout.write(self.style.WARNING(f'Added users roles.')) def superuser_role_sql(self): with connections['legacy'].cursor() as cursor: cursor.execute(''' select t.* from ( SELECT site_id, account_id, COALESCE(role, 'GUEST') as role FROM site_affiliations AS sa ) t join accounts a on a.id = t.account_id where t.role in ('admin') ''') return namedtuplefetchall(cursor) def add_superuser(self): for s in tqdm(self.superuser_role_sql(), desc='Add superuser'): users = User.objects.filter(old_id=s.account_id).select_for_update() with transaction.atomic(): for user in users: user.is_superuser = True user.save() self.stdout.write(self.style.WARNING(f'Added superuser.')) def handle(self, *args, **kwargs): self.add_old_roles() self.add_site_role() self.update_site_role() self.add_role_user() self.add_superuser()