Merge pull request #16 from ggozad/feat/db_upgrades

Store haiku.rag version & allow for future db upgrades

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "haiku.rag"
-version = "0.3.3"
+version = "0.3.4"
 description = "Retrieval Augmented Generation (RAG) with SQLite"
 authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
 license = { text = "MIT" }

src/haiku/rag/store/engine.py

@@ -1,11 +1,17 @@
 import sqlite3
 import struct
+from importlib import metadata
 from pathlib import Path
 from typing import Literal
 
 import sqlite_vec
+from packaging.version import parse
+from rich.console import Console
 
+from haiku.rag.config import Config
 from haiku.rag.embeddings import get_embedder
+from haiku.rag.store.upgrades import upgrades
+from haiku.rag.utils import int_to_semantic_version, semantic_version_to_int
 
 
 class Store:
@@ -13,7 +19,7 @@ class Store:
         self, db_path: Path | Literal[":memory:"], skip_validation: bool = False
     ):
         self.db_path: Path | Literal[":memory:"] = db_path
-        self._connection = self.create_db()
+        self.create_or_update_db()
 
         # Validate config compatibility after connection is established
         if not skip_validation:
@@ -21,12 +27,39 @@ class Store:
 
             settings_repo = SettingsRepository(self)
             settings_repo.validate_config_compatibility()
+        current_version = metadata.version("haiku.rag")
+        self.set_user_version(current_version)
 
-    def create_db(self) -> sqlite3.Connection:
+    def create_or_update_db(self):
         """Create the database and tables with sqlite-vec support for embeddings."""
+        current_version = metadata.version("haiku.rag")
+
         db = sqlite3.connect(self.db_path)
         db.enable_load_extension(True)
         sqlite_vec.load(db)
+        self._connection = db
+        existing_tables = [
+            row[0]
+            for row in db.execute(
+                "SELECT name FROM sqlite_master WHERE type='table';"
+            ).fetchall()
+        ]
+
+        # If we have a db already, perform upgrades and return
+        if self.db_path != ":memory:" and "documents" in existing_tables:
+            # Upgrade database
+            console = Console()
+            db_version = self.get_user_version()
+            for version, steps in upgrades:
+                if parse(current_version) >= parse(version) and parse(version) > parse(
+                    db_version
+                ):
+                    for step in steps:
+                        step(db)
+                        console.print(
+                            f"[green][b]DB Upgrade: [/b]{step.__doc__}[/green]"
+                        )
+            return
 
         # Create documents table
         db.execute("""
@@ -39,7 +72,6 @@ class Store:
                 updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
             )
         """)
-
         # Create chunks table
         db.execute("""
             CREATE TABLE IF NOT EXISTS chunks (
@@ -50,7 +82,6 @@ class Store:
                 FOREIGN KEY (document_id) REFERENCES documents (id) ON DELETE CASCADE
             )
         """)
-
         # Create vector table for chunk embeddings
         embedder = get_embedder()
         db.execute(f"""
@@ -59,7 +90,6 @@ class Store:
                 embedding FLOAT[{embedder._vector_dim}]
             )
         """)
-
         # Create FTS5 table for full-text search
         db.execute("""
             CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
@@ -68,7 +98,6 @@ class Store:
                 content_rowid='id'
             )
         """)
-
         # Create settings table for storing current configuration
         db.execute("""
             CREATE TABLE IF NOT EXISTS settings (
@@ -76,23 +105,35 @@ class Store:
                 settings TEXT NOT NULL DEFAULT '{}'
             )
         """)
-
-        # Create indexes for better performance
-        db.execute(
-            "CREATE INDEX IF NOT EXISTS idx_chunks_document_id ON chunks(document_id)"
-        )
-
         # Save current settings to the new database
-        from haiku.rag.config import Config
-
         settings_json = Config.model_dump_json()
         db.execute(
             "INSERT OR IGNORE INTO settings (id, settings) VALUES (1, ?)",
             (settings_json,),
         )
-
+        # Create indexes for better performance
+        db.execute(
+            "CREATE INDEX IF NOT EXISTS idx_chunks_document_id ON chunks(document_id)"
+        )
         db.commit()
-        return db
+
+    def get_user_version(self) -> str:
+        """Returns the SQLite user version"""
+        if self._connection is None:
+            raise ValueError("Store connection is not available")
+
+        cursor = self._connection.execute("PRAGMA user_version;")
+        version = cursor.fetchone()
+        return int_to_semantic_version(version[0])
+
+    def set_user_version(self, version: str) -> None:
+        """Updates the SQLite user version"""
+        if self._connection is None:
+            raise ValueError("Store connection is not available")
+
+        self._connection.execute(
+            f"PRAGMA user_version = {semantic_version_to_int(version)};"
+        )
 
     def recreate_embeddings_table(self) -> None:
         """Recreate the embeddings table with current vector dimensions."""

src/haiku/rag/store/upgrades/__init__.py

@@ -0,0 +1,3 @@
+from haiku.rag.store.upgrades.v0_3_4 import upgrades as v0_3_4_upgrades
+
+upgrades = v0_3_4_upgrades

src/haiku/rag/store/upgrades/v0_3_4.py

@@ -0,0 +1,26 @@
+from collections.abc import Callable
+from sqlite3 import Connection
+
+from haiku.rag.config import Config
+
+
+def add_settings_table(db: Connection) -> None:
+    """Create settings table for storing current configuration"""
+    db.execute("""
+        CREATE TABLE settings (
+            id INTEGER PRIMARY KEY DEFAULT 1,
+            settings TEXT NOT NULL DEFAULT '{}'
+        )
+    """)
+
+    settings_json = Config.model_dump_json()
+    db.execute(
+        "INSERT INTO settings (id, settings) VALUES (1, ?)",
+        (settings_json,),
+    )
+    db.commit()
+
+
+upgrades: list[tuple[str, list[Callable[[Connection], None]]]] = [
+    ("0.3.4", [add_settings_table])
+]

src/haiku/rag/utils.py

@@ -29,6 +29,37 @@ def get_default_data_dir() -> Path:
     return data_path
 
 
+def semantic_version_to_int(version: str) -> int:
+    """
+    Convert a semantic version string to an integer.
+
+    :param version: Semantic version string
+    :type version: str
+    :return: Integer representation of semantic version
+    :rtype: int
+    """
+    major, minor, patch = version.split(".")
+    major = int(major) << 16
+    minor = int(minor) << 8
+    patch = int(patch)
+    return major + minor + patch
+
+
+def int_to_semantic_version(version: int) -> str:
+    """
+    Convert an integer to a semantic version string.
+
+    :param version: Integer representation of semantic version
+    :type version: int
+    :return: Semantic version string
+    :rtype: str
+    """
+    major = version >> 16
+    minor = (version >> 8) & 255
+    patch = version & 255
+    return f"{major}.{minor}.{patch}"
+
+
 async def is_up_to_date() -> tuple[bool, Version, Version]:
     """
     Checks whether haiku.rag is current.

tests/test_settings.py

@@ -50,7 +50,6 @@ async def test_config_validation_on_db_load():
 
         # Create store and save settings
         store1 = Store(db_path)
-        SettingsRepository(store1)
         store1.close()
 
         # Change config

tests/test_utils.py

@@ -0,0 +1,15 @@
+from haiku.rag.utils import int_to_semantic_version, semantic_version_to_int
+
+
+def test_sqlite_user_version():
+    version = "0.1.5"
+    assert semantic_version_to_int(version) == 261
+    assert int_to_semantic_version(261) == version
+
+    version = "0.0.0"
+    assert semantic_version_to_int(version) == 0
+    assert int_to_semantic_version(0) == version
+
+    version = "255.255.255"
+    assert semantic_version_to_int(version) == 16777215
+    assert int_to_semantic_version(16777215) == version

uv.lock

@@ -816,7 +816,7 @@ wheels = [
 
 [[package]]
 name = "haiku-rag"
-version = "0.3.3"
+version = "0.3.4"
 source = { editable = "." }
 dependencies = [
     { name = "fastmcp" },