324 lines
10 KiB
Python
324 lines
10 KiB
Python
"""Minimal e621 API client for server-side matching and metadata refresh.
|
|
|
|
The SPA talks to e621 directly for browsing; this client exists for work the
|
|
browser cannot do reliably: long batch scans, staged-upload processing and
|
|
requests tied to a library item rather than an open page. It uses the
|
|
requesting user's stored credentials and a global throttle (e621 asks for at
|
|
most two requests per second, one per second sustained).
|
|
|
|
e621's load balancer also sheds load with 429s (sometimes with an HTML
|
|
"shedding" page instead of JSON) and the IQDB endpoint has its own, much
|
|
stricter throttle. Every call therefore retries with exponential backoff and
|
|
honours ``Retry-After``; only 401/403 are treated as fatal.
|
|
"""
|
|
|
|
import logging
|
|
import random
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import requests
|
|
from django.conf import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Seconds between requests, per process. e621 allows 2/s hard and 1/s
|
|
# sustained; each gunicorn worker throttles on its own, so leave enough
|
|
# headroom that combined traffic does not trip the limit.
|
|
REQUEST_INTERVAL = 1.0
|
|
# IQDB is throttled far more aggressively than the rest of the API.
|
|
IQDB_INTERVAL = 2.0
|
|
|
|
MAX_ATTEMPTS = 4
|
|
BACKOFF_BASE = 2.0
|
|
BACKOFF_CAP = 60.0
|
|
RETRYABLE_STATUSES = {429, 500, 502, 503, 504}
|
|
|
|
|
|
class E621Error(Exception):
|
|
"""A non-404 failure while talking to e621."""
|
|
|
|
|
|
class E621NotFound(E621Error):
|
|
"""The requested post does not exist (HTTP 404)."""
|
|
|
|
|
|
class E621AuthError(E621Error):
|
|
"""e621 rejected the stored credentials (401/403)."""
|
|
|
|
|
|
class E621RateLimited(E621Error):
|
|
"""e621 shed load or throttled the request after every retry."""
|
|
|
|
|
|
_throttle_lock = threading.Lock()
|
|
_last_request_at = 0.0
|
|
|
|
|
|
def credentials_configured(user):
|
|
return bool(user is not None and getattr(user, "e621_configured", False))
|
|
|
|
|
|
def _wait_for_slot(interval=REQUEST_INTERVAL):
|
|
global _last_request_at
|
|
with _throttle_lock:
|
|
delay = _last_request_at + interval - time.monotonic()
|
|
if delay > 0:
|
|
time.sleep(delay)
|
|
_last_request_at = time.monotonic()
|
|
|
|
|
|
def _retry_delay(attempt, response=None):
|
|
"""Backoff for a retryable failure, honouring ``Retry-After``."""
|
|
if response is not None:
|
|
retry_after = response.headers.get("Retry-After")
|
|
if retry_after:
|
|
try:
|
|
return max(float(retry_after), 1.0)
|
|
except (TypeError, ValueError):
|
|
pass
|
|
delay = min(BACKOFF_BASE * (2**attempt), BACKOFF_CAP)
|
|
return delay + random.uniform(0, delay * 0.25)
|
|
|
|
|
|
def _request(
|
|
user,
|
|
method,
|
|
path,
|
|
*,
|
|
params=None,
|
|
data=None,
|
|
files=None,
|
|
timeout=30,
|
|
require_auth=True,
|
|
interval=REQUEST_INTERVAL,
|
|
attempts=MAX_ATTEMPTS,
|
|
):
|
|
"""One e621 call with retries; returns the parsed JSON payload.
|
|
|
|
``files`` may be a callable returning the multipart mapping, which is
|
|
called once per attempt: streamed uploads consume their file handle, so a
|
|
retry needs a freshly opened file.
|
|
|
|
Raises E621NotFound for 404s, E621AuthError for 401/403 and
|
|
E621RateLimited when e621 keeps shedding/throttling after every attempt.
|
|
"""
|
|
configured = credentials_configured(user)
|
|
if require_auth and not configured:
|
|
raise E621Error("Configure your e621 credentials in Account first.")
|
|
base = (getattr(user, "e621_base_url", "") or "https://e621.net").rstrip("/")
|
|
auth = (
|
|
(user.e621_username, user.e621_api_key_plain) if configured else None
|
|
)
|
|
url = f"{base}{path}"
|
|
|
|
last_error = None
|
|
for attempt in range(attempts):
|
|
request_files = files() if callable(files) else files
|
|
_wait_for_slot(interval)
|
|
response = None
|
|
try:
|
|
response = requests.request(
|
|
method,
|
|
url,
|
|
params=params,
|
|
data=data,
|
|
files=request_files,
|
|
auth=auth,
|
|
headers={"User-Agent": settings.USER_AGENT},
|
|
timeout=timeout,
|
|
)
|
|
except requests.RequestException as exc:
|
|
last_error = E621Error(f"Could not reach e621: {exc}")
|
|
else:
|
|
if response.status_code == 404:
|
|
raise E621NotFound(f"e621 returned 404 for {path}")
|
|
if response.status_code in {401, 403}:
|
|
raise E621AuthError(
|
|
f"e621 rejected the request ({response.status_code}). "
|
|
"Check the stored e621 credentials."
|
|
)
|
|
if response.status_code == 429:
|
|
# Throttles and load-shedding can arrive as JSON ({"message":
|
|
# "Throttled: ..."}) or as an HTML page.
|
|
message = ""
|
|
try:
|
|
payload = response.json()
|
|
except ValueError:
|
|
payload = None
|
|
if isinstance(payload, dict):
|
|
message = str(
|
|
payload.get("message") or payload.get("error") or ""
|
|
)
|
|
last_error = E621RateLimited(
|
|
message or f"e621 throttled the request for {path}"
|
|
)
|
|
elif response.status_code < 400:
|
|
try:
|
|
return response.json()
|
|
except ValueError:
|
|
# An HTML page with a 2xx status.
|
|
last_error = E621RateLimited(
|
|
f"e621 returned an unexpected {response.status_code} response."
|
|
)
|
|
elif response.status_code in RETRYABLE_STATUSES:
|
|
last_error = E621RateLimited(
|
|
f"e621 replied {response.status_code} for {path}"
|
|
)
|
|
else:
|
|
raise E621Error(f"e621 replied {response.status_code} for {path}")
|
|
finally:
|
|
_close_upload_files(request_files)
|
|
if attempt + 1 < attempts:
|
|
delay = _retry_delay(attempt, response)
|
|
logger.info(
|
|
"e621 %s %s failed (%s); retrying in %.1fs",
|
|
method,
|
|
path,
|
|
last_error,
|
|
delay,
|
|
)
|
|
time.sleep(delay)
|
|
|
|
if last_error is None:
|
|
last_error = E621Error("e621 request failed.")
|
|
raise last_error
|
|
|
|
|
|
def _close_upload_files(files):
|
|
"""Close the handles behind a multipart mapping (see _request)."""
|
|
if not isinstance(files, dict):
|
|
return
|
|
for value in files.values():
|
|
handle = value[1] if isinstance(value, tuple) and len(value) > 1 else value
|
|
close = getattr(handle, "close", None)
|
|
if close is not None:
|
|
try:
|
|
close()
|
|
except Exception: # noqa: BLE001 - closing must never mask errors
|
|
pass
|
|
|
|
|
|
def get(user, path, params=None, timeout=30, require_auth=True):
|
|
"""GET an e621 API path using the user's credentials.
|
|
|
|
Reads that e621 serves anonymously (searches, pools, tags) can pass
|
|
require_auth=False; matching endpoints keep requiring credentials.
|
|
"""
|
|
return _request(
|
|
user,
|
|
"GET",
|
|
path,
|
|
params=params,
|
|
timeout=timeout,
|
|
require_auth=require_auth,
|
|
)
|
|
|
|
|
|
def find_post_by_md5(user, md5):
|
|
"""The e621 post with this exact MD5, or None."""
|
|
payload = get(
|
|
user,
|
|
"/posts.json",
|
|
params={"tags": f"md5:{md5}", "limit": 1},
|
|
require_auth=False,
|
|
)
|
|
posts = payload.get("posts") if isinstance(payload, dict) else None
|
|
if not posts:
|
|
return None
|
|
return posts[0]
|
|
|
|
|
|
def fetch_post(user, post_id):
|
|
"""One post by id. Raises E621NotFound when the post is gone."""
|
|
payload = get(user, f"/posts/{int(post_id)}.json")
|
|
post = payload.get("post") if isinstance(payload, dict) else None
|
|
if not isinstance(post, dict):
|
|
raise E621Error("e621 returned an unexpected post payload.")
|
|
return post
|
|
|
|
|
|
def check_md5_batch(user, md5s):
|
|
"""Look many MD5s up in one posts.json query.
|
|
|
|
Returns ``{md5: post}`` for the ones e621 knows; missing MD5s are simply
|
|
absent. Works anonymously, like the original app's batch cache command.
|
|
"""
|
|
wanted = {str(value).strip().lower() for value in md5s if value}
|
|
if not wanted:
|
|
return {}
|
|
values = sorted(wanted)
|
|
payload = get(
|
|
user,
|
|
"/posts.json",
|
|
params={
|
|
"tags": f"md5:{','.join(values)}",
|
|
"limit": min(len(values), 320),
|
|
},
|
|
require_auth=False,
|
|
)
|
|
posts = payload.get("posts") if isinstance(payload, dict) else None
|
|
found = {}
|
|
for post in posts or []:
|
|
if not isinstance(post, dict):
|
|
continue
|
|
file_data = post.get("file") or {}
|
|
md5 = str(file_data.get("md5") or "").strip().lower()
|
|
if md5 in wanted:
|
|
found[md5] = post
|
|
return found
|
|
|
|
|
|
def fetch_posts_by_ids(user, ids):
|
|
"""Fetch many posts in one query (up to 320 ids). Missing ids are absent."""
|
|
values = sorted({int(value) for value in ids})
|
|
if not values:
|
|
return []
|
|
payload = get(
|
|
user,
|
|
"/posts.json",
|
|
params={
|
|
"tags": f"id:{','.join(str(value) for value in values)}",
|
|
"limit": min(len(values), 320),
|
|
},
|
|
require_auth=False,
|
|
)
|
|
posts = payload.get("posts") if isinstance(payload, dict) else None
|
|
return [post for post in posts or [] if isinstance(post, dict)]
|
|
|
|
|
|
def iqdb_search(user, path, timeout=60):
|
|
"""Reverse-image search one file through e621's IQDB endpoint.
|
|
|
|
Returns the legacy match list. Uses the extra-strict IQDB interval and
|
|
retries through e621's throttle; raises E621RateLimited when it persists.
|
|
"""
|
|
path = Path(path)
|
|
|
|
def open_file():
|
|
# A fresh handle per attempt: the stream is consumed by the request.
|
|
return {"search[file]": (path.name, open(path, "rb"))}
|
|
|
|
payload = _request(
|
|
user,
|
|
"POST",
|
|
"/iqdb_queries.json",
|
|
files=open_file,
|
|
timeout=timeout,
|
|
require_auth=False,
|
|
interval=IQDB_INTERVAL,
|
|
)
|
|
if isinstance(payload, list):
|
|
return payload
|
|
if isinstance(payload, dict):
|
|
matches = payload.get("matches")
|
|
if isinstance(matches, list):
|
|
return matches
|
|
# e621 answers its throttle with {"success": false, "message": ...}.
|
|
message = payload.get("message") or payload.get("error")
|
|
if message:
|
|
raise E621RateLimited(str(message))
|
|
raise E621Error("e621 returned an unexpected IQDB payload.")
|
|
return []
|