import os from pathlib import Path from celery import shared_task from celery.utils.log import get_task_logger from django.conf import settings from django.contrib.auth.models import User from django.core.management import call_command from storage.export import export from storage.models import DataStorage, get_storage_path _log = get_task_logger(__name__) MEDIA_ROOT = getattr(settings, 'MEDIA_ROOT') @shared_task def export_task(export_type, user_pk, model_name, query_parameters): return export(export_type, user_pk, model_name, query_parameters) @shared_task def dump_db(pk, dump_models): # Add store record user = User.objects.get(pk=pk) _log.info(f'User [{user}] request db dump for models {dump_models}') store = DataStorage(type=DataStorage.Type.DB_DUMP, format=DataStorage.Format.JSON, user=user, size=0, file='name') store.save() file_name = get_storage_path(store, 'dump.json') fpath = os.path.join(MEDIA_ROOT, file_name) dirname = os.path.dirname(fpath) Path(dirname).mkdir(parents=True, exist_ok=True) with open(fpath, 'w') as f: call_command('dumpdata', *dump_models, stdout=f) store.size = os.path.getsize(fpath) store.file = file_name store.save() _log.info(f'Dump ready in [{store.file.name}] with size of [{store.size}]') store.update_crc() return store.pk def remove_vulnerabilities(data_dict): """ Remove column vulnerabilities from data return: DataDict without vulnerabilities column """ for item in data_dict: try: item.pop('vulnerabilities') except KeyError: continue return data_dict