fix(enterprise): migrate user models to SQLAlchemy 2.0 [6/13] (#13852)

Co-authored-by: openhands <openhands@all-hands.dev>

enterprise/storage/user.py

@@ -2,46 +2,50 @@
 SQLAlchemy model for User.
 """
 
-from uuid import uuid4
+from datetime import datetime
+from typing import TYPE_CHECKING
+from uuid import UUID, uuid4
 
-from sqlalchemy import (
-    JSON,
-    UUID,
-    Boolean,
-    Column,
-    DateTime,
-    ForeignKey,
-    Integer,
-    String,
-)
-from sqlalchemy.orm import relationship
+from sqlalchemy import DateTime, ForeignKey, String
+from sqlalchemy.dialects.postgresql import JSON
+from sqlalchemy.orm import Mapped, mapped_column, relationship
 from storage.base import Base
 
+if TYPE_CHECKING:
+    from storage.org import Org
+    from storage.org_member import OrgMember
+    from storage.role import Role
+    from storage.stored_conversation_metadata_saas import StoredConversationMetadataSaas
 
-class User(Base):  # type: ignore
+
+class User(Base):
     """User model with organizational relationships."""
 
     __tablename__ = 'user'
 
-    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid4)
-    current_org_id = Column(UUID(as_uuid=True), ForeignKey('org.id'), nullable=False)
-    role_id = Column(Integer, ForeignKey('role.id'), nullable=True)
-    accepted_tos = Column(DateTime, nullable=True)
-    enable_sound_notifications = Column(Boolean, nullable=True)
-    language = Column(String, nullable=True)
-    user_consents_to_analytics = Column(Boolean, nullable=True)
-    email = Column(String, nullable=True)
-    email_verified = Column(Boolean, nullable=True)
-    git_user_name = Column(String, nullable=True)
-    git_user_email = Column(String, nullable=True)
-    sandbox_grouping_strategy = Column(String, nullable=True)
-    disabled_skills = Column(JSON, nullable=True)
-    onboarding_completed = Column(Boolean, nullable=True, default=False)
+    id: Mapped[UUID] = mapped_column(primary_key=True, default=uuid4)
+    current_org_id: Mapped[UUID] = mapped_column(ForeignKey('org.id'), nullable=False)
+    role_id: Mapped[int | None] = mapped_column(ForeignKey('role.id'), nullable=True)
+    accepted_tos: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
+    enable_sound_notifications: Mapped[bool | None] = mapped_column(nullable=True)
+    language: Mapped[str | None] = mapped_column(String, nullable=True)
+    user_consents_to_analytics: Mapped[bool | None] = mapped_column(nullable=True)
+    email: Mapped[str | None] = mapped_column(String, nullable=True)
+    email_verified: Mapped[bool | None] = mapped_column(nullable=True)
+    git_user_name: Mapped[str | None] = mapped_column(String, nullable=True)
+    git_user_email: Mapped[str | None] = mapped_column(String, nullable=True)
+    sandbox_grouping_strategy: Mapped[str | None] = mapped_column(String, nullable=True)
+    disabled_skills: Mapped[list[str] | None] = mapped_column(JSON, nullable=True)
+    onboarding_completed: Mapped[bool | None] = mapped_column(
+        nullable=True, default=False
+    )
 
     # Relationships
-    role = relationship('Role', back_populates='users')
-    org_members = relationship('OrgMember', back_populates='user')
-    current_org = relationship('Org', back_populates='current_users')
-    stored_conversation_metadata_saas = relationship(
-        'StoredConversationMetadataSaas', back_populates='user'
+    role: Mapped['Role | None'] = relationship('Role', back_populates='users')
+    org_members: Mapped[list['OrgMember']] = relationship(
+        'OrgMember', back_populates='user'
     )
+    current_org: Mapped['Org'] = relationship('Org', back_populates='current_users')
+    stored_conversation_metadata_saas: Mapped[
+        list['StoredConversationMetadataSaas']
+    ] = relationship('StoredConversationMetadataSaas', back_populates='user')

enterprise/storage/user_authorization.py

@@ -3,7 +3,8 @@
 from datetime import UTC, datetime
 from enum import Enum
 
-from sqlalchemy import Column, DateTime, Identity, Integer, String
+from sqlalchemy import DateTime, Identity, String
+from sqlalchemy.orm import Mapped, mapped_column
 from storage.base import Base
 
 
@@ -14,7 +15,7 @@ class UserAuthorizationType(str, Enum):
     BLACKLIST = 'blacklist'
 
 
-class UserAuthorization(Base):  # type: ignore
+class UserAuthorization(Base):
     """Stores user authorization rules based on email patterns and provider types.
 
     Supports:
@@ -28,16 +29,16 @@ class UserAuthorization(Base):  # type: ignore
 
     __tablename__ = 'user_authorizations'
 
-    id = Column(Integer, Identity(), primary_key=True)
-    email_pattern = Column(String, nullable=True)
-    provider_type = Column(String, nullable=True)
-    type = Column(String, nullable=False)
-    created_at = Column(
+    id: Mapped[int] = mapped_column(Identity(), primary_key=True)
+    email_pattern: Mapped[str | None] = mapped_column(String, nullable=True)
+    provider_type: Mapped[str | None] = mapped_column(String, nullable=True)
+    type: Mapped[str] = mapped_column(String, nullable=False)
+    created_at: Mapped[datetime] = mapped_column(
         DateTime(timezone=True),
         default=lambda: datetime.now(UTC),
         nullable=False,
     )
-    updated_at = Column(
+    updated_at: Mapped[datetime] = mapped_column(
         DateTime(timezone=True),
         default=lambda: datetime.now(UTC),
         onupdate=lambda: datetime.now(UTC),

enterprise/storage/user_repo_map.py

@@ -1,4 +1,5 @@
-from sqlalchemy import Boolean, Column, Integer, String
+from sqlalchemy import String
+from sqlalchemy.orm import Mapped, mapped_column
 from storage.base import Base
 
 
@@ -8,7 +9,8 @@ class UserRepositoryMap(Base):
     """
 
     __tablename__ = 'user-repos'
-    id = Column(Integer, primary_key=True, autoincrement=True)
-    user_id = Column(String, nullable=False)
-    repo_id = Column(String, nullable=False)
-    admin = Column(Boolean, nullable=True)
+
+    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+    user_id: Mapped[str] = mapped_column(String, nullable=False)
+    repo_id: Mapped[str] = mapped_column(String, nullable=False)
+    admin: Mapped[bool | None] = mapped_column(nullable=True)

enterprise/storage/user_settings.py

@@ -1,47 +1,69 @@
 from __future__ import annotations
 
+from datetime import datetime
+from typing import Any
+
 from pydantic import SecretStr
 from server.constants import DEFAULT_BILLING_MARGIN
-from sqlalchemy import JSON, Boolean, Column, DateTime, Float, Identity, Integer, String
+from sqlalchemy import DateTime, Identity, String
+from sqlalchemy.dialects.postgresql import JSON
+from sqlalchemy.orm import Mapped, mapped_column
 from storage.base import Base
 from storage.encrypt_utils import decrypt_legacy_value, encrypt_legacy_value
 
 
-class UserSettings(Base):  # type: ignore
+class UserSettings(Base):
     __tablename__ = 'user_settings'
-    id = Column(Integer, Identity(), primary_key=True)
-    keycloak_user_id = Column(String, nullable=True, index=True)
-    language = Column(String, nullable=True)
+
+    id: Mapped[int] = mapped_column(Identity(), primary_key=True)
+    keycloak_user_id: Mapped[str | None] = mapped_column(
+        String, nullable=True, index=True
+    )
+    language: Mapped[str | None] = mapped_column(String, nullable=True)
     # Deprecated (v0): API keys now live on Org / OrgMember.
     # Kept for backward-compat during migration; do not use in new code.
-    llm_api_key = Column(String, nullable=True)
-    llm_api_key_for_byor = Column(String, nullable=True)
-    remote_runtime_resource_factor = Column(Integer, nullable=True)
-    user_consents_to_analytics = Column(Boolean, nullable=True)
-    billing_margin = Column(Float, nullable=True, default=DEFAULT_BILLING_MARGIN)
-    enable_sound_notifications = Column(Boolean, nullable=True, default=False)
-    enable_proactive_conversation_starters = Column(
-        Boolean, nullable=False, default=True
+    llm_api_key: Mapped[str | None] = mapped_column(String, nullable=True)
+    llm_api_key_for_byor: Mapped[str | None] = mapped_column(String, nullable=True)
+    remote_runtime_resource_factor: Mapped[int | None] = mapped_column(nullable=True)
+    user_consents_to_analytics: Mapped[bool | None] = mapped_column(nullable=True)
+    billing_margin: Mapped[float | None] = mapped_column(
+        nullable=True, default=DEFAULT_BILLING_MARGIN
     )
-    sandbox_base_container_image = Column(String, nullable=True)
-    sandbox_runtime_container_image = Column(String, nullable=True)
-    sandbox_grouping_strategy = Column(String, nullable=True)
-    user_version = Column(Integer, nullable=False, default=0)
-    accepted_tos = Column(DateTime, nullable=True)
+    enable_sound_notifications: Mapped[bool | None] = mapped_column(
+        nullable=True, default=False
+    )
+    enable_proactive_conversation_starters: Mapped[bool] = mapped_column(
+        nullable=False, default=True
+    )
+    sandbox_base_container_image: Mapped[str | None] = mapped_column(
+        String, nullable=True
+    )
+    sandbox_runtime_container_image: Mapped[str | None] = mapped_column(
+        String, nullable=True
+    )
+    sandbox_grouping_strategy: Mapped[str | None] = mapped_column(String, nullable=True)
+    user_version: Mapped[int] = mapped_column(nullable=False, default=0)
+    accepted_tos: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
     # Deprecated (v0): mcp_config now lives inside AgentSettings on Org / OrgMember.
-    mcp_config = Column(JSON, nullable=True)
-    disabled_skills = Column(JSON, nullable=True)
-    search_api_key = Column(String, nullable=True)
-    sandbox_api_key = Column(String, nullable=True)
-    max_budget_per_task = Column(Float, nullable=True)
-    enable_solvability_analysis = Column(Boolean, nullable=True, default=False)
-    email = Column(String, nullable=True)
-    email_verified = Column(Boolean, nullable=True)
-    git_user_name = Column(String, nullable=True)
-    git_user_email = Column(String, nullable=True)
-    v1_enabled = Column(Boolean, nullable=True)
-    agent_settings = Column(JSON, nullable=False, default=dict)
-    conversation_settings = Column(JSON, nullable=False, default=dict)
+    mcp_config: Mapped[dict[str, Any] | None] = mapped_column(JSON, nullable=True)
+    disabled_skills: Mapped[list[str] | None] = mapped_column(JSON, nullable=True)
+    search_api_key: Mapped[str | None] = mapped_column(String, nullable=True)
+    sandbox_api_key: Mapped[str | None] = mapped_column(String, nullable=True)
+    max_budget_per_task: Mapped[float | None] = mapped_column(nullable=True)
+    enable_solvability_analysis: Mapped[bool | None] = mapped_column(
+        nullable=True, default=False
+    )
+    email: Mapped[str | None] = mapped_column(String, nullable=True)
+    email_verified: Mapped[bool | None] = mapped_column(nullable=True)
+    git_user_name: Mapped[str | None] = mapped_column(String, nullable=True)
+    git_user_email: Mapped[str | None] = mapped_column(String, nullable=True)
+    v1_enabled: Mapped[bool | None] = mapped_column(nullable=True)
+    agent_settings: Mapped[dict[str, Any]] = mapped_column(
+        JSON, nullable=False, default=dict
+    )
+    conversation_settings: Mapped[dict[str, Any]] = mapped_column(
+        JSON, nullable=False, default=dict
+    )
 
     @property
     def llm_api_key_for_byor_secret(self) -> SecretStr | None:
@@ -61,8 +83,8 @@ class UserSettings(Base):  # type: ignore
         raw = value.get_secret_value() if isinstance(value, SecretStr) else value
         self.llm_api_key_for_byor = encrypt_legacy_value(raw)
 
-    already_migrated = Column(
-        Boolean, nullable=True, default=False
+    already_migrated: Mapped[bool | None] = mapped_column(
+        nullable=True, default=False
     )  # False = not migrated, True = migrated
 
     def to_settings(self):