Add Erwan's suggestions

arkindex_worker/cache.py

@@ -13,6 +13,7 @@ from peewee import (
     ForeignKeyField,
     IntegerField,
     Model,
+    OperationalError,
     SqliteDatabase,
     TextField,
     UUIDField,
@@ -217,25 +218,27 @@ def create_tables():
 
 
 def create_version_table():
+    """
+    Creates the Version table in the cache DB.
+    This step must be independent from other tables creation since we only
+    want to create the table and add the one and only Version entry when the
+    cache is created from scratch.
+    """
     db.create_tables([Version])
     Version.create(version=SQL_VERSION)
 
 
 def check_version(cache_path):
-    provided_db = sqlite3.connect(cache_path)
-    provided_db.row_factory = sqlite3.Row
-    cursor = provided_db.cursor()
-
-    try:
-        db_results = cursor.execute(SQL_VERSION_QUERY).fetchall()
-    except sqlite3.OperationalError:
-        db_results = []
-
-    assert (
-        len(db_results) == 1 and db_results[0]["version"] == SQL_VERSION
-    ), f"The SQLite database {cache_path} does not have the correct cache version, it should be {SQL_VERSION}"
-
-    provided_db.close()
+    with SqliteDatabase(cache_path) as provided_db:
+        with provided_db.bind_ctx([Version]):
+            try:
+                version = Version.get().version
+            except OperationalError:
+                version = None
+
+            assert (
+                version == SQL_VERSION
+            ), f"The SQLite database {cache_path} does not have the correct cache version, it should be {SQL_VERSION}"
 
 
 def retrieve_parents_cache_path(parent_ids, data_dir="/data", chunk=None):