from django.db.models import Q from ishtar_common.serializers_utils import generic_get_results, archive_serialization from archaeological_context_records import models from archaeological_finds.models import Find from archaeological_finds.serializers import ( generate_warehouse_queryset as finds_generate_warehouse_queryset, ) CR_MODEL_LIST = [models.Dating, models.ContextRecord, models.RecordRelations] # TODO: associated documents def generate_warehouse_queryset(ids): base_query_key = "base_finds__find" q_cr, q_record_relation = None, None for container_key in ("container", "container_ref"): for warehouse_key in ("location", "responsible"): q_c = Q( **{ "{}__{}__{}__id__in".format( base_query_key, container_key, warehouse_key ): ids } ) q_r = Q( **{ "left_record__{}__{}__{}__id__in".format( base_query_key, container_key, warehouse_key ): ids, "right_record__{}__{}__{}__id__in".format( base_query_key, container_key, warehouse_key ): ids, } ) if not q_cr: q_cr = q_c q_record_relation = q_r else: q_cr |= q_c q_record_relation |= q_r result_queryset = { models.ContextRecord.__name__: models.ContextRecord.objects.filter(q_cr), models.RecordRelations.__name__: models.RecordRelations.objects.filter(q_r), } return result_queryset def cr_serialization( archive=False, return_empty_types=False, archive_name=None, operation_queryset=None, site_queryset=None, cr_queryset=None, find_queryset=None, warehouse_queryset=None, get_queryset=False, no_geo=True, put_locks=False, lock_user=None, ): result_queryset = {} find_ids, cr_ids = None, None if operation_queryset: operation_ids = operation_queryset.values_list("id", flat=True) result_queryset = { models.ContextRecord.__name__: models.ContextRecord.objects.filter( operation_id__in=operation_ids ), models.RecordRelations.__name__: models.RecordRelations.objects.filter( left_record__operation_id__in=operation_ids, right_record__operation_id__in=operation_ids, ), } cr_ids = list( result_queryset[models.ContextRecord.__name__].values_list("id", flat=True) ) find_ids = list( Find.objects.filter( base_finds__context_record__operation_id__in=operation_ids ).values_list("id", flat=True) ) elif site_queryset: site_ids = site_queryset.values_list("id", flat=True) result_queryset = { models.ContextRecord.__name__: models.ContextRecord.objects.filter( operation__archaeological_sites__id__in=site_ids ), models.RecordRelations.__name__: models.RecordRelations.objects.filter( left_record__operation__archaeological_sites__id__in=site_ids, right_record__operation__archaeological_sites__id__in=site_ids, ), } f_q = { "base_finds__context_record__operation__archaeological_sites__id__in": site_ids } find_ids = list(Find.objects.filter(**f_q).values_list("id", flat=True)) elif cr_queryset: cr_ids = cr_queryset.values_list("id", flat=True) result_queryset = { models.ContextRecord.__name__: cr_queryset, models.RecordRelations.__name__: models.RecordRelations.objects.filter( left_record_id__in=cr_ids, right_record_id__in=cr_ids, ), } find_ids = list( Find.objects.filter(base_finds__context_record__in=cr_ids).values_list( "id", flat=True ) ) elif find_queryset: find_ids = find_queryset.values_list("id", flat=True) result_queryset = { models.ContextRecord.__name__: models.ContextRecord.objects.filter( base_finds__find__id__in=find_ids, ), models.RecordRelations.__name__: models.RecordRelations.objects.filter( left_record__base_finds__find__id__in=find_ids, right_record__base_finds__find__id__in=find_ids, ), } elif warehouse_queryset: warehouse_ids = warehouse_queryset.values_list("id", flat=True) result_queryset = generate_warehouse_queryset(warehouse_ids) w_queryset = finds_generate_warehouse_queryset(warehouse_ids) find_ids = w_queryset[Find.__name__].values_list("id", flat=True) if result_queryset: cr_ids = list( result_queryset[models.ContextRecord.__name__].values_list("id", flat=True) ) result_queryset[models.Dating.__name__] = models.Dating.objects.filter( Q(context_records__id__in=cr_ids) | Q(find__id__in=list(find_ids)) ) if get_queryset: return result_queryset result = generic_get_results( CR_MODEL_LIST, "context_records", result_queryset=result_queryset, no_geo=no_geo ) if put_locks: for model in CR_MODEL_LIST: if not hasattr(model, "locked"): continue q = model.objects if result_queryset and model.__name__ in result_queryset: q = result_queryset[model.__name__] q.update(locked=True, lock_user=lock_user) full_archive = archive_serialization( result, archive_dir="context_records", archive=archive, return_empty_types=return_empty_types, archive_name=archive_name, ) return full_archive