fix(enterprise): migrate telemetry models to SQLAlchemy 2.0 [3/13] (#13849)

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

enterprise/storage/telemetry_identity.py

@@ -5,13 +5,14 @@ for the OpenHands Enterprise Telemetry Service.
 """
 
 from datetime import UTC, datetime
-from typing import Optional
+from typing import Any
 
-from sqlalchemy import CheckConstraint, Column, DateTime, Integer, String
+from sqlalchemy import CheckConstraint, DateTime, String
+from sqlalchemy.orm import Mapped, mapped_column
 from storage.base import Base
 
 
-class TelemetryIdentity(Base):  # type: ignore
+class TelemetryIdentity(Base):
     """Stores persistent identity information for telemetry.
 
     This table is designed to contain exactly one row (enforced by database constraint)
@@ -25,15 +26,15 @@ class TelemetryIdentity(Base):  # type: ignore
     __tablename__ = 'telemetry_replicated_identity'
     __table_args__ = (CheckConstraint('id = 1', name='single_identity_row'),)
 
-    id = Column(Integer, primary_key=True, default=1)
-    customer_id = Column(String(255), nullable=True)
-    instance_id = Column(String(255), nullable=True)
-    created_at = Column(
+    id: Mapped[int] = mapped_column(primary_key=True, default=1)
+    customer_id: Mapped[str | None] = mapped_column(String(255), nullable=True)
+    instance_id: Mapped[str | None] = mapped_column(String(255), nullable=True)
+    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),
@@ -42,9 +43,9 @@ class TelemetryIdentity(Base):  # type: ignore
 
     def __init__(
         self,
-        customer_id: Optional[str] = None,
-        instance_id: Optional[str] = None,
-        **kwargs,
+        customer_id: str | None = None,
+        instance_id: str | None = None,
+        **kwargs: Any,
     ):
         """Initialize telemetry identity.
 
@@ -69,8 +70,8 @@ class TelemetryIdentity(Base):  # type: ignore
 
     def set_customer_info(
         self,
-        customer_id: Optional[str] = None,
-        instance_id: Optional[str] = None,
+        customer_id: str | None = None,
+        instance_id: str | None = None,
     ) -> None:
         """Update customer and instance identification information.
 

enterprise/storage/telemetry_metrics.py

@@ -6,13 +6,14 @@ and retry logic for the OpenHands Enterprise Telemetry Service.
 
 import uuid
 from datetime import UTC, datetime
-from typing import Any, Dict, Optional
+from typing import Any
 
-from sqlalchemy import JSON, Column, DateTime, Integer, String, Text
+from sqlalchemy import JSON, DateTime, String, Text
+from sqlalchemy.orm import Mapped, mapped_column
 from storage.base import Base
 
 
-class TelemetryMetrics(Base):  # type: ignore
+class TelemetryMetrics(Base):
     """Stores collected telemetry metrics with upload tracking.
 
     Each record represents a single metrics collection event with associated
@@ -21,23 +22,27 @@ class TelemetryMetrics(Base):  # type: ignore
 
     __tablename__ = 'telemetry_metrics'
 
-    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
-    collected_at = Column(
+    id: Mapped[str] = mapped_column(
+        String, primary_key=True, default=lambda: str(uuid.uuid4())
+    )
+    collected_at: Mapped[datetime] = mapped_column(
         DateTime(timezone=True),
         nullable=False,
         default=lambda: datetime.now(UTC),
         index=True,
     )
-    metrics_data = Column(JSON, nullable=False)
-    uploaded_at = Column(DateTime(timezone=True), nullable=True, index=True)
-    upload_attempts = Column(Integer, nullable=False, default=0)
-    last_upload_error = Column(Text, nullable=True)
-    created_at = Column(
+    metrics_data: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False)
+    uploaded_at: Mapped[datetime | None] = mapped_column(
+        DateTime(timezone=True), nullable=True, index=True
+    )
+    upload_attempts: Mapped[int] = mapped_column(nullable=False, default=0)
+    last_upload_error: Mapped[str | None] = mapped_column(Text, nullable=True)
+    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),
@@ -46,9 +51,9 @@ class TelemetryMetrics(Base):  # type: ignore
 
     def __init__(
         self,
-        metrics_data: Dict[str, Any],
-        collected_at: Optional[datetime] = None,
-        **kwargs,
+        metrics_data: dict[str, Any],
+        collected_at: datetime | None = None,
+        **kwargs: Any,
     ):
         """Initialize a new telemetry metrics record.