1 files changed, 28 insertions, 12 deletions
diff --git a/archaeological_context_records/serializers.py b/archaeological_context_records/serializers.py
index 98f7d2226..6035f0182 100644
--- a/archaeological_context_records/serializers.py
+++ b/archaeological_context_records/serializers.py
@@ -13,24 +13,40 @@ def cr_serialization(archive=False, return_empty_types=False,
                             archive_name=None, operation_queryset=None,
                             site_queryset=None, cr_queryset=None):
     result_queryset = {}
-    """
     if operation_queryset:
+        operation_ids = operation_queryset.values_list("id", flat=True)
         result_queryset = {
-            models.ContextRecord.__name__: operation_queryset,
-            models.ArchaeologicalSite.__name__:
-                models.ArchaeologicalSite.objects.filter(
-                    operations__id__in=operation_queryset.values_list(
-                        "id", flat=True)
-                    )
+            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,
+                )
         }
     elif site_queryset:
+        site_ids = site_queryset.values_list("id", flat=True)
         result_queryset = {
-            models.ArchaeologicalSite.__name__: site_queryset,
-            models.Operation.__name__: models.Operation.objects.filter(
-                archaeological_sites__id__in=site_queryset.values_list(
-                    "id", flat=True))
+            models.ContextRecord.__name__:
+                models.ContextRecord.objects.filter(
+                    archaeological_site_id__in=site_ids),
+            models.RecordRelations.__name__:
+                models.RecordRelations.objects.filter(
+                    left_record__archaeological_site_id__in=site_ids,
+                    right_record__archaeological_site_id__in=site_ids,
+                )
+        }
+    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,
+                )
         }
-    """
 
     result = generic_get_results(CR_MODEL_LIST, "context_records",
                                  result_queryset=result_queryset)