Restore set_elements on create responses

arkindex/training/serializers.py

@@ -15,6 +15,7 @@ from arkindex.documents.serializers.elements import ElementListSerializer
 from arkindex.ponos.models import Task
 from arkindex.process.models import Worker
 from arkindex.project.serializer_fields import ArchivedField, DatasetSetsCountField, EnumField
+from arkindex.project.tools import add_as_prefetch
 from arkindex.training.models import (
     Dataset,
     DatasetElement,
@@ -591,14 +592,19 @@ class DatasetSerializer(serializers.ModelSerializer):
 
     @transaction.atomic
     def create(self, validated_data):
-        sets = validated_data.pop("set_names")
+        set_names = validated_data.pop("set_names")
         dataset = Dataset.objects.create(**validated_data)
-        DatasetSet.objects.bulk_create(
+        sets = DatasetSet.objects.bulk_create(
             DatasetSet(
                 name=set_name,
                 dataset_id=dataset.id
-            ) for set_name in sets
+            ) for set_name in set_names
         )
+        # We will output set element counts in the API, but we know there are zero,
+        # so no need to make another query to prefetch the sets and count them
+        for set in sets:
+            set.element_count = 0
+        add_as_prefetch(dataset.sets, sets)
         return dataset
 
     class Meta: