diff --git a/.env_template b/.env_template index 4b110d2..f7524a8 100644 --- a/.env_template +++ b/.env_template @@ -19,7 +19,11 @@ POSTGRES_URI="postgresql://pictrsafety:ThatSecretPGSQLPass@127.0.0.1/pictrsafety # Note than 127.0.0.1 is always allowed. #KNOWN_PICTRS_IPS='["192.168.0.1"]' # If this amount of images are waiting for scanning, start retuning "OK" for each new scan after that -# This will prevent pict-rs from rejecting each and every upload because your scanning workers crashed or something +# This will prevent pict-rs from rejecting each and every upload because your scanning workers are too slow or crashed or something # and the admin is AFK # Set this value to 0 to never bypass the scan. -SCAN_BYPASS_THRESHOLD=10 \ No newline at end of file +SCAN_BYPASS_THRESHOLD=10 +# If this amount of seconds have passed since a scanning worker last checked in, start returning "OK" for each new scan after that +# This will prevent pict-rs from rejecting each and every upload because your scanning workers have gone missing +# Set this value to 0 to never bypass the scan. +MISSING_WORKER_THRESHOLD=5 \ No newline at end of file diff --git a/pictrs_safety_api/apis/v1/base.py b/pictrs_safety_api/apis/v1/base.py index 447b787..45c9578 100644 --- a/pictrs_safety_api/apis/v1/base.py +++ b/pictrs_safety_api/apis/v1/base.py @@ -13,6 +13,7 @@ from pictrs_safety_api.database import functions as database from pictrs_safety_api import exceptions as e from pictrs_safety_api import enums +from pictrs_safety_api.classes.worker import worker api = Namespace('v1', 'API Version 1' ) @@ -51,9 +52,13 @@ def post(self,pictrs_id): raise e.Unauthorized("You are not authorized to use this service", f"Unauthorized IP: {request.remote_addr}") elif pictrs_id not in json.loads(os.getenv("KNOWN_PICTRS_IDS", "[]")): raise e.Unauthorized("You are not authorized to use this service", f"Unauthorized ID: {pictrs_id}") + if worker.is_stale(): + logger.warning(f"Returning OK due to stale worker. Last seen: {worker.last_check_in}") + return {"message": "Worker Stale"}, 200 scan_threshold = int(os.getenv("SCAN_BYPASS_THRESHOLD", 10)) if scan_threshold > 0 and database.count_waiting_scan_requests() > scan_threshold: - return {"message": "Image OK"}, 200 + logger.warning(f"Returning OK due to full queue. Queue: {database.count_waiting_scan_requests()}") + return {"message": "Queue Full"}, 200 self.args = self.post_parser.parse_args() file = self.args["file"] if not file: @@ -144,10 +149,11 @@ def get(self): '''Pick up an image to safety validate ''' # I don't get why this is not using the import from earlier... - logger.debug(request.remote_addr) + # logger.debug(request.remote_addr) self.args = self.get_parser.parse_args() if os.getenv("FEDIVERSE_SAFETY_WORKER_AUTH") != self.args.apikey: raise e.Forbidden("Access Denied") + worker.check_in() pop: ScanRequest = database.find_waiting_scan_request() if not pop: return {"message": "Nothing to do"},204 diff --git a/pictrs_safety_api/classes/request.py b/pictrs_safety_api/classes/request.py index 9435ea3..5f0d4bf 100644 --- a/pictrs_safety_api/classes/request.py +++ b/pictrs_safety_api/classes/request.py @@ -1,6 +1,5 @@ -import dateutil.relativedelta from datetime import datetime -from sqlalchemy import Enum, UniqueConstraint +from sqlalchemy import Enum from loguru import logger from pictrs_safety_api.flask import db diff --git a/pictrs_safety_api/classes/worker.py b/pictrs_safety_api/classes/worker.py new file mode 100644 index 0000000..714a5bd --- /dev/null +++ b/pictrs_safety_api/classes/worker.py @@ -0,0 +1,20 @@ +import os +from datetime import datetime, timedelta + +class Worker(): + last_check_in = None + + def __init__(self): + self.check_in() + + def check_in(self): + self.last_check_in = datetime.utcnow() + + def is_stale(self): + threshold = int(os.getenv("MISSING_WORKER_THRESHOLD", 5)) + # When threshold is 0, the worker is never stale + if threshold == 0: + return False + return self.last_check_in < datetime.utcnow() - timedelta(seconds=threshold) + +worker = Worker() \ No newline at end of file diff --git a/pictrs_safety_api/database/functions.py b/pictrs_safety_api/database/functions.py index caf5eee..eabcc18 100644 --- a/pictrs_safety_api/database/functions.py +++ b/pictrs_safety_api/database/functions.py @@ -1,8 +1,4 @@ -import time -import uuid -import json from loguru import logger -from datetime import datetime, timedelta from sqlalchemy import func, or_, and_, not_, Boolean from sqlalchemy.orm import noload from pictrs_safety_api.flask import db