Store created elements in a local SQLite database

arkindex_worker/cache.py

+# -*- coding: utf-8 -*-
+import os
+import sqlite3
+
+
+class LocalDB(object):
+    def __init__(self, path):
+        if not os.path.exists(path):
+            open(path, "x").close()
+
+        self.db = sqlite3.connect(path)
+        self.cursor = self.db.cursor()
+
+    def create_elements_table(self):
+        try:
+            self.cursor.execute(
+                """CREATE TABLE elements (
+                id TEXT PRIMARY KEY,
+                parent_id TEXT,
+                name TEXT NOT NULL,
+                type TEXT NOT NULL,
+                polygon TEXT,
+                worker_version_id TEXT
+            )"""
+            )
+        except sqlite3.OperationalError:
+            print("Table 'elements' already exists")
+
+    def insert(self, table, lines):
+        self.cursor.executemany(f"INSERT INTO {table} VALUES (?,?,?,?,?,?)", lines)
+        self.db.commit()

arkindex_worker/worker.py

@@ -16,6 +16,7 @@ from apistar.exceptions import ErrorResponse
 
 from arkindex import ArkindexClient, options_from_env
 from arkindex_worker import logger
+from arkindex_worker.cache import LocalDB
 from arkindex_worker.models import Element
 from arkindex_worker.reporting import Reporter
 
@@ -451,6 +452,22 @@ class ElementsWorker(BaseWorker):
         for element in elements:
             self.report.add_element(parent.id, element["type"])
 
+        # Store elements in local cache
+        cache = LocalDB(f"/data/{os.environ.get('TASK_ID')}/db.sqlite")
+        cache.create_elements_table()
+        to_insert = [
+            (
+                created_ids[idx],
+                parent.id,
+                element["name"],
+                element["type"],
+                json.dumps(element["polygon"]),
+                self.worker_version_id,
+            )
+            for idx, element in enumerate(elements)
+        ]
+        cache.insert("elements", to_insert)
+
         return created_ids
 
     def create_transcription(self, element, text, type=None, score=None):