Merges pull request #3

Off of Amazon!!!

1 files changed, 10 insertions, 14 deletions

diff --git a/jb/managers/bonus.py b/jb/managers/bonus.py
index 89b81f0..b649103 100644
--- a/jb/managers/bonus.py
+++ b/jb/managers/bonus.py
@@ -1,8 +1,8 @@
-from typing import List, Any
+from typing import Any
 
 from psycopg import sql
 
-from jb.managers import PostgresManager
+from jb.managers.base import PostgresManager
 from jb.models.bonus import Bonus
 
 
@@ -11,8 +11,7 @@ class BonusManager(PostgresManager):
     def create(self, bonus: Bonus) -> None:
         assert bonus.id is None
         data = bonus.to_postgres()
-        query = sql.SQL(
-            """
+        query = sql.SQL("""
         INSERT INTO mtwerk_bonus
         (payout_event_id, amt_worker_id, amount, grant_time, assignment_id, reason)
         VALUES (
@@ -29,20 +28,17 @@ class BonusManager(PostgresManager):
             %(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()
+        with self.pg_config.make_connection() as conn, conn.cursor() as c:
+            c.execute(query, data)
+            res = c.fetchone()
+            conn.commit()
         bonus.id = res["id"]  # type: ignore
         bonus.assignment_id = res["assignment_id"]  # type: ignore
-        return None
 
-    def filter(self, amt_assignment_id: str) -> List[Bonus]:
-        res: List[Any] = self.pg_config.execute_sql_query(
+    def filter(self, amt_assignment_id: str) -> list[Bonus]:
+        res: list[Any] = self.pg_config.execute_sql_query(
             """
             SELECT mb.*, ma.amt_assignment_id
             FROM mtwerk_bonus mb