1 files changed, 90 insertions, 7 deletions

diff --git a/ishtar_common/serializers_utils.py b/ishtar_common/serializers_utils.py
index ada0a62cf..a4494d9c9 100644
--- a/ishtar_common/serializers_utils.py
+++ b/ishtar_common/serializers_utils.py
@@ -4,9 +4,12 @@ import datetime
 import json
 import importlib
 import os
+from rest_framework.renderers import JSONRenderer
+from rest_framework import serializers
 import tempfile
 from zipfile import ZipFile
 
+from django.contrib.contenttypes.models import ContentType
 from django.contrib.sites.models import Site
 from django.core.serializers import serialize
 from django.db.models import Q
@@ -113,6 +116,70 @@ def archive_serialization(
     return archive_name
 
 
+class ItemKeySerializer(serializers.ModelSerializer):
+    app = serializers.SerializerMethodField()
+    model = serializers.SerializerMethodField()
+    content_slug = serializers.SerializerMethodField()
+    importer_type = serializers.SerializerMethodField()
+
+    class Meta:
+        model = models.ItemKey
+        fields = ['key', 'app', 'model', 'content_slug', 'importer_type']
+
+    def get_content_slug(self, obj):
+        return obj.content_slug
+
+    def get_app(self, obj):
+        return obj.content_type.app_label
+
+    def get_model(self, obj):
+        return obj.content_type.model
+
+    def get_importer_type(self, obj):
+        return obj.importer_type.slug if obj.importer_type else ""
+
+    def save(self, data):
+        try:
+            ct = ContentType.objects.get(
+                app_label=data["app"],
+                model=data["model"],
+            )
+        except ContentType.DoesNotExist:
+            return
+        model = ct.model_class()
+        if hasattr(model, "txt_idx"):
+            slug = "txt_idx"
+        elif hasattr(model, "slug"):
+            slug = "slug"
+        else:
+            slug = "pk"
+        try:
+            value = model.objects.get(**{slug: data["content_slug"]})
+        except model.DoesNotExist:
+            return
+        importer_type = None
+        if data["importer_type"]:
+            try:
+                importer_type = models.ImporterType.objects.get(
+                    slug=data["importer_type"]
+                )
+            except models.ImporterType.DoesNotExist:
+                return
+        obj, created = models.ItemKey.objects.get_or_create(
+            key=data["key"],
+            content_type=ct,
+            importer_type=importer_type,
+            ishtar_import=None,
+            user=None,
+            group=None,
+            defaults={"object_id": value.pk}
+        )
+        if not created:
+            obj.object_id = value.pk
+            obj.save()
+        return obj
+
+
 GENERIC_QUERYSET_FILTER = {
     "JsonDataSection": {"JsonDataField": "json_data_field__pk__in"},
     "Regexp": {
@@ -170,6 +237,8 @@ GENERIC_QUERYSET_FILTER = {
     }
 }
 
+CUSTOM_SERIALIZERS = {"ishtar_common.itemkey": ItemKeySerializer}
+
 
 def generic_get_results(
     model_list, dirname, no_geo=True, result_queryset=None, serialization_include=None
@@ -212,13 +281,26 @@ def generic_get_results(
             q = q.filter(**{recursion + "__isnull": True})
 
         key = (dirname, model_name)
-        result[key] = serialize(
-            "json",
-            q.distinct().all(),
-            indent=2,
-            use_natural_foreign_keys=True,
-            use_natural_primary_keys=True,
-        )
+        model_name = f"{model._meta.app_label}.{model._meta.model_name}"
+        is_custom = model_name in CUSTOM_SERIALIZERS
+        if is_custom:
+            current_serializer = CUSTOM_SERIALIZERS[model_name]
+            result[key] = JSONRenderer().render(
+                [{"model": model_name,
+                  "fields": current_serializer(item).data}
+                 for item in q.distinct().all()]
+            )
+        else:
+            result[key] = serialize(
+                "json",
+                q.distinct().all(),
+                indent=2,
+                use_natural_foreign_keys=True,
+                use_natural_primary_keys=True,
+            )
+        if recursion and is_custom:
+            # TODO
+            raise NotImplementedError("Recursion not managed for this custom serializer")
 
         if recursion:
             serialized = [item["id"] for item in q.values("id").all()]
@@ -278,6 +360,7 @@ def generic_get_results(
             for k in serialization_include[model.__name__]:
                 if k in excluded_fields:
                     excluded_fields.pop(excluded_fields.index(k))
+
         if excluded_fields:
             new_result = json.loads(result[key])
             for idx in range(len(new_result)):