trans2es_data2es.py

# -*- coding: UTF-8 -*-

from multiprocessing import Pool
from functools import partial

from django.core.management.base import BaseCommand

from data_sync.tasks import write_to_es
from data_sync.type_info import get_type_info_map


def process(dtype, pks):
    write_to_es(dtype, pks)


class Command(BaseCommand):
    def add_arguments(self, parser):
        parser.add_argument('-t', '--type', dest='dtype')
        parser.add_argument('-p', '--procs', dest='parallel', type=int)

    def handle(self, *args, **options):
        t = options['dtype']
        p = options['parallel'] or 2
        pool = Pool(processes=p)

        dt = get_type_info_map()[t]
        pks = dt.model.objects.all().values_list('id', flat=True)
        pks = list(pks)
        n = int(len(pks) / p)
        _process = partial(process, t)
        pool.map(_process, [pks[i:i + n] for i in range(0, len(pks), n)])