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

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Co-authored-by: octodog <mu001@lablup.com>

Author: 3 people

changes/7880.deps.md

@@ -0,0 +1 @@
+Migrate to SQLAlchemy 2.0 with comprehensive type safety improvements

docs/manager/graphql-reference/supergraph.graphql

@@ -3892,7 +3892,7 @@ type ModelRevision implements Node
   modelRuntimeConfig: ModelRuntimeConfig!
 
   """Model data mount configuration."""
-  modelMountConfig: ModelMountConfig!
+  modelMountConfig: ModelMountConfig
 
   """Additional volume folder mounts."""
   extraMounts: ExtraVFolderMountConnection!
@@ -6379,7 +6379,7 @@ type Route implements Node
   trafficRatio: Float!
 
   """The timestamp when the route was created."""
-  createdAt: DateTime!
+  createdAt: DateTime
 
   """Error data if the route is in a failed state."""
   errorData: JSONString

docs/manager/graphql-reference/v2-schema.graphql

@@ -1464,7 +1464,7 @@ type ModelRevision implements Node {
   modelRuntimeConfig: ModelRuntimeConfig!
 
   """Model data mount configuration."""
-  modelMountConfig: ModelMountConfig!
+  modelMountConfig: ModelMountConfig
 
   """Additional volume folder mounts."""
   extraMounts: ExtraVFolderMountConnection!
@@ -2410,7 +2410,7 @@ type Route implements Node {
   trafficRatio: Float!
 
   """The timestamp when the route was created."""
-  createdAt: DateTime!
+  createdAt: DateTime
 
   """Error data if the route is in a failed state."""
   errorData: JSONString

pants.toml

@@ -26,6 +26,10 @@ pants_ignore = [
     "/tools/pants-plugins",
     "/wheelhouse/*/",
     "/wheelhouse/*.whl",
+    # Alembic migration files - legacy SQLAlchemy 1.x style, excluded from type checking
+    "/src/ai/backend/manager/models/alembic/versions/",
+    "/src/ai/backend/account_manager/models/alembic/versions/",
+    "/src/ai/backend/appproxy/coordinator/models/alembic/versions/",
 ]
 build_file_prelude_globs = ["tools/build-macros.py"]
 
@@ -125,6 +129,7 @@ report = ["xml", "console"]
 
 [mypy]
 install_from_resolve = "mypy"
+config = "pyproject.toml"
 
 [setuptools]
 install_from_resolve = "setuptools"

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",
@@ -3282,13 +3282,13 @@
           "artifacts": [
             {
               "algorithm": "sha256",
-              "hash": "3671a94fd25e62f5f2f554f5e95389c2294d89822378a5f2dd24353e1494a9e0",
-              "url": "https://files.pythonhosted.org/packages/bb/f5/fddaec430367be9d62a7ed125530e133bfd4a1c0350fe221149ee0f2b526/jupyter_client-8.7.0-py3-none-any.whl"
+              "hash": "f93a5b99c5e23a507b773d3a1136bd6e16c67883ccdbd9a829b0bbdb98cd7d7a",
+              "url": "https://files.pythonhosted.org/packages/2d/0b/ceb7694d864abc0a047649aec263878acb9f792e1fec3e676f22dc9015e3/jupyter_client-8.8.0-py3-none-any.whl"
             },
             {
               "algorithm": "sha256",
-              "hash": "3357212d9cbe01209e59190f67a3a7e1f387a4f4e88d1e0433ad84d7b262531d",
-              "url": "https://files.pythonhosted.org/packages/a6/27/d10de45e8ad4ce872372c4a3a37b7b35b6b064f6f023a5c14ffcced4d59d/jupyter_client-8.7.0.tar.gz"
+              "hash": "d556811419a4f2d96c869af34e854e3f059b7cc2d6d01a9cd9c85c267691be3e",
+              "url": "https://files.pythonhosted.org/packages/05/e4/ba649102a3bc3fbca54e7239fb924fd434c766f855693d86de0b1f2bec81/jupyter_client-8.8.0.tar.gz"
             }
           ],
           "project_name": "jupyter-client",
@@ -3298,8 +3298,10 @@
             "ipykernel; extra == \"docs\"",
             "ipykernel>=6.14; extra == \"test\"",
             "jupyter-core>=5.1",
-            "mypy; extra == \"test\"",
+            "msgpack; extra == \"test\"",
+            "mypy; platform_python_implementation != \"PyPy\" and extra == \"test\"",
             "myst-parser; extra == \"docs\"",
+            "orjson; extra == \"orjson\"",
             "paramiko; sys_platform == \"win32\" and extra == \"test\"",
             "pre-commit; extra == \"test\"",
             "pydata-sphinx-theme; extra == \"docs\"",
@@ -3317,7 +3319,7 @@
             "traitlets>=5.3"
           ],
           "requires_python": ">=3.10",
-          "version": "8.7.0"
+          "version": "8.8.0"
         },
         {
           "artifacts": [
@@ -6203,52 +6205,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 +7683,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/alembic/env.py

@@ -62,7 +62,7 @@ def do_run_migrations(connection: Connection) -> None:
 
 async def run_async_migrations() -> None:
     connectable = async_engine_from_config(
-        config.get_section(config.config_ini_section),
+        config.get_section(config.config_ini_section) or {},
         prefix="sqlalchemy.",
         poolclass=pool.NullPool,
     )

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

@@ -117,17 +117,17 @@ 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:
-        return self._enum_class
+    def python_type(self) -> type[T_StrEnum]:
+        return self._enum_cls
 
 
 class PasswordColumn(TypeDecorator):

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

@@ -1,41 +1,50 @@
+from datetime import datetime
+from uuid import UUID
+
 import sqlalchemy as sa
-from sqlalchemy.orm import relationship
+from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from ai.backend.account_manager.types import UserRole, UserStatus
 from ai.backend.account_manager.utils import verify_password
 
-from .base import GUID, Base, IDColumn, PasswordColumn, StrEnumType
+from .base import GUID, Base, PasswordColumn, StrEnumType
 
 __all__: tuple[str, ...] = ("UserProfileRow",)
 
 
 class UserProfileRow(Base):
     __tablename__ = "user_profiles"
-    id = IDColumn()
-    user_id = sa.Column("user_id", GUID, nullable=False)
-    username = sa.Column("username", sa.String(length=64), index=True, nullable=False, unique=True)
-    email = sa.Column("email", sa.String(length=64), index=True, nullable=False)
-    password = sa.Column("password", PasswordColumn(), nullable=False)
-    need_password_change = sa.Column("need_password_change", sa.Boolean, server_default=sa.false())
-    password_changed_at = sa.Column(
+    id: Mapped[UUID] = mapped_column(
+        GUID, primary_key=True, server_default=sa.text("uuid_generate_v4()")
+    )
+    user_id: Mapped[UUID] = mapped_column("user_id", GUID, nullable=False)
+    username: Mapped[str] = mapped_column(
+        "username", sa.String(length=64), index=True, nullable=False, unique=True
+    )
+    email: Mapped[str] = mapped_column("email", sa.String(length=64), index=True, nullable=False)
+    password: Mapped[str] = mapped_column("password", PasswordColumn(), nullable=False)
+    need_password_change: Mapped[bool | None] = mapped_column(
+        "need_password_change", sa.Boolean, server_default=sa.false()
+    )
+    password_changed_at: Mapped[datetime | None] = mapped_column(
         "password_changed_at",
         sa.DateTime(timezone=True),
         server_default=sa.func.now(),
     )
-    full_name = sa.Column("full_name", sa.String(length=64))
-    description = sa.Column("description", sa.String(length=500))
+    full_name: Mapped[str | None] = mapped_column("full_name", sa.String(length=64))
+    description: Mapped[str | None] = mapped_column("description", sa.String(length=500))
 
-    role = sa.Column(
+    role: Mapped[UserRole] = mapped_column(
         "role",
         StrEnumType(UserRole),
         default=UserRole.USER,
         server_default=UserRole.USER.value,
         nullable=False,
     )
-    status = sa.Column(
+    status: Mapped[UserStatus] = mapped_column(
         "status", StrEnumType(UserStatus), server_default=UserStatus.ACTIVE.value, nullable=False
     )
-    status_info = sa.Column("status_info", sa.Unicode(), nullable=True)
+    status_info: Mapped[str | None] = mapped_column("status_info", sa.Unicode(), nullable=True)
 
     created_at = sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now())
     modified_at = sa.Column(

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 (
@@ -229,6 +229,8 @@ async def connect_database(
     async with version_check_db.begin() as conn:
         result = await conn.execute(sa.text("show server_version"))
         version_str = result.scalar()
+        if version_str is None:
+            raise RuntimeError("Failed to get PostgreSQL version")
         major, minor, *_ = map(int, version_str.partition(" ")[0].split("."))
         if (major, minor) < (11, 0):
             pgsql_connect_opts["server_settings"].pop("jit")