66 lines
1.7 KiB
Python
66 lines
1.7 KiB
Python
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
|