diff options
Diffstat (limited to 'jb/managers')
| -rw-r--r-- | jb/managers/__init__.py | 23 | ||||
| -rw-r--r-- | jb/managers/amt.py | 47 | ||||
| -rw-r--r-- | jb/managers/assignment.py | 87 | ||||
| -rw-r--r-- | jb/managers/base.py | 16 | ||||
| -rw-r--r-- | jb/managers/bonus.py | 24 | ||||
| -rw-r--r-- | jb/managers/email_manager.py | 65 | ||||
| -rw-r--r-- | jb/managers/gr_api.py | 160 | ||||
| -rw-r--r-- | jb/managers/hit.py | 119 | ||||
| -rw-r--r-- | jb/managers/thl.py | 74 | ||||
| -rw-r--r-- | jb/managers/worker.py | 6 |
10 files changed, 364 insertions, 257 deletions
diff --git a/jb/managers/__init__.py b/jb/managers/__init__.py index e99569a..e69de29 100644 --- a/jb/managers/__init__.py +++ b/jb/managers/__init__.py @@ -1,23 +0,0 @@ -from enum import IntEnum -from typing import Collection - -from generalresearchutils.pg_helper import PostgresConfig - - -class Permission(IntEnum): - READ = 1 - UPDATE = 2 - CREATE = 3 - DELETE = 4 - - -class PostgresManager: - def __init__( - self, - pg_config: PostgresConfig, - permissions: Collection[Permission] = None, # type: ignore - **kwargs, # type: ignore - ): - super().__init__(**kwargs) - self.pg_config = pg_config - self.permissions = set(permissions) if permissions else set() diff --git a/jb/managers/amt.py b/jb/managers/amt.py index 88ae07d..17e5630 100644 --- a/jb/managers/amt.py +++ b/jb/managers/amt.py @@ -1,25 +1,24 @@ -import logging -from datetime import timezone, datetime -from typing import Tuple, Optional, List, Dict, Any +from datetime import datetime, timezone +from typing import Any import botocore.exceptions -from generalresearchutils.currency import USDCent +from generalresearch.currency import USDCent from mypy_boto3_mturk import MTurkClient from mypy_boto3_mturk.type_defs import ( AssignmentTypeDef, BonusPaymentTypeDef, CreateHITTypeResponseTypeDef, - GetHITResponseTypeDef, CreateHITWithHITTypeResponseTypeDef, + GetHITResponseTypeDef, ) from pydantic import ValidationError from jb.config import TOPIC_ARN -from jb.models import AMTAccount +from jb.models.amt import AMTAccount from jb.models.assignment import Assignment from jb.models.bonus import Bonus from jb.models.definitions import HitStatus -from jb.models.hit import HitType, HitQuestion, Hit +from jb.models.hit import Hit, HitQuestion, HitType REJECT_MESSAGE_UNKNOWN_ASSIGNMENT = "Unknown assignment" REJECT_MESSAGE_NO_WORK = "Assignment was submitted with no attempted work." @@ -56,7 +55,7 @@ class AMTManager: } ) - def get_hit_if_exists(self, amt_hit_id: str) -> Tuple[Optional[Hit], Optional[str]]: + def get_hit_if_exists(self, amt_hit_id: str) -> tuple[Hit | None, str | None]: try: res: GetHITResponseTypeDef = self.amt_client.get_hit(HITId=amt_hit_id) @@ -78,7 +77,7 @@ class AMTManager: return HitStatus.Disposed else: - logging.warning(msg) + LOG.warning(msg) return HitStatus.Unassignable return res.status @@ -137,7 +136,7 @@ class AMTManager: # Baddies have been known to submit assignments with purposely # malformed "answer" (xml) section, which will raise # a pydantic validation error. Try to parse again with no Answer. - logging.exception(e) + LOG.exception(e) ass_res["Answer"] = None # If it wasn't the Answer that caused the ValidationError, it'll raise again assignment = Assignment.from_amt_get_assignment(ass_res) @@ -146,13 +145,13 @@ class AMTManager: assert assignment.id is None return assignment - def get_assignment_if_exists(self, amt_assignment_id: str) -> Optional[Assignment]: + def get_assignment_if_exists(self, amt_assignment_id: str) -> Assignment | None: expected_err_msg = f"Assignment {amt_assignment_id} does not exist" try: return self.get_assignment(amt_assignment_id=amt_assignment_id) except botocore.exceptions.ClientError as e: - logging.warning(e) + LOG.warning(e) error_code = e.response["Error"]["Code"] error_msg = e.response["Error"]["Message"] if error_code == "RequestError" and expected_err_msg in error_msg: @@ -161,7 +160,7 @@ class AMTManager: def reject_assignment_if_possible( self, amt_assignment_id: str, msg: str = REJECT_MESSAGE_UNKNOWN_ASSIGNMENT - ) -> Optional[Dict[str, Any]]: + ) -> dict[str, Any] | None: # Unclear to me when this would fail try: @@ -170,7 +169,7 @@ class AMTManager: ) except botocore.exceptions.ClientError as e: - logging.warning(e) + LOG.warning(e) return None def approve_assignment_if_possible( @@ -178,7 +177,7 @@ class AMTManager: amt_assignment_id: str, msg: str = APPROVAL_MESSAGE, override_rejection: bool = False, - ) -> Optional[Dict[str, Any]]: + ) -> dict[str, Any] | None: # Unclear to me when this would fail try: @@ -189,7 +188,7 @@ class AMTManager: ) except botocore.exceptions.ClientError as e: - logging.warning(e) + LOG.warning(e) return None def update_hit_review_status(self, amt_hit_id: str, revert: bool = False) -> None: @@ -198,7 +197,7 @@ class AMTManager: self.amt_client.update_hit_review_status(HITId=amt_hit_id, Revert=revert) except botocore.exceptions.ClientError as e: - logging.warning(f"{amt_hit_id=}, {e}") + LOG.warning(f"{amt_hit_id=}, {e}") error_msg = e.response["Error"]["Message"] if "does not exist" in error_msg: @@ -207,8 +206,6 @@ class AMTManager: # elif "This HIT is currently in the state 'Reviewing'" in error_msg: # logging.warning(error_msg) - return None - def send_bonus( self, amt_worker_id: str, @@ -216,7 +213,7 @@ class AMTManager: amt_assignment_id: str, reason: str, unique_request_token: str, - ) -> Optional[Dict[str, Any]]: + ) -> dict[str, Any] | None: try: return self.amt_client.send_bonus( WorkerId=amt_worker_id, @@ -227,14 +224,12 @@ class AMTManager: ) except botocore.exceptions.ClientError as e: - logging.warning(f"{amt_worker_id=} {amt_assignment_id=}, {e}") + LOG.warning(f"{amt_worker_id=} {amt_assignment_id=}, {e}") return None - def get_bonus( - self, amt_assignment_id: str, payout_event_id: str - ) -> Optional[Bonus]: + def get_bonus(self, amt_assignment_id: str, payout_event_id: str) -> Bonus | None: - res: List[BonusPaymentTypeDef] = self.amt_client.list_bonus_payments( + res: list[BonusPaymentTypeDef] = self.amt_client.list_bonus_payments( AssignmentId=amt_assignment_id )["BonusPayments"] @@ -268,5 +263,3 @@ class AMTManager: self.amt_client.update_expiration_for_hit( HITId=hit["HITId"], ExpireAt=now ) - - return None diff --git a/jb/managers/assignment.py b/jb/managers/assignment.py index f6aa2ce..2740aee 100644 --- a/jb/managers/assignment.py +++ b/jb/managers/assignment.py @@ -1,11 +1,10 @@ from datetime import datetime, timezone -from typing import Optional from psycopg import sql from pydantic import NonNegativeInt, PositiveInt -from jb.managers import PostgresManager -from jb.models.assignment import AssignmentStub, Assignment +from jb.managers.base import PostgresManager +from jb.models.assignment import Assignment, AssignmentStub from jb.models.definitions import AssignmentStatus @@ -14,8 +13,7 @@ class AssignmentManager(PostgresManager): def create_stub(self, stub: AssignmentStub) -> None: assert stub.id is None data = stub.to_postgres() - query = sql.SQL( - """ + query = sql.SQL(""" INSERT INTO mtwerk_assignment (amt_assignment_id, amt_worker_id, status, created_at, modified_at, hit_id) @@ -23,16 +21,13 @@ class AssignmentManager(PostgresManager): (%(amt_assignment_id)s, %(amt_worker_id)s, %(status)s, %(created_at)s, %(modified_at)s, %(hit_id)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() stub.id = pk - return None def create(self, assignment: Assignment) -> None: # Typically this is NOT used (we'd create the stub when HIT is @@ -42,8 +37,7 @@ class AssignmentManager(PostgresManager): assert assignment.id is None data = assignment.to_postgres() - query = sql.SQL( - """ + query = sql.SQL(""" INSERT INTO mtwerk_assignment (amt_assignment_id, amt_worker_id, status, created_at, modified_at, hit_id, @@ -57,16 +51,13 @@ class AssignmentManager(PostgresManager): %(approval_time)s, %(rejection_time)s, %(requester_feedback)s, %(tsid)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() assignment.id = pk - return None def get_stub(self, amt_assignment_id: str) -> AssignmentStub: res = self.pg_config.execute_sql_query( @@ -85,7 +76,7 @@ class AssignmentManager(PostgresManager): assert len(res) == 1 return AssignmentStub.model_validate(res[0]) - def get_stub_if_exists(self, amt_assignment_id: str) -> Optional[AssignmentStub]: + def get_stub_if_exists(self, amt_assignment_id: str) -> AssignmentStub | None: try: return self.get_stub(amt_assignment_id=amt_assignment_id) except AssertionError: @@ -121,8 +112,7 @@ class AssignmentManager(PostgresManager): "amt_assignment_id": assignment.amt_assignment_id, "modified_at": now, } - query = sql.SQL( - """ + query = sql.SQL(""" UPDATE mtwerk_assignment SET submit_time = %(submit_time)s, auto_approval_time = %(auto_approval_time)s, @@ -130,15 +120,12 @@ class AssignmentManager(PostgresManager): tsid = %(tsid)s, modified_at = %(modified_at)s WHERE amt_assignment_id = %(amt_assignment_id)s - """ - ) + """) # We force this to fail if the assignment doesn't already exist in the db - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, data) - assert c.rowcount == 1, f"Expected 1 row, got {c.rowcount}" - conn.commit() - return None + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, data) + assert c.rowcount == 1, f"Expected 1 row, got {c.rowcount}" + conn.commit() def reject(self, assignment: Assignment) -> None: assert assignment.status == AssignmentStatus.Rejected @@ -156,8 +143,7 @@ class AssignmentManager(PostgresManager): "accept_time": assignment.accept_time, "modified_at": now, } - query = sql.SQL( - """ + query = sql.SQL(""" UPDATE mtwerk_assignment SET submit_time = %(submit_time)s, rejection_time = %(rejection_time)s, @@ -167,15 +153,12 @@ class AssignmentManager(PostgresManager): accept_time = %(accept_time)s, modified_at = %(modified_at)s WHERE amt_assignment_id = %(amt_assignment_id)s - """ - ) + """) # We force this to fail if the assignment doesn't already exist in the db - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, data) - assert c.rowcount == 1, f"Expected 1 row, got {c.rowcount}" - conn.commit() - return None + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, data) + assert c.rowcount == 1, f"Expected 1 row, got {c.rowcount}" + conn.commit() def approve(self, assignment: Assignment) -> None: assert assignment.status == AssignmentStatus.Approved @@ -194,8 +177,7 @@ class AssignmentManager(PostgresManager): "accept_time": assignment.accept_time, "modified_at": now, } - query = sql.SQL( - """ + query = sql.SQL(""" UPDATE mtwerk_assignment SET submit_time = %(submit_time)s, approval_time = %(approval_time)s, @@ -205,15 +187,12 @@ class AssignmentManager(PostgresManager): accept_time = %(accept_time)s, modified_at = %(modified_at)s WHERE amt_assignment_id = %(amt_assignment_id)s - """ - ) + """) # We force this to fail if the assignment doesn't already exist in the db - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, data) - assert c.rowcount == 1, f"Expected 1 row, got {c.rowcount}" - conn.commit() - return None + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, data) + assert c.rowcount == 1, f"Expected 1 row, got {c.rowcount}" + conn.commit() def missing_tsid_count( self, amt_worker_id: str, lookback_hrs: PositiveInt = 24 diff --git a/jb/managers/base.py b/jb/managers/base.py new file mode 100644 index 0000000..4b0637f --- /dev/null +++ b/jb/managers/base.py @@ -0,0 +1,16 @@ +from collections.abc import Collection + +from generalresearch.managers.base import Permission +from generalresearch.pg_helper import PostgresConfig + + +class PostgresManager: + def __init__( + self, + pg_config: PostgresConfig, + permissions: Collection[Permission] | None = None, + **kwargs, # type: ignore + ): + super().__init__(**kwargs) + self.pg_config = pg_config + self.permissions = set(permissions) if permissions else set() 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 diff --git a/jb/managers/email_manager.py b/jb/managers/email_manager.py new file mode 100644 index 0000000..7298c60 --- /dev/null +++ b/jb/managers/email_manager.py @@ -0,0 +1,65 @@ +import requests +from generalresearch.config import is_debug + +from jb.config import settings + +MAUTIC_BASE_URL = "https://mail.jamesbillings67.com" +EMAIL_TEMPLATE_ID = 1 + +assert settings.mautic_api_key +auth_headers = {"Authorization": f"Basic {settings.mautic_api_key.get_secret_value()}"} + + +def get_or_create_contact(email: str, amt_worker_id: str | None = None) -> int: + body = {"email": email} + if amt_worker_id: + body["amt_worker_id"] = amt_worker_id + res = requests.post( + url=f"{MAUTIC_BASE_URL}/api/contacts/new", + json=body, + headers=auth_headers, + ).json() + return int(res["contact"]["id"]) + + +def send_login_email_from_url(mautic_url: str, magic_link: str) -> None: + if is_debug(): + print("MAGIC_LINK: ", magic_link) + return + + email_tokens = { + "magic_link": magic_link, + } + body = {"tokens": email_tokens} + response = requests.post(url=mautic_url, json=body, headers=auth_headers) + + try: + response.raise_for_status() + except requests.exceptions.HTTPError: + print(f"Failed to send email. Status code: {response.status_code}") + print(response.text) + raise + d = response.json() + + assert d.get("success"), f"Failed to send email: {d.get('failed')}" + print("Email sent successfully") + + +def send_login_email(email: str, magic_token: str) -> None: + contact_id = get_or_create_contact(email=email) + mautic_url = ( + f"{MAUTIC_BASE_URL}/api/emails/{EMAIL_TEMPLATE_ID}/contact/{contact_id}/send" + ) + magic_link = f"{settings.base_url}auth/magic-link/?token={magic_token}" + send_login_email_from_url(mautic_url, magic_link) + + +def send_amt_link_email(email: str, magic_token: str) -> None: + # Don't actually associate the email with the worker ID + # until they click the link + contact_id = get_or_create_contact(email=email) + mautic_url = ( + f"{MAUTIC_BASE_URL}/api/emails/{EMAIL_TEMPLATE_ID}/contact/{contact_id}/send" + ) + magic_link = f"{settings.base_url}auth/link-amt/?token={magic_token}" + send_login_email_from_url(mautic_url, magic_link) diff --git a/jb/managers/gr_api.py b/jb/managers/gr_api.py new file mode 100644 index 0000000..f0be1a7 --- /dev/null +++ b/jb/managers/gr_api.py @@ -0,0 +1,160 @@ +"""Client for General Research's product-user API.""" + +import logging +from typing import Any + +import requests + +from jb.models.auth import User + +logger = logging.getLogger("amtjb") + + +class GRApiError(RuntimeError): + """The General Research API could not satisfy a request.""" + + +class GRApiNotFoundError(GRApiError): + """The requested General Research API resource does not exist.""" + + +class GRApiManager: + def __init__( + self, + base_url: str, + token: str, + product_id: str, + timeout: float = 5.0, + session: requests.Session | None = None, + ) -> None: + if not token: + raise ValueError("General Research API token must not be empty") + + self.base_url = base_url.rstrip("/") + self.product_id = product_id + self.timeout = timeout + self.session = session or requests.Session() + self.session.headers.update( + { + "Authorization": token, + "Accept": "application/json", + "Content-Type": "application/json", + } + ) + + def _request(self, method: str, url: str, **kwargs: Any) -> requests.Response: + try: + response = self.session.request(method, url, timeout=self.timeout, **kwargs) + response.raise_for_status() + return response + except requests.HTTPError as exc: + if exc.response is not None and exc.response.status_code == 404: + raise GRApiNotFoundError( + f"General Research API resource not found: {method} {url}" + ) from exc + raise GRApiError( + f"General Research API request failed: {method} {url}" + ) from exc + except requests.RequestException as exc: + raise GRApiError( + f"General Research API request failed: {method} {url}" + ) from exc + + def _parse_user_response(self, res: dict[str, Any]) -> User: + return User.model_validate( + { + "product_user_id": res["product_user_id"], + "email": res["metadata"].get("email_address"), + "display_name": res["metadata"].get("display_name"), + "blocked": res["blocked"], + } + ) + + def ensure_user_exists(self, user: User) -> User: + """Idempotently create the product user if it does not already exist.""" + url = f"{self.base_url}/{self.product_id}/user/{user.product_user_id}/" + res = self._request("PUT", url).json() + if res["metadata"].get("email_address") is None: + self.set_user_email(user) + return self.get_user(user.product_user_id) + user_thl = self._parse_user_response(res) + if user_thl.email != user.email: + raise ValueError( + f"user {user.product_user_id} already exists with email {user_thl.email}" + ) + return user_thl + + def get_user(self, product_user_id: str) -> User: + """Retrieve the user's email address and display name.""" + url = f"{self.base_url}/{self.product_id}/user/{product_user_id}/" + res = self._request("GET", url).json() + return self._parse_user_response(res) + + def get_user_by_email(self, email: str) -> User: + user = User.model_validate({"email": email}) + return self.get_user(user.product_user_id) + + def set_user_email(self, user: User) -> None: + """This should only be called once per user upon account creation. + A user cannot change their email address.""" + url = f"{self.base_url}/{self.product_id}/user/{user.product_user_id}/metadata/" + _ = self._request( + "PATCH", + url, + json={"email_address": str(user.email)}, + ).json() + + def set_user_display_name(self, user: User) -> User: + """Can be called as many times as needed. Multiple + users can have the same display name.""" + if user.display_name is None: + return user + url = f"{self.base_url}/{self.product_id}/user/{user.product_user_id}/metadata/" + _ = self._request( + "PATCH", + url, + json={"display_name": user.display_name}, + ).json() + return self.get_user(user.product_user_id) + + def transition_user_from_amt(self, user: User, amt_worker_id: str) -> User: + """This should only be called once upon transition from an + AMT account to a General Research account.""" + url = f"{self.base_url}/{self.product_id}/user/{amt_worker_id}/" + try: + self._request("PATCH", url, json={"product_user_id": user.product_user_id}) + except GRApiNotFoundError as exc: + raise ValueError(f"User {amt_worker_id} does not exist") from exc + except GRApiError as exc: + http_error = exc.__cause__ + response = ( + http_error.response + if isinstance(http_error, requests.HTTPError) + else None + ) + if response is not None and response.status_code == 400: + try: + detail = response.json().get("detail") + except (ValueError, AttributeError): + detail = None + if detail == "Unable to update User": + raise ValueError( + "unable to update user, probably another user already " + "exists with this email" + ) from exc + raise + self.set_user_email(user) + transitioned_user = self.get_user(user.product_user_id) + logger.warning( + "Transitioned product user from AMT worker %s to %s with email %s", + amt_worker_id, + transitioned_user.product_user_id, + transitioned_user.email, + extra={ + "event": "transition_user_from_amt", + "amt_worker_id": amt_worker_id, + "product_user_id": transitioned_user.product_user_id, + "email": str(transitioned_user.email), + }, + ) + return transitioned_user 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 diff --git a/jb/managers/thl.py b/jb/managers/thl.py index 6e8effc..85e0697 100644 --- a/jb/managers/thl.py +++ b/jb/managers/thl.py @@ -1,64 +1,18 @@ -from generalresearchutils.models.thl.payout import UserPayoutEvent -from generalresearchutils.models.thl.task_status import TaskStatusResponse -from generalresearchutils.models.thl.wallet.cashout_method import ( - CashoutRequestResponse, +import requests +from generalresearch.currency import USDCent +from generalresearch.models.thl.definitions import PayoutStatus +from generalresearch.models.thl.payout import UserPayoutEvent +from generalresearch.models.thl.task_status import TaskStatusResponse +from generalresearch.models.thl.wallet.cashout_method import ( CashoutRequestInfo, + CashoutRequestResponse, ) -from generalresearchutils.models.thl.user_profile import UserProfile -from generalresearchutils.currency import USDCent - from jb.config import settings +from jb.models.auth import User -from generalresearchutils.models.thl.definitions import PayoutStatus - - -from typing import Optional -import requests - -# TODO: Organize this more with other endpoints (offerwall, cashout -# requests/approvals, etc). - - -def get_user_profile(amt_worker_id: str) -> UserProfile: - url = f"{settings.fsb_host}{settings.product_id}/user/{amt_worker_id}/profile/" - res = requests.get(url).json() - - if res.get("detail") == "user not found": - raise ValueError("user not found") - user_profile = res["user_profile"] - # todo: these are computed fields, need a UserProfile parser - user_profile.pop("email_md5", None) - user_profile.pop("email_sha1", None) - user_profile.pop("email_sha256", None) - # todo: this contains computed fields inside each streak object - user_profile.pop("streaks", None) - # todo: this shouldn't be in here anyways ---v - user_profile["user"].pop("id", None) - - return UserProfile.model_validate(user_profile) - - -def get_user_blocked(amt_worker_id: str) -> bool: - # Not blocked if None - res = get_user_profile(amt_worker_id=amt_worker_id) - return res.user.blocked if res.user.blocked is not None else False - - -def get_user_blocked_or_not_exists(amt_worker_id: str) -> Optional[bool]: - try: - res = get_user_profile(amt_worker_id=amt_worker_id) - return res.user.blocked if res.user.blocked is not None else False - - except ValueError as e: - if e.args[0] == "user not found": - return True - - return None - - -def get_task_status(tsid: str) -> Optional[TaskStatusResponse]: +def get_task_status(tsid: str) -> TaskStatusResponse | None: url = f"{settings.fsb_host}{settings.product_id}/status/{tsid}/" d = requests.get(url).json() if d.get("msg") == "invalid tsid": @@ -68,18 +22,14 @@ def get_task_status(tsid: str) -> Optional[TaskStatusResponse]: def user_cashout_request( - amt_worker_id: str, amount: USDCent, cashout_method_id: str + user: User, amount: USDCent, cashout_method_id: str ) -> CashoutRequestInfo: - assert cashout_method_id in { - settings.amt_assignment_cashout_method, - settings.amt_bonus_cashout_method, - } assert isinstance(amount, USDCent) assert USDCent(0) < amount < USDCent(10_00) url = f"{settings.fsb_host}{settings.product_id}/cashout/" body: dict[str, str | int] = { - "bpuid": amt_worker_id, + "bpuid": user.product_user_id, "amount": int(amount), "cashout_method_id": cashout_method_id, } @@ -110,7 +60,7 @@ def get_wallet_balance(amt_worker_id: str) -> USDCent: return USDCent(requests.get(url, params=params).json()["wallet"]["amount"]) -def get_wallet_balance_if_non_negative(amt_worker_id: str) -> Optional[USDCent]: +def get_wallet_balance_if_non_negative(amt_worker_id: str) -> USDCent | None: url = f"{settings.fsb_host}{settings.product_id}/wallet/" params = {"bpuid": amt_worker_id} amt = requests.get(url, params=params).json()["wallet"]["amount"] diff --git a/jb/managers/worker.py b/jb/managers/worker.py index e2d7237..9e7d5e3 100644 --- a/jb/managers/worker.py +++ b/jb/managers/worker.py @@ -1,5 +1,3 @@ -from typing import List - from mypy_boto3_mturk.type_defs import WorkerBlockTypeDef from jb.decorators import AMT_CLIENT @@ -8,9 +6,9 @@ from jb.decorators import AMT_CLIENT class WorkerManager: @staticmethod - def fetch_worker_blocks() -> List[WorkerBlockTypeDef]: + def fetch_worker_blocks() -> list[WorkerBlockTypeDef]: p = AMT_CLIENT.get_paginator("list_worker_blocks") - res: List[WorkerBlockTypeDef] = [] + res: list[WorkerBlockTypeDef] = [] for item in p.paginate(): res.extend(item["WorkerBlocks"]) return res |