1 files changed, 54 insertions, 0 deletions
diff --git a/jb/managers/bonus.py b/jb/managers/bonus.py
new file mode 100644
index 0000000..0cb8b02
--- /dev/null
+++ b/jb/managers/bonus.py
@@ -0,0 +1,54 @@
+from typing import List
+
+from psycopg import sql
+
+from jb.managers import PostgresManager
+from jb.models.bonus import Bonus
+
+
+class BonusManager(PostgresManager):
+
+    def create(self, bonus: Bonus) -> None:
+        assert bonus.id is None
+        data = bonus.to_postgres()
+        query = sql.SQL(
+            """
+        INSERT INTO mtwerk_bonus
+        (payout_event_id, amt_worker_id, amount, grant_time, assignment_id, reason)
+        VALUES (
+            %(payout_event_id)s,
+            %(amt_worker_id)s,
+            %(amount)s,
+            %(grant_time)s,
+            (
+                SELECT id
+                FROM mtwerk_assignment
+                WHERE amt_assignment_id = %(amt_assignment_id)s
+                LIMIT 1
+            ),
+            %(reason)s
+        )
+        RETURNING id, assignment_id;
+        """
+        )
+
+        with self.pg_config.make_connection() as conn:
+            with conn.cursor() as c:
+                c.execute(query, data)
+                res = c.fetchone()
+                conn.commit()
+        bonus.id = res["id"]
+        bonus.assignment_id = res["assignment_id"]
+        return None
+
+    def filter(self, amt_assignment_id: str) -> List[Bonus]:
+        res = self.pg_config.execute_sql_query(
+            """
+            SELECT mb.*, ma.amt_assignment_id
+            FROM mtwerk_bonus mb
+            JOIN mtwerk_assignment ma ON ma.id = mb.assignment_id
+            WHERE amt_assignment_id = %(amt_assignment_id)s;
+        """,
+            params={"amt_assignment_id": amt_assignment_id},
+        )
+        return [Bonus.from_postgres(x) for x in res]