from django.db.models import Q from ishtar_common.serializers import generic_get_results, \ archive_serialization from archaeological_warehouse import models WAREHOUSE_MODEL_LIST = [ models.Warehouse, models.Container, models.WarehouseDivisionLink, models.ContainerLocalisation ] # TODO: associated documents def generate_queryset(base_query_key, ids): warehouse_division_q, warehouse_q, container_q = None, None, None container_loca_q = None for find_key in ("finds", "finds_ref"): base_key = "{}__{}".format(find_key, base_query_key) q = Q(**{base_key: ids}) q_loca = Q( **{"container__{}__{}".format(find_key, base_query_key): ids}) if not container_q: container_q = q container_loca_q = q_loca else: container_q |= q container_loca_q |= q_loca for container_key in ("containers", "owned_containers"): q = Q(** {"{}__{}__{}".format(container_key, find_key, base_query_key): ids}) q_div = Q(** {"warehouse__{}__{}__{}".format( container_key, find_key, base_query_key): ids}) if not warehouse_q: warehouse_q = q warehouse_division_q = q_div else: warehouse_q |= q warehouse_division_q |= q_div result_queryset = { models.Warehouse.__name__: models.Warehouse.objects.filter(warehouse_q), models.Container.__name__: models.Container.objects.filter(container_q), models.WarehouseDivisionLink.__name__: models.WarehouseDivisionLink.objects.filter( warehouse_division_q), models.ContainerLocalisation.__name__: models.ContainerLocalisation.objects.filter( container_loca_q) } return result_queryset def warehouse_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): result_queryset = {} if operation_queryset: operation_ids = operation_queryset.values_list("id", flat=True) base_query_key = "base_finds__context_record__operation_id__in" result_queryset = generate_queryset(base_query_key, operation_ids) elif site_queryset: site_ids = site_queryset.values_list("id", flat=True) base_query_key = "base_finds__context_record__"\ "archaeological_site_id__in" result_queryset = generate_queryset(base_query_key, site_ids) elif cr_queryset: cr_ids = cr_queryset.values_list("id", flat=True) base_query_key = "base_finds__context_record__id__in" result_queryset = generate_queryset(base_query_key, cr_ids) elif find_queryset: find_ids = find_queryset.values_list("id", flat=True) base_query_key = "id__in" result_queryset = generate_queryset(base_query_key, find_ids) elif warehouse_queryset: warehouse_ids = warehouse_queryset.values_list("id", flat=True) result_queryset = { models.Warehouse.__name__: warehouse_queryset, models.Container.__name__: models.Container.objects.filter( Q(location__id__in=warehouse_ids) | Q(responsible__id__in=warehouse_ids) ), models.WarehouseDivisionLink.__name__: models.WarehouseDivisionLink.objects.filter( warehouse_id__in=warehouse_ids), models.ContainerLocalisation.__name__: models.ContainerLocalisation.objects.filter( division__warehouse_id__in=warehouse_ids) } result = generic_get_results(WAREHOUSE_MODEL_LIST, "warehouse", result_queryset=result_queryset) full_archive = archive_serialization( result, archive_dir="warehouse", archive=archive, return_empty_types=return_empty_types, archive_name=archive_name, ) return full_archive