403 lines
13 KiB
Python
403 lines
13 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
请求客户端模块
|
||
"""
|
||
|
||
import json
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Any, Dict, Generator, Literal, Optional, Tuple, Union
|
||
from xml.etree import ElementTree
|
||
|
||
from pydantic import BaseModel, Field, HttpUrl, model_validator
|
||
from requests import Response, Session
|
||
from requests.adapters import HTTPAdapter
|
||
from urllib3.util.retry import Retry
|
||
|
||
sys.path.append(Path(__file__).parent.as_posix())
|
||
from sqlite import SQLite
|
||
|
||
|
||
class Parameters(BaseModel):
|
||
"""
|
||
请求参数模型
|
||
"""
|
||
|
||
url: HttpUrl = Field(default=..., description="统一资源定位符，基于HttpUrl自动校验")
|
||
params: Optional[Dict[str, Any]] = Field(
|
||
default=None, description="统一资源定位符的查询参数"
|
||
)
|
||
headers: Optional[Dict[str, str]] = Field(default=None, description="请求头")
|
||
data: Optional[Dict[str, Any]] = Field(default=None, description="表单数据")
|
||
json_: Optional[Dict[str, Any]] = Field(
|
||
default=None, alias="json", description="JSON数据"
|
||
)
|
||
files: Optional[
|
||
Dict[
|
||
str,
|
||
Union[
|
||
Tuple[str, bytes],
|
||
Tuple[str, bytes, str],
|
||
Tuple[str, bytes, str, Dict[str, str]],
|
||
],
|
||
]
|
||
] = Field(
|
||
default=None,
|
||
description="上传文件，{字段名: (文件名, 字节数据, 内容类型, 请求头)}",
|
||
)
|
||
stream_enabled: Optional[bool] = Field(default=None, description="使用流式传输")
|
||
guid: Optional[str] = Field(default=None, description="缓存全局唯一标识")
|
||
|
||
@model_validator(mode="after")
|
||
def validate_data(self):
|
||
"""校验：表单数据和JSON数据互斥"""
|
||
if self.data is not None and self.json_ is not None:
|
||
raise ValueError("表单数据和JSON数据不能同时使用")
|
||
return self
|
||
|
||
@model_validator(mode="after")
|
||
def validate_files(self):
|
||
if self.files is not None and self.stream_enabled:
|
||
raise ValueError("上传文件和使用流式传输不能同时使用")
|
||
return self
|
||
|
||
|
||
class RequestException(Exception):
|
||
"""请求异常"""
|
||
|
||
def __init__(
|
||
self,
|
||
status: Optional[int] = 400,
|
||
code: int = 0,
|
||
message: str = "请求发生异常",
|
||
):
|
||
"""
|
||
:param status: 状态码
|
||
:param code: 错误码
|
||
:param message: 错误信息
|
||
"""
|
||
self.status = status
|
||
self.code = code
|
||
self.message = message
|
||
super().__init__(self.message)
|
||
|
||
|
||
class Caches(SQLite):
|
||
"""
|
||
缓存，支持：
|
||
query：查询并返回单条缓存
|
||
update：新增或更新单条缓存
|
||
"""
|
||
|
||
def __init__(self, cache_ttl: int):
|
||
"""
|
||
初始化
|
||
:param cache_ttl: 缓存生存时间，单位为秒
|
||
"""
|
||
# 初始化
|
||
super().__init__(database=Path(__file__).parent.resolve() / "caches.db")
|
||
self.cache_ttl = cache_ttl
|
||
|
||
# 初始化缓存表（不清理过期缓存）
|
||
try:
|
||
with self:
|
||
self.execute(
|
||
sql="""
|
||
CREATE TABLE IF NOT EXISTS caches
|
||
(
|
||
--缓存唯一标识
|
||
guid TEXT PRIMARY KEY,
|
||
--缓存（JSON序列化）
|
||
cache TEXT NOT NULL,
|
||
--缓存时间
|
||
timestamp REAL NOT NULL
|
||
)
|
||
"""
|
||
)
|
||
self.execute(
|
||
sql="""
|
||
CREATE INDEX IF NOT EXISTS idx_timestamp ON caches(timestamp)
|
||
"""
|
||
)
|
||
except Exception as exception:
|
||
raise RuntimeError(f"初始化缓存表发生异常：{str(exception)}") from exception
|
||
|
||
def query(self, guid: str) -> Optional[Dict[str, Any]]:
|
||
"""
|
||
查询并返回单条缓存
|
||
:param guid: 缓存唯一标识
|
||
:return: 缓存
|
||
"""
|
||
try:
|
||
with self:
|
||
result = self.query_one(
|
||
sql="""
|
||
SELECT cache
|
||
FROM caches
|
||
WHERE guid = ? AND timestamp >= ?
|
||
""",
|
||
parameters=(guid, time.time() - self.cache_ttl),
|
||
)
|
||
return None if result is None else json.loads(result["cache"])
|
||
except Exception as exception:
|
||
raise RuntimeError(
|
||
f"查询并获取单条缓存发生异常：{str(exception)}"
|
||
) from exception
|
||
|
||
def update(self, guid: str, cache: Dict) -> Optional[bool]:
|
||
"""
|
||
新增或更新单条缓存（若无则新增缓存，若有则更新缓存）
|
||
:param guid: 缓存唯一标识
|
||
:param cache: 缓存
|
||
:return: 成功返回True，失败返回False
|
||
"""
|
||
try:
|
||
with self:
|
||
return self.execute(
|
||
sql="""
|
||
INSERT OR REPLACE INTO caches (guid, cache, timestamp) VALUES (?, ?, ?)
|
||
""",
|
||
parameters=(
|
||
guid,
|
||
json.dumps(cache, ensure_ascii=False),
|
||
time.time(),
|
||
),
|
||
)
|
||
except Exception as exception:
|
||
raise RuntimeError("新增或更新缓存发生异常") from exception
|
||
|
||
|
||
class Request:
|
||
"""
|
||
请求客户端，支持：
|
||
get：GET请求
|
||
post：POST请求
|
||
download：下载
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
default_headers: Optional[Dict[str, str]] = None,
|
||
total: int = 3,
|
||
backoff_factor: float = 0.5,
|
||
timeout: int = 60,
|
||
cache_enabled: bool = False,
|
||
cache_ttl: int = 360,
|
||
):
|
||
"""
|
||
:param default_headers: 默认请求头
|
||
:param total: 最大重试次数，默认 3
|
||
:param backoff_factor: 重试间隔退避因子，默认 0.5
|
||
:param timeout: 超时时间（单位为秒），默认为 60
|
||
:param cache_enabled: 使用缓存，默认 False
|
||
:param cache_ttl: 缓存生存时间（单位为天），默认为 360
|
||
"""
|
||
# 创建请求会话并挂载适配器
|
||
self.session = self._create_session(
|
||
default_headers=default_headers, total=total, backoff_factor=backoff_factor
|
||
)
|
||
# 初始化超时时间
|
||
self.timeout = timeout
|
||
# 实例化缓存
|
||
self.caches = Caches(cache_ttl=cache_ttl * 86400) if cache_enabled else None
|
||
|
||
@staticmethod
|
||
def _create_session(
|
||
total: int,
|
||
backoff_factor: float,
|
||
default_headers: Optional[Dict[str, str]] = None,
|
||
) -> Session:
|
||
"""
|
||
创建请求会话并挂载适配器
|
||
:param default_headers: 默认请求头
|
||
:param total: 最大重试次数
|
||
:param backoff_factor: 重试间隔退避因子
|
||
:return Session: 请求会话实例
|
||
"""
|
||
# 实例化请求会话
|
||
session = Session()
|
||
|
||
# 设置默认请求头
|
||
if default_headers:
|
||
session.headers.update(default_headers)
|
||
|
||
# 设置重试策略并挂载适配器
|
||
adapter = HTTPAdapter(
|
||
max_retries=Retry(
|
||
allowed_methods=["HEAD", "GET", "POST", "PUT", "DELETE", "PATCH"],
|
||
status_forcelist=[
|
||
408,
|
||
502,
|
||
503,
|
||
504,
|
||
], # 408为请求超时，502为网关错误，503为服务不可用，504为网关超时
|
||
total=total,
|
||
respect_retry_after_header=True,
|
||
backoff_factor=backoff_factor,
|
||
)
|
||
)
|
||
session.mount("http://", adapter)
|
||
session.mount("https://", adapter)
|
||
|
||
return session
|
||
|
||
def get(self, **kwargs) -> Any:
|
||
"""
|
||
GET请求
|
||
:param kwargs: 请求参数
|
||
:return: 响应内容
|
||
"""
|
||
return self._request(method="GET", parameters=Parameters(**kwargs))
|
||
|
||
def post(self, **kwargs) -> Any:
|
||
"""
|
||
POST请求
|
||
:param kwargs: 请求参数
|
||
:return: 响应内容
|
||
"""
|
||
return self._request(method="POST", parameters=Parameters(**kwargs))
|
||
|
||
def download(
|
||
self, stream_enabled: bool = False, chunk_size: int = 1024, **kwargs
|
||
) -> Any:
|
||
"""
|
||
下载
|
||
:param stream_enabled: 使用流式传输
|
||
:param chunk_size: 流式传输的分块大小
|
||
:param kwargs: 请求参数
|
||
:return: 响应内容
|
||
"""
|
||
response = self._request(
|
||
method="GET",
|
||
parameters=Parameters(**{"stream_enabled": stream_enabled, **kwargs}),
|
||
)
|
||
# 若使用流式传输则处理流式传输响应
|
||
if stream_enabled:
|
||
return self._process_stream_response(
|
||
response=response, chunk_size=chunk_size
|
||
)
|
||
return response
|
||
|
||
def _request(self, method: Literal["GET", "POST"], parameters: Parameters) -> Any:
|
||
"""
|
||
请求
|
||
:param method: 请求方法
|
||
:param parameters: 请求参数模型
|
||
:return: 响应内容
|
||
"""
|
||
# 将请求参数模型转为请求参数字典
|
||
kwargs = parameters.model_dump(exclude_none=True, by_alias=True)
|
||
|
||
# 将统一资源定位符转为字符串
|
||
url = str(kwargs.pop("url"))
|
||
|
||
# 过滤表单数据中空值
|
||
if kwargs.get("data"):
|
||
kwargs["data"] = {k: v for k, v in kwargs["data"].items() if v}
|
||
|
||
# 过滤JSON数据中空值
|
||
if kwargs.get("json"):
|
||
kwargs["json"] = {k: v for k, v in kwargs["json"].items() if v}
|
||
|
||
# 使用流式传输
|
||
stream_enabled = kwargs.pop("stream_enabled", False)
|
||
|
||
# 缓存全局唯一标识
|
||
guid = kwargs.pop("guid", None)
|
||
# 若缓存非空且缓存全局唯一标识非空则查询并获取单条缓存
|
||
if self.caches and guid:
|
||
cache = self.caches.query(guid)
|
||
if cache:
|
||
return cache
|
||
|
||
# 发送请求并处理响应
|
||
try:
|
||
response = self.session.request(
|
||
method=method, url=url, timeout=self.timeout, **kwargs
|
||
)
|
||
response.raise_for_status() # 若返回非2??状态码则抛出异常
|
||
|
||
# 若使用流式传输则直接返回响应对象（不缓存）
|
||
if stream_enabled:
|
||
return response
|
||
|
||
# 处理响应对象
|
||
response = self._process_response(response=response)
|
||
# 若使用缓存且缓存全局唯一标识非空则新增或更新缓存
|
||
if self.caches and guid:
|
||
self.caches.update(guid, response)
|
||
|
||
return response
|
||
|
||
# 重构异常信息
|
||
except Exception as exception:
|
||
try:
|
||
response = getattr(exception, "response", None)
|
||
status = (
|
||
response.json().get("status", response.status_code)
|
||
if response
|
||
else None
|
||
)
|
||
message = (
|
||
response.json().get("message", response.text)
|
||
if response
|
||
else str(exception).splitlines()[0]
|
||
)
|
||
except Exception:
|
||
status = None
|
||
message = f"{method} {kwargs["url"]} 请求发生异常：{str(exception).splitlines()[0]}"
|
||
return RequestException(status=status, message=message).__dict__
|
||
|
||
@staticmethod
|
||
def _process_response(
|
||
response: Response,
|
||
) -> Any:
|
||
"""
|
||
处理响应对象
|
||
:param response: 响应对象
|
||
:return: 响应内容
|
||
"""
|
||
content = response.content
|
||
if not content:
|
||
return None
|
||
|
||
# 响应类型
|
||
response_type = (
|
||
response.headers.get("Content-Type", "").split(";")[0].strip().lower()
|
||
)
|
||
# 根据响应类型匹配响应内容解析方法并返回
|
||
match response_type:
|
||
# 若为JSON则反序列化
|
||
case "application/json" | "text/json":
|
||
return response.json()
|
||
# 若为XML解析为Element对象
|
||
case "application/xml" | "text/xml":
|
||
return ElementTree.fromstring(content)
|
||
# 若为影像件格式则返回影像件格式和响应内容
|
||
case _ if response_type.startswith("image/"):
|
||
return response_type.split(sep="/", maxsplit=1)[1], content
|
||
case _:
|
||
try:
|
||
return content.decode("utf-8")
|
||
except UnicodeDecodeError:
|
||
return content
|
||
|
||
# 处理流式传输响应
|
||
@staticmethod
|
||
def _process_stream_response(
|
||
response: Response, chunk_size: int
|
||
) -> Generator[bytes, None, None]:
|
||
"""
|
||
处理流式响应
|
||
:param response: 响应对象
|
||
:param chunk_size: 分块大小
|
||
:return: 响应内容迭代器
|
||
"""
|
||
try:
|
||
for chunk in response.iter_content(chunk_size=chunk_size):
|
||
if chunk:
|
||
yield chunk
|
||
finally:
|
||
response.close()
|