Log Warning if connecting to DB with older lightly-studio version

CHANGELOG.md

@@ -9,6 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added a warning when connecting to a database with a different LightlyStudio version.
+
+
 ### Changed
 
 ### Deprecated

lightly_studio/src/lightly_studio/api/db_tables.py

@@ -16,6 +16,9 @@
 from lightly_studio.models.dataset import (
     DatasetTable,  # noqa: F401, required for SQLModel to work properly
 )
+from lightly_studio.models.database_version import (
+    DatabaseVersionTable,  # noqa: F401, required for SQLModel to work properly
+)
 from lightly_studio.models.embedding_model import (
     EmbeddingModelTable,  # noqa: F401, required for SQLModel to work properly
 )

lightly_studio/src/lightly_studio/database_version_validation.py

@@ -0,0 +1,71 @@
+"""Validation of LightlyStudio database version metadata."""
+
+from __future__ import annotations
+
+import logging
+from importlib import metadata
+
+from sqlalchemy import inspect
+from sqlalchemy.engine import Engine
+from sqlmodel import Session, SQLModel, select
+
+from lightly_studio.models.database_version import DatabaseVersionTable
+
+
+def validate_or_initialize_database_version(engine: Engine) -> None:
+    """Validate database schema version or initialize it for fresh databases.
+
+    These cases are handled:
+    1. Fresh database with no managed tables: Insert the expected version.
+    2. Existing database without a `database_version` table: Warn.
+    3. Existing database with a `database_version` table:
+         a. With 0 rows: Raise an error, this should not happen.
+         b. With >1 rows: Raise an error, this should not happen.
+         c. With 1 row and a different version: Warn about incompatible version.
+         d. With 1 row and the expected version: Do nothing.
+
+    Note that only in case 1 the version row is written. In all other cases the
+    version metadata is left unchanged.
+
+    The DB queries are executed here and not in the resolvers, as they are only used here.
+    """
+    existing_table_names = set(inspect(engine).get_table_names())
+    SQLModel.metadata.create_all(bind=engine)
+    if len(existing_table_names) == 0:
+        # Case 1: Fresh database, insert the expected version.
+        with Session(engine) as session:
+            session.add(DatabaseVersionTable(version=metadata.version("lightly-studio")))
+            session.commit()
+        return
+
+    if DatabaseVersionTable.__tablename__ not in existing_table_names:
+        # Case 2: Existing database without a `database_version` table: Warn.
+        logging.warning(
+            f"Incompatible database schema version. Expected version "
+            f"'{metadata.version('lightly-studio')}', but found missing version metadata."
+        )
+        return
+
+    with Session(engine) as session:
+        version_rows = session.exec(select(DatabaseVersionTable)).all()
+        if len(version_rows) == 0:
+            # Case 3a: Existing database with a `database_version` table but 0 rows.
+            raise RuntimeError(
+                "Expected exactly one row in 'database_version' when the table already "
+                "exists, got 0."
+            )
+        if len(version_rows) > 1:
+            # Case 3b: Existing database with a `database_version` table but >1 rows.
+            raise RuntimeError(
+                f"Expected at most one row in 'database_version' when the table already "
+                f"exists, got {len(version_rows)}."
+            )
+        version_row = version_rows[0]
+        expected_version = metadata.version("lightly-studio")
+        if version_row.version != expected_version:
+            # Case 3c: Existing database with a different schema version.
+            logging.warning(
+                f"Incompatible database schema version. Expected version "
+                f"'{expected_version}', got '{version_row.version}'."
+            )
+        # Case 3d: Existing database with the expected schema version.

lightly_studio/src/lightly_studio/db_manager.py

@@ -28,6 +28,7 @@
 from sqlmodel import Session, SQLModel, create_engine
 
 import lightly_studio.api.db_tables  # noqa: F401, required for SQLModel to work properly
+from lightly_studio.database_version_validation import validate_or_initialize_database_version
 from lightly_studio.dataset.env import LIGHTLY_STUDIO_DATABASE_URL
 
 
@@ -111,7 +112,7 @@ def __init__(
             SQLModel.metadata.drop_all(bind=self._engine)
             logging.info("Dropped all tables in PostgreSQL database.")
 
-        SQLModel.metadata.create_all(self._engine)
+        validate_or_initialize_database_version(engine=self._engine)
 
     @contextmanager
     def session(self) -> Generator[Session, None, None]:

lightly_studio/src/lightly_studio/models/database_version.py

@@ -0,0 +1,13 @@
+"""Model for tracking the LightlyStudio database schema version."""
+
+from __future__ import annotations
+
+from sqlmodel import Field, SQLModel
+
+
+class DatabaseVersionTable(SQLModel, table=True):
+    """Stores the schema version expected by the running LightlyStudio package."""
+
+    __tablename__ = "database_version"
+
+    version: str = Field(primary_key=True)

lightly_studio/src/lightly_studio/resolvers/dataset_resolver/table_coverage_utils.py

@@ -13,7 +13,8 @@
 # Tables not relevant for collection operations:
 # - setting (application-level, not collection-specific)
 # - two_dim_embeddings (cached projections, regenerated as needed)
-_EXCLUDED_TABLES_COUNT = 2
+# - database_version (application-level schema version metadata, not collection-specific)
+_EXCLUDED_TABLES_COUNT = 3
 
 _TOTAL_TABLES_COUNT = _HANDLED_TABLES_COUNT + _EXCLUDED_TABLES_COUNT
 

lightly_studio/tests/conftest.py

@@ -74,6 +74,12 @@ def _use_postgres(request: pytest.FixtureRequest) -> bool:
     return bool(request.config.getoption("--postgres"))
 
 
+@pytest.fixture
+def patch_engine_singleton(mocker: MockerFixture) -> None:
+    """Patch the db_manager engine singleton to simulate a fresh module state."""
+    mocker.patch.object(db_manager, "_engine", new=None)
+
+
 @pytest.fixture(scope="session")
 def postgres_url(_use_postgres: bool) -> Generator[str | None, None, None]:
     """Start a Postgres container and yield its URL, or None for DuckDB."""

lightly_studio/tests/test_database_version_validation.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+from importlib import metadata
+from pathlib import Path
+
+import pytest
+from sqlmodel import select
+
+from lightly_studio import db_manager
+from lightly_studio.models.database_version import DatabaseVersionTable
+
+PACKAGE_VERSION = metadata.version("lightly-studio")
+
+
+def _create_database(db_url: str, version: str | None) -> None:
+    db_manager.connect(db_url=db_url)
+    with db_manager.session() as session:
+        db_version = session.exec(select(DatabaseVersionTable)).first()
+        assert db_version is not None
+        if version != db_version.version:
+            session.delete(db_version)
+            if version is not None:
+                session.add(DatabaseVersionTable(version=version))
+    db_manager.close()
+
+
+def _drop_database_version_table(db_url: str) -> None:
+    _create_database(db_url=db_url, version=PACKAGE_VERSION)
+    db_manager.connect(db_url=db_url)
+    with db_manager.session() as session:
+        session.connection().exec_driver_sql("DROP TABLE database_version")
+    db_manager.close()
+
+
+@pytest.mark.usefixtures("patch_engine_singleton")
+class TestValidateOrInitializeDatabaseVersion:
+    """Tests for `_validate_or_initialize_database_version`.
+
+    Covered cases:
+    1. Fresh database with no managed tables: insert the expected version.
+    2. Existing database without a `database_version` table: warn.
+    3a. Existing database with a `database_version` table but 0 rows: raise.
+    3b. Existing database with a `database_version` table but >1 rows: raise.
+    3c. Existing database with 1 row and a different version: warn.
+    3d. Existing database with 1 row and the expected version: do nothing.
+    """
+
+    def test_connect__case_1__fresh_database_initializes_expected_version(
+        self,
+        tmp_path: Path,
+        caplog: pytest.LogCaptureFixture,
+    ) -> None:
+        db_url = f"duckdb:///{tmp_path / 'no_db.db'}"
+        caplog.clear()
+        with caplog.at_level("WARNING"):
+            db_manager.connect(db_url=db_url, must_exist=False)
+        assert caplog.messages == []
+        with db_manager.session() as session:
+            assert session.exec(select(DatabaseVersionTable.version)).first() == PACKAGE_VERSION
+        db_manager.close()
+
+    def test_connect__case_2__existing_database_without_database_version_table_warns(
+        self,
+        tmp_path: Path,
+        caplog: pytest.LogCaptureFixture,
+    ) -> None:
+        db_url = f"duckdb:///{tmp_path / 'db_without_version.db'}"
+        _drop_database_version_table(db_url=db_url)
+        expected_warning = (
+            f"Incompatible database schema version. Expected version "
+            f"'{PACKAGE_VERSION}', but found missing version metadata."
+        )
+        caplog.clear()
+        with caplog.at_level("WARNING"):
+            db_manager.connect(db_url=db_url, must_exist=True)
+        assert caplog.messages == [expected_warning]
+        with db_manager.session() as session:
+            assert session.exec(select(DatabaseVersionTable.version)).first() is None
+        db_manager.close()
+
+    def test_connect__case_3a__database_version_table_without_rows_raises(
+        self,
+        tmp_path: Path,
+    ) -> None:
+        db_url = f"duckdb:///{tmp_path / 'empty_database_version.db'}"
+        _create_database(db_url=db_url, version=None)
+
+        with pytest.raises(
+            RuntimeError,
+            match=(
+                r"Expected exactly one row in 'database_version' when the table already "
+                r"exists, got 0."
+            ),
+        ):
+            db_manager.connect(db_url=db_url, must_exist=True)
+
+    def test_connect__case_3b__database_version_table_with_multiple_rows_raises(
+        self,
+        tmp_path: Path,
+    ) -> None:
+        db_url = f"duckdb:///{tmp_path / 'db_with_multiple_versions.db'}"
+        _create_database(db_url=db_url, version=PACKAGE_VERSION)
+
+        db_manager.connect(db_url=db_url)
+        with db_manager.session() as session:
+            session.add(DatabaseVersionTable(version="0.0.0"))
+        db_manager.close()
+
+        with pytest.raises(
+            RuntimeError,
+            match=(
+                r"Expected at most one row in 'database_version' when the table already "
+                r"exists, got 2."
+            ),
+        ):
+            db_manager.connect(db_url=db_url, must_exist=True)
+
+    def test_connect__case_3c__database_version_mismatch_warns(
+        self,
+        tmp_path: Path,
+        caplog: pytest.LogCaptureFixture,
+    ) -> None:
+        db_url = f"duckdb:///{tmp_path / 'db_with_other_version.db'}"
+        _create_database(db_url=db_url, version="0.0.0")
+        expected_warning = (
+            f"Incompatible database schema version. Expected version '{PACKAGE_VERSION}', "
+            "got '0.0.0'."
+        )
+        caplog.clear()
+        with caplog.at_level("WARNING"):
+            db_manager.connect(db_url=db_url, must_exist=True)
+        assert caplog.messages == [expected_warning]
+        with db_manager.session() as session:
+            assert session.exec(select(DatabaseVersionTable.version)).first() == "0.0.0"
+        db_manager.close()
+
+    def test_connect__case_3d__database_version_match_does_nothing(
+        self,
+        tmp_path: Path,
+        caplog: pytest.LogCaptureFixture,
+    ) -> None:
+        db_url = f"duckdb:///{tmp_path / 'db_with_same_version.db'}"
+        _create_database(db_url=db_url, version=PACKAGE_VERSION)
+        caplog.clear()
+        with caplog.at_level("WARNING"):
+            db_manager.connect(db_url=db_url, must_exist=True)
+        assert caplog.messages == []
+        with db_manager.session() as session:
+            assert session.exec(select(DatabaseVersionTable.version)).first() == PACKAGE_VERSION
+        db_manager.close()

lightly_studio/tests/test_db_manager.py

@@ -27,15 +27,6 @@
 )
 
 
-@pytest.fixture
-def patch_engine_singleton(mocker: MockerFixture) -> None:
-    """Patch the _engine variable to None before each test.
-
-    This simulates a fresh load of the module.
-    """
-    mocker.patch.object(db_manager, "_engine", new=None)
-
-
 def test_get_engine__default(
     mocker: MockerFixture,
     patch_engine_singleton: None,  # noqa ARG001