Merges pull request #3

Off of Amazon!!!

1 files changed, 46 insertions, 73 deletions

diff --git a/jb/managers/hit.py b/jb/managers/hit.py
index 63af2d4..bbae92b 100644
--- a/jb/managers/hit.py
+++ b/jb/managers/hit.py
@@ -1,11 +1,10 @@
 from datetime import datetime, timezone
-from typing import Optional, List
 
 from psycopg import sql
 
-from jb.managers import PostgresManager
+from jb.managers.base import PostgresManager
 from jb.models.definitions import HitStatus
-from jb.models.hit import HitQuestion, HitType, Hit
+from jb.models.hit import Hit, HitQuestion, HitType
 
 
 class HitQuestionManager(PostgresManager):
@@ -13,21 +12,17 @@ class HitQuestionManager(PostgresManager):
     def create(self, question: HitQuestion) -> None:
         assert question.id is None
         data = question.to_postgres()
-        query = sql.SQL(
-            """
+        query = sql.SQL("""
         INSERT INTO mtwerk_question (url, height)
         VALUES (%(url)s, %(height)s)
         RETURNING id;
-        """
-        )
+        """)
 
-        with self.pg_config.make_connection() as conn:
-            with conn.cursor() as c:
-                c.execute(query, data)
-                pk = c.fetchone()["id"]  # type: ignore
-                conn.commit()
+        with self.pg_config.make_connection() as conn, conn.cursor() as c:
+            c.execute(query, data)
+            pk = c.fetchone()["id"]  # type: ignore
+            conn.commit()
         question.id = pk
-        return None
 
     def get_by_id(self, question_id: int) -> HitQuestion:
         res = self.pg_config.execute_sql_query(
@@ -42,7 +37,7 @@ class HitQuestionManager(PostgresManager):
         assert len(res) == 1
         return HitQuestion.model_validate(res[0])
 
-    def get_by_values_if_exists(self, url: str, height: int) -> Optional[HitQuestion]:
+    def get_by_values_if_exists(self, url: str, height: int) -> HitQuestion | None:
         res = self.pg_config.execute_sql_query(
             """
             SELECT *
@@ -73,8 +68,7 @@ class HitTypeManager(PostgresManager):
         assert hit_type.amt_hit_type_id is not None
 
         data = hit_type.to_postgres()
-        query = sql.SQL(
-            """
+        query = sql.SQL("""
         INSERT INTO mtwerk_hittype (
             amt_hit_type_id,
             title,
@@ -96,27 +90,21 @@ class HitTypeManager(PostgresManager):
             %(min_active)s
         )
         RETURNING id;
-        """
-        )
+        """)
 
-        with self.pg_config.make_connection() as conn:
-            with conn.cursor() as c:
-                c.execute(query, data)
-                pk = c.fetchone()["id"]  # type: ignore
-                conn.commit()
+        with self.pg_config.make_connection() as conn, conn.cursor() as c:
+            c.execute(query, data)
+            pk = c.fetchone()["id"]  # type: ignore
+            conn.commit()
         hit_type.id = pk
 
-        return None
-
-    def filter_active(self) -> List[HitType]:
-        res = self.pg_config.execute_sql_query(
-            """
+    def filter_active(self) -> list[HitType]:
+        res = self.pg_config.execute_sql_query("""
             SELECT *
             FROM mtwerk_hittype
             WHERE min_active > 0
             LIMIT 50
-        """
-        )
+        """)
 
         if len(res) == 50:
             raise ValueError("Too many HitTypes!")
@@ -135,7 +123,7 @@ class HitTypeManager(PostgresManager):
         assert len(res) == 1
         return HitType.from_postgres(res[0])
 
-    def get_if_exists(self, amt_hit_type_id: str) -> Optional[HitType]:
+    def get_if_exists(self, amt_hit_type_id: str) -> HitType | None:
         try:
             return self.get(amt_hit_type_id=amt_hit_type_id)
         except AssertionError:
@@ -151,22 +139,18 @@ class HitTypeManager(PostgresManager):
 
     def set_min_active(self, hit_type: HitType) -> None:
         assert hit_type.id, "must be in the db first!"
-        query = sql.SQL(
-            """
+        query = sql.SQL("""
         UPDATE mtwerk_hittype
         SET min_active = %(min_active)s
         WHERE id = %(id)s
-        """
-        )
+        """)
         data = {"id": hit_type.id, "min_active": hit_type.min_active}
-        with self.pg_config.make_connection() as conn:
-            with conn.cursor() as c:
-                c.execute(query, data)
-                conn.commit()
-                row_cnt = c.rowcount
+        with self.pg_config.make_connection() as conn, conn.cursor() as c:
+            c.execute(query, data)
+            conn.commit()
+            row_cnt = c.rowcount
 
         assert row_cnt == 1, f"Expected 1 row updated, got {row_cnt}"
-        return None
 
 
 class HitManager(PostgresManager):
@@ -175,8 +159,7 @@ class HitManager(PostgresManager):
         assert hit.amt_hit_id is not None
         assert hit.id is None
         data = hit.to_postgres()
-        query = sql.SQL(
-            """
+        query = sql.SQL("""
         INSERT INTO mtwerk_hit (
             amt_hit_id,
             hit_type_id,
@@ -210,38 +193,32 @@ class HitManager(PostgresManager):
             %(assignment_available_count)s
         )
         RETURNING id;
-        """
-        )
+        """)
 
-        with self.pg_config.make_connection() as conn:
-            with conn.cursor() as c:
-                c.execute(query, data)
-                pk = c.fetchone()["id"]  # type: ignore
-                conn.commit()
+        with self.pg_config.make_connection() as conn, conn.cursor() as c:
+            c.execute(query, data)
+            pk = c.fetchone()["id"]  # type: ignore
+            conn.commit()
         hit.id = pk
         return hit
 
     def update_status(self, amt_hit_id: str, hit_status: HitStatus):
         now = datetime.now(tz=timezone.utc)
-        query = sql.SQL(
-            """
+        query = sql.SQL("""
         UPDATE mtwerk_hit
         SET status = %(status)s, modified_at = %(modified_at)s
         WHERE amt_hit_id = %(amt_hit_id)s;
-        """
-        )
+        """)
 
         data = {
             "amt_hit_id": amt_hit_id,
             "status": hit_status.value,
             "modified_at": now,
         }
-        with self.pg_config.make_connection() as conn:
-            with conn.cursor() as c:
-                c.execute(query, data)
-                conn.commit()
-                assert c.rowcount == 1, c.rowcount
-        return None
+        with self.pg_config.make_connection() as conn, conn.cursor() as c:
+            c.execute(query, data)
+            conn.commit()
+            assert c.rowcount == 1, c.rowcount
 
     def update_hit(self, hit: Hit):
         hit.modified_at = datetime.now(tz=timezone.utc)
@@ -255,8 +232,7 @@ class HitManager(PostgresManager):
             "amt_hit_id",
             "modified_at",
         }
-        query = sql.SQL(
-            """
+        query = sql.SQL("""
         UPDATE mtwerk_hit
         SET status = %(status)s, review_status = %(review_status)s,
             assignment_pending_count = %(assignment_pending_count)s,
@@ -265,17 +241,14 @@ class HitManager(PostgresManager):
             modified_at = %(modified_at)s
         WHERE amt_hit_id = %(amt_hit_id)s
         RETURNING id;
-        """
-        )
+        """)
 
         data = hit.model_dump(mode="json", include=fields)
-        with self.pg_config.make_connection() as conn:
-            with conn.cursor() as c:
-                c.execute(query, data)
-                conn.commit()
-                assert c.rowcount == 1, c.rowcount
-                hit.id = c.fetchone()["id"]  # type: ignore
-        return None
+        with self.pg_config.make_connection() as conn, conn.cursor() as c:
+            c.execute(query, data)
+            conn.commit()
+            assert c.rowcount == 1, c.rowcount
+            hit.id = c.fetchone()["id"]  # type: ignore
 
     def get_from_amt_id(self, amt_hit_id: str) -> Hit:
         res = self.pg_config.execute_sql_query(
@@ -324,7 +297,7 @@ class HitManager(PostgresManager):
 
         return Hit.from_postgres(res)
 
-    def get_from_amt_id_if_exists(self, amt_hit_id: str) -> Optional[Hit]:
+    def get_from_amt_id_if_exists(self, amt_hit_id: str) -> Hit | None:
         try:
             return self.get_from_amt_id(amt_hit_id=amt_hit_id)
 
@@ -333,7 +306,7 @@ class HitManager(PostgresManager):
 
     def get_active_count(self, hit_type_id: int) -> int:
         return self.pg_config.execute_sql_query(
-            """
+            query="""
             SELECT COUNT(1) as active_count
             FROM mtwerk_hit
             WHERE status = %(status)s