tests-passed

orm/invite.py

@@ -1,7 +1,7 @@
 import enum
 
-from sqlalchemy import Column, ForeignKey, String
-from sqlalchemy.orm import relationship
+from sqlalchemy import ForeignKey, Index, Integer, String, UniqueConstraint
+from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from orm.base import BaseModel as Base
 
@@ -12,24 +12,33 @@ class InviteStatus(enum.Enum):
     REJECTED = "REJECTED"
 
     @classmethod
-    def from_string(cls, value: str) -> "Invite":
+    def from_string(cls, value: str) -> "InviteStatus":
         return cls(value)
 
 
 class Invite(Base):
     __tablename__ = "invite"
 
-    inviter_id = Column(ForeignKey("author.id"), primary_key=True)
-    author_id = Column(ForeignKey("author.id"), primary_key=True)
-    shout_id = Column(ForeignKey("shout.id"), primary_key=True)
-    status = Column(String, default=InviteStatus.PENDING.value)
+    id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    inviter_id: Mapped[int] = mapped_column(ForeignKey("author.id"))
+    author_id: Mapped[int] = mapped_column(ForeignKey("author.id"))
+    shout_id: Mapped[int] = mapped_column(ForeignKey("shout.id"))
+    status: Mapped[str] = mapped_column(String, default=InviteStatus.PENDING.value)
 
     inviter = relationship("Author", foreign_keys=[inviter_id])
     author = relationship("Author", foreign_keys=[author_id])
     shout = relationship("Shout")
 
-    def set_status(self, status: InviteStatus):
-        self.status = status.value  # type: ignore[assignment]
+    __table_args__ = (
+        UniqueConstraint(inviter_id, author_id, shout_id),
+        Index("idx_invite_inviter_id", "inviter_id"),
+        Index("idx_invite_author_id", "author_id"),
+        Index("idx_invite_shout_id", "shout_id"),
+        {"extend_existing": True},
+    )
+
+    def set_status(self, status: InviteStatus) -> None:
+        self.status = status.value
 
     def get_status(self) -> InviteStatus:
-        return InviteStatus.from_string(self.status)
+        return InviteStatus.from_string(str(self.status))