from django.core.management.base import BaseCommand, CommandError from products.models import * from django.db import transaction class Command(BaseCommand): help = 'Manage serial code in batch' def add_arguments(self, parser): parser.add_argument('--tenant-id', '-t', required=True, type=int) parser.add_argument('--job', '-j', type=int) parser.add_argument('--all', '-A', action="store_true", help="apply to all code") parser.add_argument('--seq-range', "-r", help="code by seq range") parser.add_argument('--code-file', "-f", help="code list from file") parser.add_argument('--activate', "-a", action="store_true", help="activate code") parser.add_argument('--deactivate', "-d", action="store_true", help="deactivate code") parser.add_argument('--bind-product', "-b", type=int, help="bind to product by id") parser.add_argument('--unbind-product', "-u", action="store_true", help="unbind product") def handle(self, *args, **options): if options.get('job'): job = Job.objects.get(pk=options['job']) else: job = Job.objects.create(name='bind-batch.%d' % options.get('batch')) try: job.update('running', 0.0) tenant = Tenant.objects.get(pk=options['tenant_id']) query = self.build_query(options) total = query.count() done = 0 prod = None if options.get("bind_product"): prod = Product.objects.get(tenant=tenant, pk=options['bind_product']) for batch_ids in self.iterate_batches(query, 10000): uq = SerialCode.objects.filter(pk__in=batch_ids) if options.get("activate"): uq.update(is_active=True) elif options.get("deactivate"): uq.update(is_active=False) elif options.get("unbind_product"): uq.update(product=None) elif options.get("bind_product"): uq.update(product=prod) transaction.commit() done += len(batch_ids) print("code batch op", done, total) perc = done * 100.0 / total job.update('running', perc) job.update('done', 100.0) except Exception as e: job.update('error', message=str(e)) raise def iterate_batches(self, query, batchsize): total = query.count() cur = 0 while cur < total: yield [x.pk for x in query[cur : cur + batchsize]] cur += batchsize def build_query(self, options): tenant = Tenant.objects.get(pk=options['tenant_id']) query = SerialCode.objects.filter(tenant=tenant).order_by('pk') if options.get("seq_range"): begin, end = options['seq_range'].split(',', maxsplit=1) query = query.filter(seq_num__gte=int(begin), seq_num__lte=int(end)) elif options.get("code_file"): codes = self.read_code_file(options['code_file']) query = query.filter(code__in=codes) elif options.get("all"): pass else: raise Exception("Code not specified") return query def read_code_file(self, cf): with open(cf, 'r') as f: return [x for x in f.read().splitlines() if not x.strip().startswith('#')]