Models, Project files, some pytests, requirements.. etcHEADmaster

1 files changed, 48 insertions, 0 deletions

diff --git a/jb/models/bonus.py b/jb/models/bonus.py
new file mode 100644
index 0000000..564a32d
--- /dev/null
+++ b/jb/models/bonus.py
@@ -0,0 +1,48 @@
+from typing import Optional, Dict
+
+from pydantic import BaseModel, Field, ConfigDict, PositiveInt
+from typing_extensions import Self
+
+from jb.models.currency import USDCent
+from jb.models.custom_types import AMTBoto3ID, AwareDatetimeISO, UUIDStr
+from jb.models.definitions import PayoutStatus
+
+
+class Bonus(BaseModel):
+    """
+    A Bonus is created (in our DB) ONLY associated with an APPROVED
+        thl-payout-event, AFTER the bonus has actually been sent to
+        the worker.
+    We have the payout_event uuid as the unique request token to make
+        sure it only gets sent once (param in the boto request).
+    """
+
+    model_config = ConfigDict(
+        extra="forbid",
+        validate_assignment=True,
+    )
+    id: Optional[PositiveInt] = Field(default=None)
+    assignment_id: Optional[PositiveInt] = Field(default=None)
+
+    amt_worker_id: str = Field(min_length=3, max_length=50)
+    amt_assignment_id: AMTBoto3ID = Field()
+
+    amount: USDCent = Field()
+    reason: str = Field(min_length=5)
+    grant_time: AwareDatetimeISO = Field()
+
+    # -- GRL Specific ---
+    payout_event_id: UUIDStr = Field()
+    # created: Optional[AwareDatetimeISO] = Field(default=None)
+
+    def to_postgres(self):
+        d = self.model_dump(mode="json")
+        d["amount"] = self.amount.to_usd()
+        return d
+
+    @classmethod
+    def from_postgres(cls, data: Dict) -> Self:
+        data["amount"] = USDCent(round(data["amount"] * 100))
+        fields = set(cls.model_fields.keys())
+        data = {k: v for k, v in data.items() if k in fields}
+        return cls.model_validate(data)