Validate SQLAlchemy 2.1 server compatibility

changelog.d/sqlalchemy-21-compatibility.fixed.md

@@ -0,0 +1 @@
+- Allowed SQLAlchemy 2.1 for server installations while excluding the broken 2.1.0 source distribution.

pyproject.toml

@@ -60,7 +60,7 @@ examples = [
 ]
 server = [
     "fastapi>=0.139",
-    "sqlalchemy[asyncio]>=2.0.51,<2.1",
+    "sqlalchemy[asyncio]>=2.0.51,!=2.1.0",
     "alembic>=1.18",
     "uvicorn>=0.51",
     "pydantic-settings>=2.14",

scripts/smoke_plugin_artifacts.py

@@ -403,6 +403,7 @@ def _smoke_candidate_wheels(
             cwd=root,
         ),
     )
+    _probe_sqlalchemy(python, root, environment)
     probe = repository / "scripts" / "probe_plugin_artifacts.py"
     command: list[str | Path] = [python, probe]
     for requirement in requirements:
@@ -415,6 +416,28 @@ def _smoke_candidate_wheels(
     )
 
 
+def _probe_sqlalchemy(python: Path, root: Path, environment: dict[str, str]) -> None:
+    """Exercise the installed SQLAlchemy async runtime and report its versions."""
+    result = _expect_success(
+        "probe installed SQLAlchemy async runtime",
+        _run(
+            [
+                python,
+                "-c",
+                "import asyncio; "
+                "from importlib.metadata import version; "
+                "from sqlalchemy.util.concurrency import greenlet_spawn; "
+                "assert asyncio.run(greenlet_spawn(lambda: 42)) == 42; "
+                'print(f\'SQLAlchemy=={version("sqlalchemy")} '
+                'greenlet=={version("greenlet")}\')',
+            ],
+            environment=environment,
+            cwd=root,
+        ),
+    )
+    print(result.stdout.strip())
+
+
 def main() -> int:
     """Build plugin wheels and validate their installed contracts."""
     parser = argparse.ArgumentParser()

src/kitaru/server/adapters/db/pagination.py

@@ -276,7 +276,7 @@ def _translate_query_timeout(error: DBAPIError) -> None:
 
 async def paginate(
     session: AsyncSession,
-    statement: Select[tuple[RowT]],
+    statement: Select[RowT],
     list_filter: ListFilter,
     order: PageOrder,
 ) -> tuple[Sequence[RowT], str | None]:

src/kitaru/server/adapters/db/repositories/plugin_repository.py

@@ -243,7 +243,7 @@ async def create_version(
 
     async def _get_version_row(
         self,
-        statement: Select[tuple[PluginVersionORM]],
+        statement: Select[PluginVersionORM],
         plugin_id: uuid.UUID,
         version: int,
     ) -> PluginVersionORM:

src/kitaru/server/adapters/db/repositories/replay_repository.py

@@ -16,7 +16,7 @@
 import uuid
 from collections.abc import Callable, Mapping, Sequence
 
-from sqlalchemy import func, select
+from sqlalchemy import NotNullable, func, select
 
 from kitaru.api_models.v1.replay import ReplayStatus
 from kitaru.server.adapters.db.filtering import (
@@ -315,7 +315,9 @@ async def count_by_status_many(
             return {}
         statement = (
             select(
-                ReplayORM.experiment_run_id, ReplayORM.status, func.count(ReplayORM.id)
+                NotNullable(ReplayORM.experiment_run_id),
+                ReplayORM.status,
+                func.count(ReplayORM.id),
             )
             .where(ReplayORM.experiment_run_id.in_(experiment_run_ids))
             .group_by(ReplayORM.experiment_run_id, ReplayORM.status)

src/kitaru/server/adapters/db/repositories/session_node_repository.py

@@ -203,7 +203,7 @@ async def exists_in_session(
         return bool(await self._session.scalar(statement))
 
     async def _latest_match(
-        self, statement: Select[tuple[SessionNodeORM]]
+        self, statement: Select[SessionNodeORM]
     ) -> SessionNode | None:
         """Run a cache-key search statement and return its newest match.
 

src/kitaru/server/adapters/db/repositories/task_repository.py

@@ -448,7 +448,7 @@ async def stamp_heartbeats(
         )
         rows = (await self._session.execute(statement)).all()
         await self._session.flush()
-        stamped = {
+        stamped: dict[uuid.UUID, datetime | None] = {
             task_id: cancel_requested_at for task_id, cancel_requested_at in rows
         }
         return stamped, candidate_ids - stamped.keys()

tests/scripts/test_smoke_plugin_artifacts.py

@@ -1,9 +1,12 @@
+import subprocess
 from pathlib import Path
 from zipfile import ZipFile
 
 import pytest
 from scripts.smoke_plugin_artifacts import SmokeFailure, _validate_wheel_metadata
 
+from scripts import smoke_plugin_artifacts
+
 VALID_METADATA = """Metadata-Version: 2.4
 Name: kitaru-example
 Version: 1.0.0
@@ -65,3 +68,39 @@ def test_wheel_metadata_rejects_an_invalid_project_url(tmp_path: Path) -> None:
 
     with pytest.raises(SmokeFailure, match="invalid Project-URL"):
         _validate_wheel_metadata(wheel, "kitaru-example", "1.0.0")
+
+
+def test_sqlalchemy_probe_reports_installed_versions(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
+) -> None:
+    def run(
+        command: list[str | Path], *, environment: dict[str, str], cwd: Path
+    ) -> subprocess.CompletedProcess[str]:
+        return subprocess.CompletedProcess(
+            command, 0, stdout="SQLAlchemy==2.1.1 greenlet==3.5.3\n", stderr=""
+        )
+
+    monkeypatch.setattr(smoke_plugin_artifacts, "_run", run)
+
+    smoke_plugin_artifacts._probe_sqlalchemy(tmp_path / "python", tmp_path, {})
+
+    assert capsys.readouterr().out == "SQLAlchemy==2.1.1 greenlet==3.5.3\n"
+
+
+def test_sqlalchemy_probe_rejects_missing_async_runtime(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    def run(
+        command: list[str | Path], *, environment: dict[str, str], cwd: Path
+    ) -> subprocess.CompletedProcess[str]:
+        return subprocess.CompletedProcess(
+            command,
+            1,
+            stdout="",
+            stderr="ValueError: the greenlet library is required to use this function",
+        )
+
+    monkeypatch.setattr(smoke_plugin_artifacts, "_run", run)
+
+    with pytest.raises(SmokeFailure, match="the greenlet library is required"):
+        smoke_plugin_artifacts._probe_sqlalchemy(tmp_path / "python", tmp_path, {})