New Process Dataset Sets management

closes #1719 (closed)

requires !2256 (merged)

arkindex/documents/management/commands/load_export.py

@@ -37,7 +37,7 @@ from arkindex.process.models import (
     WorkerType,
     WorkerVersion,
 )
-from arkindex.training.models import Dataset, DatasetElement, Model
+from arkindex.training.models import Dataset, DatasetElement, DatasetSet, Model
 from arkindex.users.models import Role, User
 
 EXPORT_VERSION = 8
@@ -320,17 +320,30 @@ class Command(BaseCommand):
             id=row["id"],
             corpus=self.corpus,
             name=row["name"],
-            sets=[r.strip() for r in row["sets"].split(",")],
             creator=self.user,
             description="Imported dataset",
         )]
 
+    def convert_dataset_sets(self, row):
+        return [
+            DatasetSet(
+                name=set_name.strip(),
+                dataset_id=row["id"]
+            )
+            for set_name in row["sets"].split(",")
+        ]
+
+    def map_dataset_sets(self):
+        return {
+            (str(set.dataset_id), set.name): set.id
+            for set in DatasetSet.objects.filter(dataset__corpus=self.corpus)
+        }
+
     def convert_dataset_elements(self, row):
         return [DatasetElement(
             id=row["id"],
             element_id=row["element_id"],
-            dataset_id=row["dataset_id"],
-            set=row["set_name"],
+            set_id=self.dataset_sets_map[(row["dataset_id"], row["set_name"])]
         )]
 
     def bulk_create_objects(self, ModelClass, convert_method, sql_query, ignore_conflicts=True):
@@ -603,6 +616,12 @@ class Command(BaseCommand):
             # Create datasets
             self.bulk_create_objects(Dataset, self.convert_datasets, SQL_DATASET_QUERY)
 
+            # Create dataset sets
+            self.bulk_create_objects(DatasetSet, self.convert_dataset_sets, SQL_DATASET_QUERY)
+
+            # Create dataset sets mapping
+            self.dataset_sets_map = self.map_dataset_sets()
+
             # Create dataset elements
             self.bulk_create_objects(DatasetElement, self.convert_dataset_elements, SQL_ELEMENT_DATASET_QUERY)