from django.core.management.base import BaseCommand

from establishment.models import Establishment
from transfer.models import Establishments as OldEst


class Command(BaseCommand):
    help = 'Update transportation column in establishment'

    def handle(self, *args, **kwargs):
        count = 0
        raw_qs = OldEst.objects.raw(
            '''
                select
                    e.id,
                    l.transportation 
                from establishments e 
                join locations l on l.id = e.location_id 
                where trim(l.transportation) is not null 
                    and length(trim(l.transportation)) >1;
            '''
        )
        queryset = [vars(query) for query in raw_qs]
        for obj in queryset:
            establishment = Establishment.objects.filter(old_id=obj['id']).first()
            if establishment:
                establishment.transportation = obj['transportation']
                count += 1
                establishment.save()
        self.stdout.write(self.style.WARNING(f'Updated {count} objects.'))