Teknium
440 lines
17 KiB
Python
440 lines
17 KiB
Python
"""QQ Bot chunked upload flow.
|
|
|
|
The QQ v2 API caps inline base64 uploads (``file_data`` / ``url``) at ~10 MB.
|
|
Files between 10 MB and ~100 MB use the three-step chunked flow::
|
|
|
|
1. POST /v2/{users|groups}/{id}/upload_prepare
|
|
→ upload_id, block_size, and pre-signed COS part URLs.
|
|
2. Per part: PUT bytes to the COS URL, then POST .../upload_part_finish.
|
|
3. POST /v2/{users|groups}/{id}/files with {"upload_id": ...}
|
|
→ ``file_info`` token used in a RichMedia message.
|
|
|
|
Error codes (QQ Bot v2 spec): ``40093001`` — ``upload_part_finish`` retryable
|
|
until the server's ``retry_timeout`` (or a local cap) elapses; ``40093002`` —
|
|
daily upload quota exceeded, surfaced as :class:`UploadDailyLimitExceededError`.
|
|
Other API/I/O failures raise ``RuntimeError``.
|
|
|
|
Ported from WideLee's qqbot-agent-sdk v1.2.2 (``media_loader.py::ChunkedUploader``).
|
|
Authorship preserved via Co-authored-by.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Awaitable, Callable, Dict, List
|
|
|
|
from gateway.platforms.qqbot.constants import FILE_UPLOAD_TIMEOUT
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
_BIZ_CODE_DAILY_LIMIT = 40093002 # upload_prepare: daily cumulative limit
|
|
_BIZ_CODE_PART_RETRYABLE = 40093001 # upload_part_finish: transient
|
|
|
|
_DEFAULT_CONCURRENT_PARTS = 1
|
|
_MAX_CONCURRENT_PARTS = 10
|
|
|
|
_PART_UPLOAD_TIMEOUT = 300.0 # 5 minutes per COS PUT
|
|
_PART_UPLOAD_MAX_RETRIES = 2
|
|
_PART_FINISH_RETRY_INTERVAL = 1.0
|
|
_PART_FINISH_DEFAULT_TIMEOUT = 120.0
|
|
_PART_FINISH_MAX_TIMEOUT = 600.0
|
|
|
|
_COMPLETE_UPLOAD_MAX_RETRIES = 2
|
|
_COMPLETE_UPLOAD_BASE_DELAY = 2.0
|
|
|
|
# First 10,002,432 bytes used for the ``md5_10m`` hash (per QQ API spec).
|
|
_MD5_10M_SIZE = 10_002_432
|
|
|
|
|
|
# ── Exceptions ───────────────────────────────────────────────────────
|
|
|
|
class UploadDailyLimitExceededError(Exception):
|
|
"""Raised when ``upload_prepare`` returns biz_code 40093002 (daily quota hit)."""
|
|
|
|
def __init__(self, file_name: str, file_size: int, message: str = "") -> None:
|
|
self.file_name = file_name
|
|
self.file_size = file_size
|
|
super().__init__(message or f"Daily upload limit exceeded for {file_name!r}")
|
|
|
|
@property
|
|
def file_size_human(self) -> str:
|
|
return format_size(self.file_size)
|
|
|
|
|
|
class UploadFileTooLargeError(Exception):
|
|
"""Raised when a file exceeds the platform per-file size limit."""
|
|
|
|
def __init__(self, file_name: str, file_size: int, limit_bytes: int = 0, message: str = "") -> None:
|
|
self.file_name = file_name
|
|
self.file_size = file_size
|
|
self.limit_bytes = limit_bytes
|
|
limit_str = f" ({format_size(limit_bytes)})" if limit_bytes else ""
|
|
super().__init__(
|
|
message or f"File {file_name!r} ({format_size(file_size)}) exceeds platform limit{limit_str}"
|
|
)
|
|
|
|
@property
|
|
def file_size_human(self) -> str:
|
|
return format_size(self.file_size)
|
|
|
|
@property
|
|
def limit_human(self) -> str:
|
|
return format_size(self.limit_bytes) if self.limit_bytes else "unknown"
|
|
|
|
|
|
# ── Prepare-response shape ───────────────────────────────────────────
|
|
|
|
@dataclass
|
|
class _PreparePart:
|
|
index: int
|
|
presigned_url: str
|
|
block_size: int = 0
|
|
|
|
|
|
@dataclass
|
|
class _PrepareResult:
|
|
upload_id: str
|
|
block_size: int
|
|
parts: List[_PreparePart]
|
|
concurrency: int = _DEFAULT_CONCURRENT_PARTS
|
|
retry_timeout: float = 0.0
|
|
|
|
|
|
def _parse_prepare_response(raw: Dict[str, Any]) -> _PrepareResult:
|
|
"""Parse upload_prepare response (either bare or wrapped in ``data``)."""
|
|
src = raw.get("data") if isinstance(raw.get("data"), dict) else raw
|
|
upload_id = str(src.get("upload_id", ""))
|
|
if not upload_id:
|
|
raise ValueError(f"upload_prepare response missing upload_id: {str(raw)[:200]}")
|
|
block_size = int(src.get("block_size", 0))
|
|
raw_parts = src.get("parts") or src.get("part_list") or []
|
|
if not isinstance(raw_parts, list) or not raw_parts:
|
|
raise ValueError(f"upload_prepare response missing parts: {str(raw)[:200]}")
|
|
parts = [
|
|
_PreparePart(
|
|
index=int(p.get("part_index") or p.get("index") or 0),
|
|
presigned_url=str(p.get("presigned_url") or p.get("url") or ""),
|
|
block_size=int(p.get("block_size", 0)),
|
|
)
|
|
for p in raw_parts
|
|
if isinstance(p, dict)
|
|
]
|
|
return _PrepareResult(
|
|
upload_id=upload_id,
|
|
block_size=block_size,
|
|
parts=parts,
|
|
concurrency=int(src.get("concurrency", _DEFAULT_CONCURRENT_PARTS)) or _DEFAULT_CONCURRENT_PARTS,
|
|
retry_timeout=float(src.get("retry_timeout", 0.0) or 0.0),
|
|
)
|
|
|
|
|
|
def _api_path(chat_type: str, target_id: str, endpoint: str) -> str:
|
|
base = "/v2/users" if chat_type == "c2c" else "/v2/groups"
|
|
return f"{base}/{target_id}/{endpoint}"
|
|
|
|
|
|
# ── Chunked upload driver ────────────────────────────────────────────
|
|
|
|
class ChunkedUploader:
|
|
"""Run the prepare → PUT parts → complete sequence.
|
|
|
|
:param api_request: Bound ``_api_request(method, path, body=..., timeout=...)``
|
|
coroutine from the adapter (passed in rather than imported to avoid
|
|
circular imports). Must raise ``RuntimeError`` with the biz_code in
|
|
the message on API errors.
|
|
:param http_put: Coroutine ``(url, data, headers) -> httpx-like response``
|
|
for COS part uploads.
|
|
:param log_tag: Log prefix.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
api_request: Callable[..., Awaitable[Dict[str, Any]]],
|
|
http_put: Callable[..., Awaitable[Any]],
|
|
log_tag: str = "QQBot",
|
|
) -> None:
|
|
self._api_request = api_request
|
|
self._http_put = http_put
|
|
self._log_tag = log_tag
|
|
|
|
async def upload(
|
|
self,
|
|
chat_type: str,
|
|
target_id: str,
|
|
file_path: str,
|
|
file_type: int,
|
|
file_name: str,
|
|
) -> Dict[str, Any]:
|
|
"""Run the full chunked upload and return the ``complete_upload`` response.
|
|
|
|
:param chat_type: ``'c2c'`` or ``'group'``.
|
|
:param file_type: ``MEDIA_TYPE_*`` constant.
|
|
:returns: Raw ``complete_upload`` response dict (contains ``file_info``).
|
|
:raises UploadDailyLimitExceededError: On biz_code 40093002.
|
|
:raises UploadFileTooLargeError: When the file exceeds the platform limit.
|
|
:raises RuntimeError: On other API or I/O failures.
|
|
"""
|
|
if chat_type not in {"c2c", "group"}:
|
|
raise ValueError(f"ChunkedUploader: unsupported chat_type {chat_type!r}")
|
|
|
|
file_size = Path(file_path).stat().st_size
|
|
logger.info(
|
|
"[%s] Chunked upload start: file=%s size=%s type=%d",
|
|
self._log_tag, file_name, format_size(file_size), file_type,
|
|
)
|
|
|
|
# Hashing is blocking I/O → executor.
|
|
hashes = await asyncio.get_running_loop().run_in_executor(
|
|
None, _compute_file_hashes, file_path, file_size
|
|
)
|
|
prepare = await self._prepare(chat_type, target_id, file_type, file_name, file_size, hashes)
|
|
max_concurrent = min(prepare.concurrency, _MAX_CONCURRENT_PARTS)
|
|
retry_timeout = min(
|
|
prepare.retry_timeout if prepare.retry_timeout > 0 else _PART_FINISH_DEFAULT_TIMEOUT,
|
|
_PART_FINISH_MAX_TIMEOUT,
|
|
)
|
|
logger.info(
|
|
"[%s] Prepared: upload_id=%s block_size=%s parts=%d concurrency=%d",
|
|
self._log_tag, prepare.upload_id, format_size(prepare.block_size),
|
|
len(prepare.parts), max_concurrent,
|
|
)
|
|
|
|
total_parts = len(prepare.parts)
|
|
completed = [0] # shared counter for progress logging
|
|
sem = asyncio.Semaphore(max(max_concurrent, 1))
|
|
|
|
async def _run(part: _PreparePart) -> None:
|
|
async with sem:
|
|
await self._upload_one_part(
|
|
chat_type, target_id, file_path, file_size, prepare.upload_id,
|
|
prepare.block_size, part, retry_timeout, total_parts, completed,
|
|
)
|
|
|
|
await asyncio.gather(*(_run(p) for p in prepare.parts))
|
|
logger.info("[%s] All %d parts uploaded, completing…", self._log_tag, total_parts)
|
|
return await self._complete(chat_type, target_id, prepare.upload_id)
|
|
|
|
async def _prepare(
|
|
self,
|
|
chat_type: str,
|
|
target_id: str,
|
|
file_type: int,
|
|
file_name: str,
|
|
file_size: int,
|
|
hashes: Dict[str, str],
|
|
) -> _PrepareResult:
|
|
body = {
|
|
"file_type": file_type,
|
|
"file_name": file_name,
|
|
"file_size": file_size,
|
|
"md5": hashes["md5"],
|
|
"sha1": hashes["sha1"],
|
|
"md5_10m": hashes["md5_10m"],
|
|
}
|
|
try:
|
|
raw = await self._api_request(
|
|
"POST", _api_path(chat_type, target_id, "upload_prepare"),
|
|
body=body, timeout=FILE_UPLOAD_TIMEOUT,
|
|
)
|
|
except RuntimeError as exc:
|
|
err_msg = str(exc)
|
|
if f"{_BIZ_CODE_DAILY_LIMIT}" in err_msg:
|
|
raise UploadDailyLimitExceededError(file_name, file_size, err_msg) from exc
|
|
raise
|
|
return _parse_prepare_response(raw)
|
|
|
|
async def _upload_one_part(
|
|
self,
|
|
chat_type: str,
|
|
target_id: str,
|
|
file_path: str,
|
|
file_size: int,
|
|
upload_id: str,
|
|
rsp_block_size: int,
|
|
part: _PreparePart,
|
|
retry_timeout: float,
|
|
total_parts: int,
|
|
completed: List[int],
|
|
) -> None:
|
|
"""PUT one part to COS, then call ``upload_part_finish``."""
|
|
part_index = part.index
|
|
# Per-part block_size wins; fall back to the response-level value.
|
|
actual_block_size = part.block_size if part.block_size > 0 else rsp_block_size
|
|
offset = (part_index - 1) * rsp_block_size
|
|
length = min(actual_block_size, file_size - offset)
|
|
|
|
data = await asyncio.get_running_loop().run_in_executor(
|
|
None, _read_file_chunk, file_path, offset, length
|
|
)
|
|
md5_hex = hashlib.md5(data).hexdigest()
|
|
logger.debug(
|
|
"[%s] Part %d/%d: uploading %s (offset=%d md5=%s)",
|
|
self._log_tag, part_index, total_parts, format_size(length), offset, md5_hex,
|
|
)
|
|
|
|
await self._put_to_presigned_url(part.presigned_url, data, part_index, total_parts)
|
|
await self._part_finish_with_retry(
|
|
chat_type, target_id, upload_id, part_index, length, md5_hex, retry_timeout,
|
|
)
|
|
|
|
completed[0] += 1
|
|
logger.debug(
|
|
"[%s] Part %d/%d done (%d/%d total)",
|
|
self._log_tag, part_index, total_parts, completed[0], total_parts,
|
|
)
|
|
|
|
async def _with_retries(
|
|
self,
|
|
attempt_fn: Callable[[], Awaitable[Any]],
|
|
*,
|
|
max_retries: int,
|
|
base_delay: float,
|
|
label: str,
|
|
failure_label: str,
|
|
) -> Any:
|
|
"""Run *attempt_fn* up to ``max_retries + 1`` times with exponential backoff."""
|
|
last_exc: Exception | None = None
|
|
for attempt in range(max_retries + 1):
|
|
try:
|
|
return await attempt_fn()
|
|
except Exception as exc:
|
|
last_exc = exc
|
|
if attempt < max_retries:
|
|
delay = base_delay * (2 ** attempt)
|
|
logger.warning(
|
|
"[%s] %s attempt %d failed, retry in %.1fs: %s",
|
|
self._log_tag, label, attempt + 1, delay, exc,
|
|
)
|
|
await asyncio.sleep(delay)
|
|
raise RuntimeError(f"{failure_label} failed after {max_retries + 1} attempts: {last_exc}")
|
|
|
|
async def _put_to_presigned_url(self, url: str, data: bytes, part_index: int, total_parts: int) -> None:
|
|
"""PUT part data to a pre-signed COS URL with retry."""
|
|
|
|
async def _attempt() -> None:
|
|
resp = await asyncio.wait_for(
|
|
self._http_put(url, data=data, headers={"Content-Length": str(len(data))}),
|
|
timeout=_PART_UPLOAD_TIMEOUT,
|
|
)
|
|
status = getattr(resp, "status_code", 0)
|
|
if 200 <= status < 300:
|
|
logger.debug("[%s] PUT part %d/%d: %d OK", self._log_tag, part_index, total_parts, status)
|
|
return
|
|
body_preview = ""
|
|
try:
|
|
body_preview = getattr(resp, "text", "")[:200]
|
|
except Exception: # pragma: no cover — defensive
|
|
pass
|
|
raise RuntimeError(f"COS PUT returned {status}: {body_preview}")
|
|
|
|
await self._with_retries(
|
|
_attempt, max_retries=_PART_UPLOAD_MAX_RETRIES, base_delay=1.0,
|
|
label=f"PUT part {part_index}/{total_parts}",
|
|
failure_label=f"Part {part_index}/{total_parts} upload",
|
|
)
|
|
|
|
async def _part_finish_with_retry(
|
|
self,
|
|
chat_type: str,
|
|
target_id: str,
|
|
upload_id: str,
|
|
part_index: int,
|
|
block_size: int,
|
|
md5: str,
|
|
retry_timeout: float,
|
|
) -> None:
|
|
"""Call ``upload_part_finish``, retrying on biz_code 40093001 until *retry_timeout*."""
|
|
path = _api_path(chat_type, target_id, "upload_part_finish")
|
|
body = {"upload_id": upload_id, "part_index": part_index, "block_size": block_size, "md5": md5}
|
|
loop = asyncio.get_running_loop()
|
|
start = loop.time()
|
|
attempt = 0
|
|
while True:
|
|
try:
|
|
await self._api_request("POST", path, body=body, timeout=FILE_UPLOAD_TIMEOUT)
|
|
return
|
|
except RuntimeError as exc:
|
|
if f"{_BIZ_CODE_PART_RETRYABLE}" not in str(exc):
|
|
raise
|
|
elapsed = loop.time() - start
|
|
if elapsed >= retry_timeout:
|
|
raise RuntimeError(
|
|
f"upload_part_finish persistent retry timed out "
|
|
f"after {retry_timeout:.0f}s ({attempt} retries): {exc}"
|
|
) from exc
|
|
attempt += 1
|
|
logger.debug(
|
|
"[%s] part_finish retryable error, attempt %d, elapsed=%.1fs: %s",
|
|
self._log_tag, attempt, elapsed, exc,
|
|
)
|
|
await asyncio.sleep(_PART_FINISH_RETRY_INTERVAL)
|
|
|
|
async def _complete(self, chat_type: str, target_id: str, upload_id: str) -> Dict[str, Any]:
|
|
"""Call ``complete_upload`` with retry.
|
|
|
|
Reuses the ``/files`` endpoint (same as the simple URL-based upload) but
|
|
signals the chunked-completion path by sending only ``upload_id``.
|
|
"""
|
|
path = _api_path(chat_type, target_id, "files")
|
|
body = {"upload_id": upload_id}
|
|
return await self._with_retries(
|
|
lambda: self._api_request("POST", path, body=body, timeout=FILE_UPLOAD_TIMEOUT),
|
|
max_retries=_COMPLETE_UPLOAD_MAX_RETRIES, base_delay=_COMPLETE_UPLOAD_BASE_DELAY,
|
|
label="complete_upload", failure_label="complete_upload",
|
|
)
|
|
|
|
|
|
# ── Helpers (module-level for testability) ───────────────────────────
|
|
|
|
def format_size(size_bytes: int) -> str:
|
|
"""Return a human-readable file size string (e.g. ``'12.3 MB'``)."""
|
|
size = float(size_bytes)
|
|
for unit in ("B", "KB", "MB", "GB"):
|
|
if size < 1024.0:
|
|
return f"{size:.1f} {unit}"
|
|
size /= 1024.0
|
|
return f"{size:.1f} TB"
|
|
|
|
|
|
def _read_file_chunk(file_path: str, offset: int, length: int) -> bytes:
|
|
"""Read *length* bytes at *offset*; raises IOError on a short read (truncated file)."""
|
|
with open(file_path, "rb") as fh:
|
|
fh.seek(offset)
|
|
data = fh.read(length)
|
|
if len(data) != length:
|
|
raise IOError(
|
|
f"Short read from {file_path}: expected {length} bytes at "
|
|
f"offset {offset}, got {len(data)} (file may be truncated)"
|
|
)
|
|
return data
|
|
|
|
|
|
def _compute_file_hashes(file_path: str, file_size: int) -> Dict[str, str]:
|
|
"""Compute md5, sha1, and md5_10m in a single pass."""
|
|
md5 = hashlib.md5()
|
|
sha1 = hashlib.sha1()
|
|
md5_10m = hashlib.md5()
|
|
need_10m = file_size > _MD5_10M_SIZE
|
|
bytes_read = 0
|
|
with open(file_path, "rb") as fh:
|
|
while chunk := fh.read(65536):
|
|
md5.update(chunk)
|
|
sha1.update(chunk)
|
|
if need_10m:
|
|
remaining = _MD5_10M_SIZE - bytes_read
|
|
if remaining > 0:
|
|
md5_10m.update(chunk[:remaining])
|
|
bytes_read += len(chunk)
|
|
full_md5 = md5.hexdigest()
|
|
return {
|
|
"md5": full_md5,
|
|
"sha1": sha1.hexdigest(),
|
|
# For small files the "10m" hash is just the full md5.
|
|
"md5_10m": md5_10m.hexdigest() if need_10m else full_md5,
|
|
}
|