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

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

@@ -99,7 +99,7 @@ async def get_rpc_args(self, agent_id: AgentId) -> tuple[str, PublicKey | None]:
         async def _fetch_agent() -> Row:
             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,7 @@ async def _fetch_agent() -> Row:
                 return result.first()
 
         agent = await execute_with_retry(_fetch_agent)
-        return agent["addr"], agent["public_key"]
+        return agent.addr, agent.public_key
 
     @actxmgr
     async def rpc_context(

src/ai/backend/manager/api/auth.py

@@ -443,7 +443,7 @@ async def _query_cred_by_access_key(
             keypairs.c.resource_policy == keypair_resource_policies.c.name,
         )
         query = (
-            sa.select([keypairs, keypair_resource_policies], use_labels=True)
+            sa.select(keypairs, keypair_resource_policies, use_labels=True)
             .select_from(j)
             .where(
                 (keypairs.c.access_key == access_key) & (keypairs.c.is_active.is_(True)),
@@ -464,7 +464,7 @@ async def _query_cred_by_access_key(
             users.c.uuid == keypairs.c.user,
         )
         query = (
-            sa.select([users, user_resource_policies], use_labels=True)
+            sa.select(users, user_resource_policies, use_labels=True)
             .select_from(j)
             .where(keypairs.c.access_key == access_key)
         )
@@ -500,7 +500,7 @@ def _populate_auth_result(
             for col in users.c
             if col.name not in ("password", "description", "created_at")
         },
-        "is_admin": keypair_row["keypairs_is_admin"],
+        "is_admin": keypair_row.keypairs_is_admin,
     }
 
     validate_ip(request, auth_result["user"])
@@ -512,7 +512,7 @@ def _populate_auth_result(
     auth_result["user"]["resource_policy"] = {
         col.name: user_row[f"user_resource_policies_{col.name}"] for col in user_resource_policies.c
     }
-    auth_result["user"]["id"] = keypair_row["keypairs_user_id"]  # legacy
+    auth_result["user"]["id"] = keypair_row.keypairs_user_id  # legacy
     auth_result["is_superadmin"] = auth_result["user"]["role"] == "superadmin"
 
     # Populate the result to the per-request state dict
@@ -560,7 +560,7 @@ async def _authenticate_via_jwt(
             raise AuthorizationFailed("Access key not found in database")
 
         # 3. Validate JWT token using user's secret key
-        secret_key = keypair_row["keypairs_secret_key"]
+        secret_key = keypair_row.keypairs_secret_key
         root_ctx.jwt_validator.validate_token(jwt_token, secret_key)
 
         # 4. Populate authentication result
@@ -613,7 +613,7 @@ async def _authenticate_via_hmac(
         raise AuthorizationFailed("Access key not found in HMAC")
 
     # 4. Verify HMAC signature
-    my_signature = await sign_request(sign_method, request, keypair_row["keypairs_secret_key"])
+    my_signature = await sign_request(sign_method, request, keypair_row.keypairs_secret_key)
 
     if not secrets.compare_digest(my_signature, signature):
         raise AuthorizationFailed("HMAC signature mismatch")

src/ai/backend/manager/api/cluster_template.py

@@ -69,23 +69,23 @@ async def create(request: web.Request, params: Any) -> web.Response:
             # Admin or superadmin is creating sessions for another user.
             # The check for admin privileges is already done in get_access_key_scope().
             query = (
-                sa.select([keypairs.c.user, users.c.role, users.c.domain_name])
+                sa.select(keypairs.c.user, users.c.role, users.c.domain_name)
                 .select_from(sa.join(keypairs, users, keypairs.c.user == users.c.uuid))
                 .where(keypairs.c.access_key == owner_access_key)
             )
             result = await conn.execute(query)
             row = result.first()
-            owner_domain = row["domain_name"]
-            owner_uuid = row["user"]
-            owner_role = row["role"]
+            owner_domain = row.domain_name
+            owner_uuid = row.user
+            owner_role = row.role
         else:
             # Normal case when the user is creating her/his own session.
             owner_domain = request["user"]["domain_name"]
             owner_uuid = requester_uuid
             owner_role = UserRole.USER
 
         query = (
-            sa.select([domains.c.name])
+            sa.select(domains.c.name)
             .select_from(domains)
             .where(
                 (domains.c.name == owner_domain) & (domains.c.is_active),
@@ -99,7 +99,7 @@ async def create(request: web.Request, params: Any) -> web.Response:
         if owner_role == UserRole.SUPERADMIN:
             # superadmin can spawn container in any designated domain/group.
             query = (
-                sa.select([groups.c.id])
+                sa.select(groups.c.id)
                 .select_from(groups)
                 .where(
                     (groups.c.domain_name == params["domain"])
@@ -114,7 +114,7 @@ async def create(request: web.Request, params: Any) -> web.Response:
             if params["domain"] != owner_domain:
                 raise InvalidAPIParameters("You can only set the domain to the owner's domain.")
             query = (
-                sa.select([groups.c.id])
+                sa.select(groups.c.id)
                 .select_from(groups)
                 .where(
                     (groups.c.domain_name == owner_domain)
@@ -129,7 +129,7 @@ async def create(request: web.Request, params: Any) -> web.Response:
             if params["domain"] != owner_domain:
                 raise InvalidAPIParameters("You can only set the domain to your domain.")
             query = (
-                sa.select([agus.c.group_id])
+                sa.select(agus.c.group_id)
                 .select_from(agus.join(groups, agus.c.group_id == groups.c.id))
                 .where(
                     (agus.c.user_id == owner_uuid)
@@ -196,7 +196,7 @@ async def list_template(request: web.Request, params: Any) -> web.Response:
                 users, session_templates.c.user_uuid == users.c.uuid, isouter=True
             ).join(groups, session_templates.c.group_id == groups.c.id, isouter=True)
             query = (
-                sa.select([session_templates, users.c.email, groups.c.name], use_labels=True)
+                sa.select(session_templates, users.c.email, groups.c.name, use_labels=True)
                 .select_from(j)
                 .where(
                     (session_templates.c.is_active)
@@ -277,7 +277,7 @@ async def get(request: web.Request, params: Any) -> web.Response:
 
     async with root_ctx.db.begin() as conn:
         query = (
-            sa.select([session_templates.c.template])
+            sa.select(session_templates.c.template)
             .select_from(session_templates)
             .where(
                 (session_templates.c.id == template_id)
@@ -316,7 +316,7 @@ async def put(request: web.Request, params: Any) -> web.Response:
 
     async with root_ctx.db.begin() as conn:
         query = (
-            sa.select([session_templates.c.id])
+            sa.select(session_templates.c.id)
             .select_from(session_templates)
             .where(
                 (session_templates.c.id == template_id)
@@ -366,7 +366,7 @@ async def delete(request: web.Request, params: Any) -> web.Response:
 
     async with root_ctx.db.begin() as conn:
         query = (
-            sa.select([session_templates.c.id])
+            sa.select(session_templates.c.id)
             .select_from(session_templates)
             .where(
                 (session_templates.c.id == template_id)

src/ai/backend/manager/api/events.py

@@ -138,12 +138,12 @@ async def push_session_events(
         group_id = WILDCARD
     else:
         async with root_ctx.db.begin_readonly(isolation_level="READ COMMITTED") as conn:
-            query = sa.select([groups.c.id]).select_from(groups).where(groups.c.name == group_name)
+            query = sa.select(groups.c.id).select_from(groups).where(groups.c.name == group_name)
             result = await conn.execute(query)
             row = result.first()
             if row is None:
                 raise ProjectNotFound
-            group_id = row["id"]
+            group_id = row.id
 
     filters = {
         "user_role": user_role,