(re)use cache helpers + split into several functions

worker_init_elements/worker.py

@@ -6,6 +6,14 @@ from collections import OrderedDict
 from logging import Logger, getLogger
 from time import sleep
 
+from arkindex_worker.cache import (
+    CachedElement,
+    CachedImage,
+    create_tables,
+    create_version_table,
+    init_cache_db,
+)
+from arkindex_worker.models import Element as BaseElement
 from arkindex_worker.worker.base import BaseWorker
 
 logger: Logger = getLogger(__name__)
@@ -16,33 +24,6 @@ INIT_PAGE_SIZE = 500
 PENDING_STATE = "pending"
 ERROR_STATE = "error"
 
-SQL_TABLES = """
-CREATE TABLE version AS SELECT 3 AS version;
-CREATE TABLE images (
-    id VARCHAR(32) PRIMARY KEY,
-    width INTEGER NOT NULL,
-    height INTEGER NOT NULL,
-    url TEXT NOT NULL
-);
-CREATE TABLE elements (
-    id VARCHAR(32) PRIMARY KEY,
-    parent_id VARCHAR(32),
-    type TEXT NOT NULL,
-    image_id VARCHAR(32),
-    polygon TEXT,
-    rotation_angle INTEGER NOT NULL DEFAULT 0,
-    mirrored INTEGER NOT NULL DEFAULT 0,
-    initial BOOLEAN DEFAULT 0 NOT NULL,
-    worker_version_id VARCHAR(32),
-    worker_run_id VARCHAR(32),
-    confidence REAL,
-    FOREIGN KEY(image_id) REFERENCES images(id),
-    CHECK (rotation_angle >= 0 AND rotation_angle <= 359),
-    CHECK (mirrored = 0 OR mirrored = 1),
-    CHECK (confidence IS NULL OR (confidence >= 0 AND confidence <= 1))
-);
-"""
-
 
 def split_chunks(items, n):
     """
@@ -53,6 +34,21 @@ def split_chunks(items, n):
         yield items[i::n]
 
 
+class Element(BaseElement):
+    """
+    Some attributes in the `ListProcessElements` response conflict with the `BaseElement` property/function.
+    Override them to access the API response directly.
+    """
+
+    @property
+    def polygon(self) -> list[float]:
+        return self["polygon"]
+
+    @property
+    def image_url(self) -> list[float]:
+        return self["image_url"]
+
+
 class InitElementWorker(BaseWorker):
     def configure(self) -> None:
         # CLI args are stored on the instance so that implementations can access them
@@ -73,87 +69,114 @@ class InitElementWorker(BaseWorker):
         self.use_cache = self.config["use_cache"]
         self.api_client.sleep_duration = self.config["sleep"]
 
-        assert self.worker_run_id, "Missing ARKINDEX_WORKER_RUN_ID environment variable, cannot list the elements to process"
-
+        assert self.worker_run_id, "Missing ARKINDEX_WORKER_RUN_ID environment variable, cannot retrieve process information"
         self.process = self.request("RetrieveWorkerRun", id=self.worker_run_id)[
             "process"
         ]
 
-    def dump_json(self, elements, filename="elements.json"):
-        file_path = self.work_dir / filename
-        assert len(elements), f"No elements could be written to {file_path}"
+    def dump_json(
+        self, elements: list[Element], filename: str = "elements.json"
+    ) -> None:
+        path = self.work_dir / filename
+        assert not path.exists(), f"JSON at {path} already exists"
 
-        with file_path.open("w") as f:
-            json.dump(elements, f, indent=4)
+        path.write_text(json.dumps(elements, indent=4))
 
-    def dump_sqlite(self, elements, filename="db.sqlite"):
+    def dump_sqlite(self, elements: list[Element], filename: str = "db.sqlite") -> None:
         if not self.use_cache:
             return
 
         path = self.work_dir / filename
         assert not path.exists(), f"Database at {path} already exists"
+
         db = sqlite3.connect(str(path))
-        db.executescript(SQL_TABLES)
 
-        # Set of unique images found in the elements
-        image_rows = {
-            (
-                uuid.UUID(element["image_id"]).hex,
-                element["image_width"],
-                element["image_height"],
-                element["image_url"],
-            )
-            for element in elements
-            if element.get("image_id") is not None
-        }
-        db.executemany(
-            "INSERT INTO images (id, width, height, url) VALUES (?,?,?,?)", image_rows
-        )
+        init_cache_db(path)
+        create_version_table()
+        create_tables()
 
-        element_rows = [
-            (
-                uuid.UUID(element["id"]).hex,
-                self.type_slugs[element["type_id"]],
-                uuid.UUID(element["image_id"]).hex if element.get("image_id") else None,
-                json.dumps(element["polygon"]) if element.get("polygon") else None,
-                element.get("rotation_angle") or 0,
-                element.get("mirrored") or 0,
-                element.get("confidence"),
-            )
-            for element in elements
-        ]
-        db.executemany(
-            "INSERT INTO elements (id, type, image_id, polygon, rotation_angle, mirrored, confidence, initial) VALUES (?,?,?,?,?,?,?,1)",
-            element_rows,
-        )
+        # Set of unique images found in the elements
+        CachedImage.insert_many(
+            [
+                {
+                    "id": uuid.UUID(element.image_id).hex,
+                    "width": element.image_width,
+                    "height": element.image_height,
+                    "url": element.image_url,
+                }
+                for element in elements
+                if element.get("image_id") is not None
+            ]
+        ).on_conflict_ignore(ignore=True).execute()
+
+        # Fastest way to INSERT multiple rows.
+        CachedElement.insert_many(
+            [
+                {
+                    "id": uuid.UUID(element.id).hex,
+                    "type": element.type,
+                    "image_id": uuid.UUID(element.image_id).hex
+                    if element.get("image_id")
+                    else None,
+                    "polygon": (element.polygon if element.get("polygon") else None),
+                    "rotation_angle": element.get("rotation_angle") or 0,
+                    "mirrored": element.get("mirrored") or False,
+                    "confidence": element.get("confidence"),
+                    "initial": True,
+                }
+                for element in elements
+            ]
+        ).execute()
 
-        db.commit()
         db.close()
 
-    def check_worker_activity(self):
-        # Check if workers activity associated to this process is in a pending state
-        process = self.request("RetrieveProcess", id=self.process["id"])
-        if process.get("activity_state") == ERROR_STATE:
-            logger.error(
-                "Workers activity could not be initialized. Please report this incident to an instance administrator."
+    def dump_chunks(self, elements: list[Element]) -> None:
+        assert (
+            len(elements) >= self.chunks_number
+        ), f"Too few elements have been retrieved to distribute workflow among {self.chunks_number} branches"
+
+        for index, chunk_elts in enumerate(
+            split_chunks(elements, self.chunks_number),
+            start=1,
+        ):
+            self.dump_json(
+                elements=[
+                    {
+                        "id": element.id,
+                        "type": element.type,
+                    }
+                    for element in chunk_elts
+                ],
+                **(
+                    {"filename": f"elements_chunk_{index}.json"}
+                    if self.chunks_number > 1
+                    else {}
+                ),
+            )
+            self.dump_sqlite(
+                elements=chunk_elts,
+                **(
+                    {"filename": f"db_{index}.sqlite"} if self.chunks_number > 1 else {}
+                ),
             )
-            sys.exit(1)
-        return process.get("activity_state") != PENDING_STATE
 
-    def run(self):
-        assert self.worker_run_id, "Missing ARKINDEX_WORKER_RUN_ID environment variable, cannot list the elements to process"
+        logger.info(
+            f"Added {len(elements)} element{'s'[:len(elements) > 1]} to workflow configuration"
+        )
 
+    def list_process_elements(self) -> list[Element]:
         assert self.process.get(
             "corpus"
         ), "init_elements only supports processes on corpora."
 
         corpus = self.request("RetrieveCorpus", id=self.process["corpus"])
-        self.type_slugs = {
+        type_slugs = {
             element_type["id"]: element_type["slug"] for element_type in corpus["types"]
         }
 
         elements = list(
-            self.api_client.paginate(
+            Element(**element, type=type_slugs[element["type_id"]])
+            for element in self.api_client.paginate(
                 "ListProcessElements",
                 id=self.process["id"],
                 with_image=self.use_cache,
@@ -178,42 +201,19 @@ class InitElementWorker(BaseWorker):
             logger.error("No elements found, aborting workflow.")
             sys.exit(1)
 
-        assert (
-            len(unique_elements) >= self.chunks_number
-        ), f"Too few elements have been retrieved to distribute workflow among {self.chunks_number} branches"
+        return list(unique_elements.values())
 
-        if self.chunks_number == 1:
-            self.dump_json(
-                elements=[
-                    {"id": element_id, "type": self.type_slugs[element["type_id"]]}
-                    for element_id, element in unique_elements.items()
-                ]
+    def check_worker_activity(self):
+        # Check if workers activity associated to this process is in a pending state
+        process = self.request("RetrieveProcess", id=self.process["id"])
+        if process.get("activity_state") == ERROR_STATE:
+            logger.error(
+                "Workers activity could not be initialized. Please report this incident to an instance administrator."
             )
-            self.dump_sqlite(unique_elements.values())
-        else:
-            for index, chunk_elts in enumerate(
-                split_chunks(list(unique_elements.values()), self.chunks_number),
-                start=1,
-            ):
-                self.dump_json(
-                    elements=[
-                        {
-                            "id": element["id"],
-                            "type": self.type_slugs[element["type_id"]],
-                        }
-                        for element in chunk_elts
-                    ],
-                    filename=f"elements_chunk_{index}.json",
-                )
-                self.dump_sqlite(
-                    elements=chunk_elts,
-                    filename=f"db_{index}.sqlite",
-                )
-
-        logger.info(
-            f"Added {len(unique_elements)} element{'s'[:len(unique_elements) > 1]} to workflow configuration"
-        )
+            sys.exit(1)
+        return process.get("activity_state") != PENDING_STATE
 
+    def await_worker_activity(self) -> None:
         logger.info("Awaiting workers activity initialization")
         # Await worker activities to be initialized for 0, 2, 4, 8 seconds up to an hour
         timer = 1
@@ -229,6 +229,12 @@ class InitElementWorker(BaseWorker):
                 raise Exception("Worker activity timeout")
             sleep(timer)
 
+    def run(self):
+        elements = self.list_process_elements()
+        self.dump_chunks(elements)
+
+        self.await_worker_activity()
+
 
 def main() -> None:
     InitElementWorker(