Store Classification in local cache during create_classification

Closes #56 (closed)

arkindex_worker/worker/classification.py

@@ -2,9 +2,10 @@
 import os
 
 from apistar.exceptions import ErrorResponse
+from peewee import IntegrityError
 
 from arkindex_worker import logger
-from arkindex_worker.cache import CachedElement
+from arkindex_worker.cache import CachedClassification, CachedElement
 from arkindex_worker.models import Element
 
 
@@ -84,18 +85,36 @@ class ClassificationMixin(object):
             return
 
         try:
-            self.request(
+            created = self.request(
                 "CreateClassification",
                 body={
-                    "element": element.id,
+                    "element": str(element.id),
                     "ml_class": self.get_ml_class_id(None, ml_class),
                     "worker_version": self.worker_version_id,
                     "confidence": confidence,
                     "high_confidence": high_confidence,
                 },
             )
-        except ErrorResponse as e:
 
+            if self.use_cache:
+                # Store classification in local cache
+                try:
+                    to_insert = [
+                        {
+                            "id": created["id"],
+                            "element_id": element.id,
+                            "class_name": ml_class,
+                            "confidence": created["confidence"],
+                            "state": created["state"],
+                            "worker_version_id": self.worker_version_id,
+                        }
+                    ]
+                    CachedClassification.insert_many(to_insert).execute()
+                except IntegrityError as e:
+                    logger.warning(
+                        f"Couldn't save created classification in local cache: {e}"
+                    )
+        except ErrorResponse as e:
             # Detect already existing classification
             if (
                 e.status_code == 400