85 lines
3.1 KiB
Python
85 lines
3.1 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
|
||
|
|
from django.core.management.base import BaseCommand
|
||
|
|
from collections import defaultdict
|
||
|
|
from django.apps import apps
|
||
|
|
from django.db import transaction
|
||
|
|
|
||
|
|
from django.contrib.auth.models import User
|
||
|
|
from django.db.models import Model
|
||
|
|
from polymorphic.models import PolymorphicModel
|
||
|
|
|
||
|
|
|
||
|
|
class ImportCommand(BaseCommand):
|
||
|
|
"""
|
||
|
|
Generic command for import of NK15 database
|
||
|
|
"""
|
||
|
|
|
||
|
|
def print_success(self, to_print):
|
||
|
|
return self.stdout.write(self.style.SUCCESS(to_print))
|
||
|
|
|
||
|
|
def print_error(self, to_print):
|
||
|
|
return self.stdout.write(self.style.ERROR(to_print))
|
||
|
|
|
||
|
|
def update_line(self, n, total, content):
|
||
|
|
n = str(n)
|
||
|
|
total = str(total)
|
||
|
|
n.rjust(len(total))
|
||
|
|
print(f"\r ({n}/{total}) {content:10.10}", end="")
|
||
|
|
|
||
|
|
def create_parser(self, prog_name, subcommand, **kwargs):
|
||
|
|
parser = super().create_parser(prog_name, subcommand, **kwargs)
|
||
|
|
parser.add_argument('--nk15db', action='store', default='nk15', help='NK15 database name')
|
||
|
|
parser.add_argument('--nk15user', action='store', default='nk15_user', help='NK15 database owner')
|
||
|
|
parser.add_argument('-s', '--save', action='store', help="save mapping of idbde")
|
||
|
|
parser.add_argument('-m', '--map', action='store', help="import mapping of idbde")
|
||
|
|
return parser
|
||
|
|
|
||
|
|
|
||
|
|
class BulkCreateManager(object):
|
||
|
|
"""
|
||
|
|
This helper class keeps track of ORM objects to be created for multiple
|
||
|
|
model classes, and automatically creates those objects with `bulk_create`
|
||
|
|
when the number of objects accumulated for a given model class exceeds
|
||
|
|
`chunk_size`.
|
||
|
|
Upon completion of the loop that's `add()`ing objects, the developer must
|
||
|
|
call `done()` to ensure the final set of objects is created for all models.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, chunk_size=100):
|
||
|
|
self._create_queues = defaultdict(list)
|
||
|
|
self.chunk_size = chunk_size
|
||
|
|
|
||
|
|
def _commit(self, model_class):
|
||
|
|
model_key = model_class._meta.label
|
||
|
|
if model_class.__base__ in [Model, PolymorphicModel] or model_class is User:
|
||
|
|
model_class.objects.bulk_create(self._create_queues[model_key])
|
||
|
|
else:
|
||
|
|
# ensure that parents models exists
|
||
|
|
self._commit(model_class.__base__)
|
||
|
|
with transaction.atomic():
|
||
|
|
for obj in self._create_queues[model_key]:
|
||
|
|
obj.save_base(raw=True)
|
||
|
|
self._create_queues[model_key] = []
|
||
|
|
|
||
|
|
def add(self, *args):
|
||
|
|
"""
|
||
|
|
Add an object to the queue to be created, and call bulk_create if we
|
||
|
|
have enough objs.
|
||
|
|
"""
|
||
|
|
for obj in args:
|
||
|
|
model_class = type(obj)
|
||
|
|
model_key = model_class._meta.label
|
||
|
|
self._create_queues[model_key].append(obj)
|
||
|
|
if len(self._create_queues[model_key]) >= self.chunk_size:
|
||
|
|
self._commit(model_class)
|
||
|
|
|
||
|
|
def done(self):
|
||
|
|
"""
|
||
|
|
Always call this upon completion to make sure the final partial chunk
|
||
|
|
is saved.
|
||
|
|
"""
|
||
|
|
for model_name, objs in self._create_queues.items():
|
||
|
|
if len(objs) > 0:
|
||
|
|
self._commit(apps.get_model(model_name))
|