feat(BA-3802): Migrate to SQLAlchemy 2.0 with comprehensive type safety improvements

.claude/tasks/sqlalchemy-2.0-migration.md

@@ -0,0 +1,163 @@
+# SQLAlchemy 2.0 Migration Task List
+
+## Current Status (2026-01-08)
+- **Total Errors**: 1,374 errors in 175 files
+- **Checked Files**: 3,138 source files
+
+## Completed Work
+
+### Phase 1: Basic Syntax Migration (Completed)
+- [x] Row indexing: `row["column"]` → `row.column` (~360 errors)
+- [x] Select list args: `select([col1, col2])` → `select(col1, col2)` (~260 errors)
+- [x] Insert args: `insert(table, values)` → `insert(table).values(...)` (~10 errors)
+- [x] Result.rowcount: Added `cast(CursorResult, result)` (~36 errors)
+
+### Phase 2: ORM Row 2.0 Migration (Completed - 6 batches)
+- [x] Batch 1: 4 files (commit: 204e4f8ca)
+- [x] Batch 2: endpoint/row.py (commit: 367ee0361)
+- [x] Batch 3: session/row.py + kernel/row.py (commit: f19a2fcf7)
+- [x] Batch 4: 12 files (commit: e2620ab38)
+- [x] Batch 5: 15 files (commit: 21e605789)
+- [x] Batch 6: IDColumn patterns (commit: 99cf04673)
+
+### Phase 3: Infrastructure (Completed)
+- [x] TypeDecorator signature fixes
+- [x] async_sessionmaker migration
+
+---
+
+## Remaining Work
+
+### High Priority - Core Files (by error count)
+
+| File | Errors | Primary Issues |
+|------|--------|----------------|
+| src/ai/backend/manager/registry.py | 105 | `Mapping[str, Any]` attribute access, AgentId type |
+| src/ai/backend/manager/api/vfolder.py | 71 | Row attribute access, type issues |
+| src/ai/backend/manager/repositories/permission_controller/db_source/db_source.py | 66 | ORM query types |
+| src/ai/backend/manager/models/session/row.py | 41 | Column vs runtime value types |
+| src/ai/backend/manager/models/endpoint/row.py | 37 | Column vs runtime value types |
+| src/ai/backend/manager/api/stream.py | 34 | Row attribute access |
+| src/ai/backend/manager/repositories/schedule/repository.py | 29 | ORM query types |
+| src/ai/backend/appproxy/coordinator/models/worker.py | 29 | ORM migration needed |
+| src/ai/backend/appproxy/coordinator/models/circuit.py | 27 | ORM migration needed |
+| src/ai/backend/manager/repositories/vfolder/repository.py | 26 | ORM query types |
+| src/ai/backend/manager/models/user/row.py | 26 | Column vs runtime value types |
+
+### Error Type Categories
+
+#### Category 1: Row/Mapping Attribute Access (High Priority)
+**Total: ~60 errors**
+```
+"Mapping[str, Any]" has no attribute "id" (42)
+"Mapping[str, Any]" has no attribute "host" (6)
+etc.
+```
+**Fix**: Use proper Row class types instead of Mapping[str, Any]
+
+#### Category 2: String Index on Rows (Medium Priority)
+**Total: ~40 errors**
+```
+Invalid index type "str" for "str"; expected type "SupportsIndex | slice"
+No overload variant of "__getitem__" of "Row" matches argument type "str"
+```
+**Fix**: Change `row["column"]` to `row.column`
+
+#### Category 3: Boolean Expression Types (Medium Priority)
+**Total: ~30 errors**
+```
+Argument 1 to "where" has incompatible type "bool"
+Argument 2 to "and_" has incompatible type "bool"
+Argument 2 to "join" has incompatible type "bool"
+```
+**Fix**: Use explicit comparison operators (e.g., `column == value` instead of `value`)
+
+#### Category 4: select() List Args (Low Priority - Alembic)
+**Total: ~20 errors**
+```
+No overload variant of "select" matches argument type "list[...]"
+```
+**Fix**: `select([cols])` → `select(*cols)` (mostly in alembic files - skip)
+
+#### Category 5: Column vs Runtime Value Types (Medium Priority)
+**Total: ~40 errors**
+```
+Incompatible types in assignment (expression has type "str", variable has type "Column[str]")
+Argument "name" has incompatible type "Column[str]"; expected "str"
+```
+**Fix**: ORM 2.0 `Mapped[T]` type annotations fix these automatically
+
+#### Category 6: Nullable Type Issues (Low Priority)
+**Total: ~30 errors**
+```
+Item "None" of "Row[Any] | None" has no attribute "uuid"
+Item "None" of "KernelRow | None" has no attribute "status"
+```
+**Fix**: Add proper None checks or use `cast()` where appropriate
+
+#### Category 7: AppProxy Models (Separate Task)
+**Total: ~100 errors**
+Files:
+- src/ai/backend/appproxy/coordinator/models/worker.py (29)
+- src/ai/backend/appproxy/coordinator/models/circuit.py (27)
+- src/ai/backend/appproxy/coordinator/api/worker_v2.py (20)
+- src/ai/backend/appproxy/coordinator/api/health.py (14)
+- etc.
+
+**Fix**: Need ORM 2.0 migration for appproxy models
+
+### Alembic Files (Skip - Legacy Code)
+~200+ errors in alembic migration files
+- `Inspector.from_engine(Connection)` type issues
+- `existing_server_default` type issues
+- `select([...])` list args
+- `Row["column"]` indexing
+
+**Decision**: These are legacy migration files, type errors don't affect runtime.
+
+---
+
+## Next Steps (Recommended Order)
+
+### Step 1: Fix remaining Row["column"] patterns
+Files to check:
+```bash
+grep -r 'row\["' src/ai/backend/manager/ --include="*.py" | grep -v alembic | grep -v __pycache__
+```
+
+### Step 2: Fix Boolean expression issues
+Search pattern:
+```bash
+grep -rE '\.where\((True|False|[a-z_]+)\)' src/ai/backend/manager/ --include="*.py"
+```
+
+### Step 3: Fix Mapping[str, Any] attribute access
+Primary file: `registry.py` (105 errors)
+- Change function return types from `Mapping[str, Any]` to proper Row types
+
+### Step 4: AppProxy ORM Migration (Separate Task)
+- appproxy/coordinator/models/* need same ORM 2.0 migration as manager
+
+---
+
+## Commands
+
+### Run full type check
+```bash
+pants --no-colors --no-dynamic-ui check :: 2>&1 | tee /tmp/typecheck-results.txt
+```
+
+### Count errors by file
+```bash
+grep ": error:" /tmp/typecheck-results.txt | cut -d: -f1 | sort | uniq -c | sort -rn | head -30
+```
+
+### Count error types
+```bash
+grep ": error:" /tmp/typecheck-results.txt | sed 's/.*: error: //' | sed 's/  \[.*//' | sort | uniq -c | sort -rn | head -30
+```
+
+### Check specific file errors
+```bash
+grep "specific/file.py" /tmp/typecheck-results.txt
+```

python.lock

@@ -12,7 +12,7 @@
 //     "Jinja2~=3.1.6",
 //     "PyJWT~=2.10.1",
 //     "PyYAML~=6.0",
-//     "SQLAlchemy[postgresql_asyncpg]~=1.4.54",
+//     "SQLAlchemy[postgresql_asyncpg]~=2.0.45",
 //     "aioboto3~=15.0.0",
 //     "aiodataloader~=0.4.2",
 //     "aiodns==3.2",
@@ -6203,52 +6203,82 @@
           "artifacts": [
             {
               "algorithm": "sha256",
-              "hash": "4470fbed088c35dc20b78a39aaf4ae54fe81790c783b3264872a0224f437c31a",
-              "url": "https://files.pythonhosted.org/packages/ce/af/20290b55d469e873cba9d41c0206ab5461ff49d759989b3fe65010f9d265/sqlalchemy-1.4.54.tar.gz"
+              "hash": "5225a288e4c8cc2308dbdd874edad6e7d0fd38eac1e9e5f23503425c8eee20d0",
+              "url": "https://files.pythonhosted.org/packages/bf/e1/3ccb13c643399d22289c6a9786c1a91e3dcbb68bce4beb44926ac2c557bf/sqlalchemy-2.0.45-py3-none-any.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "672c45cae53ba88e0dad74b9027dddd09ef6f441e927786b05bec75d949fbb2e",
+              "url": "https://files.pythonhosted.org/packages/0e/50/80a8d080ac7d3d321e5e5d420c9a522b0aa770ec7013ea91f9a8b7d36e4a/sqlalchemy-2.0.45-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "83d7009f40ce619d483d26ac1b757dfe3167b39921379a8bd1b596cf02dab4a6",
+              "url": "https://files.pythonhosted.org/packages/3d/8d/bb40a5d10e7a5f2195f235c0b2f2c79b0bf6e8f00c0c223130a4fbd2db09/sqlalchemy-2.0.45-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "fe187fc31a54d7fd90352f34e8c008cf3ad5d064d08fedd3de2e8df83eb4a1cf",
+              "url": "https://files.pythonhosted.org/packages/6a/c8/7cc5221b47a54edc72a0140a1efa56e0a2730eefa4058d7ed0b4c4357ff8/sqlalchemy-2.0.45-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "9c6378449e0940476577047150fd09e242529b761dc887c9808a9a937fe990c8",
+              "url": "https://files.pythonhosted.org/packages/74/04/891b5c2e9f83589de202e7abaf24cd4e4fa59e1837d64d528829ad6cc107/sqlalchemy-2.0.45-cp313-cp313-musllinux_1_2_x86_64.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "d8a2ca754e5415cde2b656c27900b19d50ba076aa05ce66e2207623d3fe41f5a",
+              "url": "https://files.pythonhosted.org/packages/75/a5/346128b0464886f036c039ea287b7332a410aa2d3fb0bb5d404cb8861635/sqlalchemy-2.0.45-cp313-cp313t-musllinux_1_2_x86_64.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "1632a4bda8d2d25703fdad6363058d882541bdaaee0e5e3ddfa0cd3229efce88",
+              "url": "https://files.pythonhosted.org/packages/be/f9/5e4491e5ccf42f5d9cfc663741d261b3e6e1683ae7812114e7636409fcc6/sqlalchemy-2.0.45.tar.gz"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "470daea2c1ce73910f08caf10575676a37159a6d16c4da33d0033546bddebc9b",
+              "url": "https://files.pythonhosted.org/packages/da/4c/13dab31266fc9904f7609a5dc308a2432a066141d65b857760c3bef97e69/sqlalchemy-2.0.45-cp313-cp313-musllinux_1_2_aarch64.whl"
             }
           ],
           "project_name": "sqlalchemy",
           "requires_dists": [
-            "aiomysql>=0.2.0; python_version >= \"3\" and extra == \"aiomysql\"",
-            "aiosqlite; python_version >= \"3\" and extra == \"aiosqlite\"",
-            "asyncmy!=0.2.4,>=0.2.3; python_version >= \"3\" and extra == \"asyncmy\"",
-            "asyncpg; python_version >= \"3\" and extra == \"postgresql-asyncpg\"",
-            "asyncpg; python_version >= \"3\" and extra == \"postgresql-asyncpg\"",
-            "cx_oracle<8,>=7; python_version < \"3\" and extra == \"oracle\"",
-            "cx_oracle>=7; python_version >= \"3\" and extra == \"oracle\"",
-            "greenlet!=0.4.17; python_version >= \"3\" and (platform_machine == \"aarch64\" or (platform_machine == \"ppc64le\" or (platform_machine == \"x86_64\" or (platform_machine == \"amd64\" or (platform_machine == \"AMD64\" or (platform_machine == \"win32\" or platform_machine == \"WIN32\"))))))",
-            "greenlet!=0.4.17; python_version >= \"3\" and extra == \"aiomysql\"",
-            "greenlet!=0.4.17; python_version >= \"3\" and extra == \"aiosqlite\"",
-            "greenlet!=0.4.17; python_version >= \"3\" and extra == \"asyncio\"",
-            "greenlet!=0.4.17; python_version >= \"3\" and extra == \"asyncmy\"",
-            "greenlet!=0.4.17; python_version >= \"3\" and extra == \"postgresql-asyncpg\"",
-            "greenlet!=0.4.17; python_version >= \"3\" and extra == \"postgresql-asyncpg\"",
+            "aiomysql>=0.2.0; extra == \"aiomysql\"",
+            "aioodbc; extra == \"aioodbc\"",
+            "aiosqlite; extra == \"aiosqlite\"",
+            "asyncmy!=0.2.4,!=0.2.6,>=0.2.3; extra == \"asyncmy\"",
+            "asyncpg; extra == \"postgresql-asyncpg\"",
+            "cx_oracle>=8; extra == \"oracle\"",
+            "greenlet>=1; extra == \"aiomysql\"",
+            "greenlet>=1; extra == \"aioodbc\"",
+            "greenlet>=1; extra == \"aiosqlite\"",
+            "greenlet>=1; extra == \"asyncio\"",
+            "greenlet>=1; extra == \"asyncmy\"",
+            "greenlet>=1; extra == \"postgresql-asyncpg\"",
+            "greenlet>=1; platform_machine == \"aarch64\" or (platform_machine == \"ppc64le\" or (platform_machine == \"x86_64\" or (platform_machine == \"amd64\" or (platform_machine == \"AMD64\" or (platform_machine == \"win32\" or platform_machine == \"WIN32\")))))",
             "importlib-metadata; python_version < \"3.8\"",
-            "mariadb!=1.1.2,>=1.0.1; python_version >= \"3\" and extra == \"mariadb-connector\"",
-            "mariadb!=1.1.2,>=1.0.1; python_version >= \"3\" and extra == \"mariadb-connector\"",
-            "mypy>=0.910; python_version >= \"3\" and extra == \"mypy\"",
-            "mysql-connector-python; extra == \"mysql-connector\"",
+            "mariadb!=1.1.10,!=1.1.2,!=1.1.5,>=1.0.1; extra == \"mariadb-connector\"",
+            "mypy>=0.910; extra == \"mypy\"",
             "mysql-connector-python; extra == \"mysql-connector\"",
-            "mysqlclient<2,>=1.4.0; python_version < \"3\" and extra == \"mysql\"",
-            "mysqlclient>=1.4.0; python_version >= \"3\" and extra == \"mysql\"",
-            "pg8000!=1.29.0,>=1.16.6; python_version >= \"3\" and extra == \"postgresql-pg8000\"",
-            "pg8000!=1.29.0,>=1.16.6; python_version >= \"3\" and extra == \"postgresql-pg8000\"",
+            "mysqlclient>=1.4.0; extra == \"mysql\"",
+            "oracledb>=1.0.1; extra == \"oracle-oracledb\"",
+            "pg8000>=1.29.1; extra == \"postgresql-pg8000\"",
             "psycopg2-binary; extra == \"postgresql-psycopg2binary\"",
             "psycopg2>=2.7; extra == \"postgresql\"",
             "psycopg2cffi; extra == \"postgresql-psycopg2cffi\"",
+            "psycopg>=3.0.7; extra == \"postgresql-psycopg\"",
+            "psycopg[binary]>=3.0.7; extra == \"postgresql-psycopgbinary\"",
             "pymssql; extra == \"mssql-pymssql\"",
-            "pymssql; extra == \"mssql-pymssql\"",
-            "pymysql; python_version >= \"3\" and extra == \"pymysql\"",
-            "pymysql<1; python_version < \"3\" and extra == \"pymysql\"",
+            "pymysql; extra == \"pymysql\"",
             "pyodbc; extra == \"mssql\"",
             "pyodbc; extra == \"mssql-pyodbc\"",
-            "pyodbc; extra == \"mssql-pyodbc\"",
-            "sqlalchemy2-stubs; extra == \"mypy\"",
-            "sqlcipher3_binary; python_version >= \"3\" and extra == \"sqlcipher\"",
+            "sqlcipher3_binary; extra == \"sqlcipher\"",
+            "typing-extensions>=4.6.0",
             "typing_extensions!=3.10.0.1; extra == \"aiosqlite\""
           ],
-          "requires_python": "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,>=2.7",
-          "version": "1.4.54"
+          "requires_python": ">=3.7",
+          "version": "2.0.45"
         },
         {
           "artifacts": [
@@ -7651,7 +7681,7 @@
     "Jinja2~=3.1.6",
     "PyJWT~=2.10.1",
     "PyYAML~=6.0",
-    "SQLAlchemy[postgresql_asyncpg]~=1.4.54",
+    "SQLAlchemy[postgresql_asyncpg]~=2.0.45",
     "aioboto3~=15.0.0",
     "aiodataloader~=0.4.2",
     "aiodns==3.2",

requirements.txt

@@ -78,7 +78,7 @@ huggingface-hub~=0.34.3
 redis[hiredis]~=7.1.0
 rich~=13.6
 ruamel.yaml~=0.18.10
-SQLAlchemy[postgresql_asyncpg]~=1.4.54
+SQLAlchemy[postgresql_asyncpg]~=2.0.45
 setproctitle~=1.3.5
 setuptools~=80.0.0
 strawberry-graphql~=0.278.0

src/ai/backend/account_manager/models/base.py

@@ -117,13 +117,13 @@ def process_bind_param(
 
     def process_result_value(
         self,
-        value: str,
+        value: Any | None,
         dialect: Dialect,
     ) -> T_StrEnum | None:
         return self._enum_cls(value) if value is not None else None
 
-    def copy(self, **kw) -> type[Self]:
-        return StrEnumType(self._enum_cls, **self._opts)
+    def copy(self, **kw) -> Self:
+        return StrEnumType(self._enum_cls, **self._opts)  # type: ignore[return-value]
 
     @property
     def python_type(self) -> T_StrEnum:

src/ai/backend/account_manager/models/utils.py

@@ -17,7 +17,7 @@
 from sqlalchemy.ext.asyncio import AsyncConnection as SAConnection
 from sqlalchemy.ext.asyncio import AsyncEngine as SAEngine
 from sqlalchemy.ext.asyncio import AsyncSession as SASession
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.ext.asyncio import async_sessionmaker
 from tenacity import (
     AsyncRetrying,
     RetryError,
@@ -53,8 +53,8 @@ def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self._readonly_txn_count = 0
         self._generic_txn_count = 0
-        self._sess_factory = sessionmaker(self, expire_on_commit=False, class_=SASession)
-        self._readonly_sess_factory = sessionmaker(self, class_=SASession)
+        self._sess_factory = async_sessionmaker(self, expire_on_commit=False)
+        self._readonly_sess_factory = async_sessionmaker(self)
 
     def _check_generic_txn_cnt(self) -> None:
         if (

src/ai/backend/appproxy/coordinator/models/utils.py

@@ -24,7 +24,7 @@
 from sqlalchemy.ext.asyncio import AsyncConnection as SAConnection
 from sqlalchemy.ext.asyncio import AsyncEngine as SAEngine
 from sqlalchemy.ext.asyncio import AsyncSession as SASession
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.ext.asyncio import async_sessionmaker
 from tenacity import (
     AsyncRetrying,
     RetryError,
@@ -66,8 +66,8 @@ def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self._readonly_txn_count = 0
         self._generic_txn_count = 0
-        self._sess_factory = sessionmaker(self, expire_on_commit=False, class_=SASession)
-        self._readonly_sess_factory = sessionmaker(self, class_=SASession)
+        self._sess_factory = async_sessionmaker(self, expire_on_commit=False)
+        self._readonly_sess_factory = async_sessionmaker(self)
 
     def _check_generic_txn_cnt(self) -> None:
         if (

src/ai/backend/manager/agent_cache.py

@@ -96,10 +96,10 @@ async def get_rpc_args(self, agent_id: AgentId) -> tuple[str, PublicKey | None]:
         if cached_args:
             return cached_args
 
-        async def _fetch_agent() -> Row:
+        async def _fetch_agent() -> Row | None:
             async with self.db.begin_readonly() as conn:
                 query = (
-                    sa.select([agents.c.addr, agents.c.public_key])
+                    sa.select(agents.c.addr, agents.c.public_key)
                     .select_from(agents)
                     .where(
                         agents.c.id == agent_id,
@@ -109,7 +109,9 @@ async def _fetch_agent() -> Row:
                 return result.first()
 
         agent = await execute_with_retry(_fetch_agent)
-        return agent["addr"], agent["public_key"]
+        if agent is None:
+            raise ValueError(f"Agent not found: {agent_id}")
+        return agent.addr, agent.public_key
 
     @actxmgr
     async def rpc_context(